快速构建 yolo 训练集 测试集

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

YOLO简介

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

Show Me the Code

标签合法性检测

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

import os

def 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_basenames
    missing_images = label_basenames - image_basenames
  
    if 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 random

def 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 False
  
    os.makedirs(train_dir)
    os.makedirs(val_dir)
    return True

def 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训练的高质量数据集。希望这篇文章能帮助您更高效地进行数据准备工作！

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