在机器学习工作流程中，数据处理是一个关键步骤。通常我们会使用不同的工具来标注数据，而每种工具都有其特定的格式。在这篇文章中，我们将展示如何从数据集中快速抽取样本，构建训练集和测试集。
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)
代码解释
- 创建目录：
create_directories函数用于在指定的基础目录下创建train和val目录。如果目录已存在，函数将退出以防止覆盖现有数据。 - 数据分割：
split_data函数负责将数据集按指定比例分为训练集和验证集。首先，获取所有图像文件并随机打乱顺序。然后，根据比例计算训练集的大小，将图像和对应的标签文件复制到相应的目录中。
通过这些步骤，您可以快速构建适用于YOLO训练的高质量数据集。希望这篇文章能帮助您更高效地进行数据准备工作！
欢迎关注作者公众号，获取更多技术分享
