快速构建 yolo 训练集 测试集

在机器学习工作流程中，数据处理是一个关键步骤。通常我们会使用不同的工具来标注数据，而每种工具都有其特定的格式。在这篇文章中，我们将展示如何从数据集中快速抽取样本，构建训练集和测试集。

YOLO简介

YOLO（You Only Look Once）是一种实时对象检测系统，以其速度和精度著称。YOLO将对象检测视为一个回归问题，直接从图像像素中预测边界框和类别概率。与其他检测器不同，YOLO在单一神经网络中完成整个检测过程，因而速度非常快，适用于实时应用。

Show Me the Code

标签合法性检测

在处理数据集时，确保每个图像都有对应的标签文件是至关重要的。以下代码用于检查数据集中图像和标签文件的配对情况：

import osdef check_file_pairs(directory):files = os.listdir(directory)images = [f for f in files if f.endswith('.jpg')]labels = [f for f in files if f.endswith('.txt')]image_basenames = set(os.path.splitext(f)[0] for f in images)label_basenames = set(os.path.splitext(f)[0] for f in labels)missing_labels = image_basenames - label_basenamesmissing_images = label_basenames - image_basenamesif missing_labels:print("Missing label files for images:", missing_labels)if missing_images:print("Missing image files for labels:", missing_images)if not missing_labels and not missing_images:print("All image and label files are correctly paired.")# Example usage
directory = 'F:\\watermeter\\data\\ocr_data\\all'
check_file_pairs(directory)

代码解释

该函数遍历给定目录中的所有文件，分别列出以.jpg结尾的图像文件和以.txt结尾的标签文件。通过集合操作，我们可以找出哪些图像文件缺少对应的标签文件，反之亦然。这样，我们可以确保数据集的完整性，避免后续训练过程中出现错误。

复制样本数据

为了构建训练集和验证集，我们需要将数据集保存在all文件夹中，并按预设比例随机选取样本，分别复制到train和val文件夹中。

import os
import shutil
import randomdef create_directories(base_dir):train_dir = os.path.join(base_dir, 'train')val_dir = os.path.join(base_dir, 'val')if os.path.exists(train_dir) or os.path.exists(val_dir):print("Train or Val directory already exists. Exiting.")return Falseos.makedirs(train_dir)os.makedirs(val_dir)return Truedef split_data(base_dir, train_ratio=0.8):all_dir = os.path.join(base_dir, 'all')train_dir = os.path.join(base_dir, 'train')val_dir = os.path.join(base_dir, 'val')images = [f for f in os.listdir(all_dir) if f.endswith('.jpg')]random.shuffle(images)train_size = int(len(images) * train_ratio)train_images = images[:train_size]val_images = images[train_size:]for image in train_images:label = image.replace('.jpg', '.txt')shutil.copy(os.path.join(all_dir, image), train_dir)shutil.copy(os.path.join(all_dir, label), train_dir)for image in val_images:label = image.replace('.jpg', '.txt')shutil.copy(os.path.join(all_dir, image), val_dir)shutil.copy(os.path.join(all_dir, label), val_dir)base_dir = 'path_to_data'
if create_directories(base_dir):split_data(base_dir, train_ratio=0.8)

代码解释

1. 创建目录：create_directories函数用于在指定的基础目录下创建train和val目录。如果目录已存在，函数将退出以防止覆盖现有数据。
2. 数据分割：split_data函数负责将数据集按指定比例分为训练集和验证集。首先，获取所有图像文件并随机打乱顺序。然后，根据比例计算训练集的大小，将图像和对应的标签文件复制到相应的目录中。

通过这些步骤，您可以快速构建适用于YOLO训练的高质量数据集。希望这篇文章能帮助您更高效地进行数据准备工作！

---

欢迎关注作者公众号，获取更多技术分享