数据集制作:从原始图片到YOLO能吃的数据

做目标检测,算法选得再好,模型调得再花哨,最后拼的还是数据。我见过太多项目,算法工程师花三个月调模型,最后发现是标注数据出了问题——标签漏了、格式错了、类别对不上。嗯,这章我们就来把数据集制作这件事彻底讲透。

核心观点:数据集的质量直接决定模型的天花板。标注工具、格式转换、数据增强,每一步都是坑,踩一个就得返工。

3.1 LabelImg标注工具:入门与避坑

LabelImg 是我用得最多的标注工具。开源、轻量、支持 YOLO 格式直接导出。说白了,它就是给图片画框、打标签的。

安装与启动

# 直接用 pip 安装
pip install labelImg

# 启动
labelImg

启动后界面很简洁。左边是图片列表,中间是画布,右边是标签列表。我个人习惯把图片和标注文件放在同一个目录下,方便管理。

标注流程

  1. 打开图片目录:File → Open Dir,选择你的图片文件夹
  2. 选择保存目录:File → Change Save Dir,指定标注文件存放位置
  3. 开始标注:按 w 键进入画框模式,拖拽鼠标框选目标
  4. 输入标签:弹窗输入类别名称,比如 person、car
  5. 保存:按 Ctrl+S,生成对应的 .xml 或 .txt 文件

快捷键小技巧:

  • w:画框
  • d:下一张
  • a:上一张
  • Ctrl+U:自动保存模式,省得一张张按保存

我曾经踩过的坑

有一次项目赶工期,我让实习生用 LabelImg 标了 5000 张图。结果训练时发现模型死活不收敛。排查了半天,原来是标注文件里类别索引写错了——YOLO 要求类别从 0 开始,他写成了从 1 开始。嗯,从那以后我定了个规矩:标注完成后必须做一次格式校验。

注意:LabelImg 默认输出 Pascal VOC 格式(.xml),YOLO 需要的是 .txt 格式。记得在 View → Auto Saving 里勾选 YOLO 格式,或者在保存时手动选择。

3.2 COCO数据集格式解析

COCO 是目前目标检测领域最主流的数据集格式。很多预训练模型都是在 COCO 上训的。理解它的结构,你才能把自己的数据转成模型能吃的格式。

COCO 的 JSON 结构

COCO 用单个 JSON 文件描述整个数据集。核心字段就三个:

字段 说明 示例
images 图片信息列表 id, file_name, width, height
annotations 标注信息列表 id, image_id, category_id, bbox, area
categories 类别信息列表 id, name, supercategory

bbox 的格式是 [x, y, width, height],注意是左上角坐标加宽高,不是中心点坐标。这个和 YOLO 不一样,转换时要小心。

为什么需要理解 COCO?

你想想看,现在主流的目标检测框架(MMDetection、Detectron2)都默认用 COCO 格式。如果你想把自定义数据集跑在这些框架上,就必须转成 COCO 格式。我建议你直接把自己的数据做成 COCO 格式,这样以后换框架也方便。

3.3 自定义数据集转换脚本

实际项目中,标注数据可能是各种格式:LabelImg 的 XML、LabelMe 的 JSON、或者自己写的 CSV。我们需要一个统一的转换脚本。

从 LabelImg XML 转 YOLO TXT

import xml.etree.ElementTree as ET
import os

def convert_xml_to_yolo(xml_path, output_dir, class_names):
    tree = ET.parse(xml_path)
    root = tree.getroot()
    
    # 获取图片尺寸
    width = int(root.find('size/width').text)
    height = int(root.find('size/height').text)
    
    # 获取图片文件名(不含扩展名)
    image_name = root.find('filename').text
    base_name = os.path.splitext(image_name)[0]
    
    # 写入 YOLO 格式
    with open(os.path.join(output_dir, base_name + '.txt'), 'w') as f:
        for obj in root.findall('object'):
            class_name = obj.find('name').text
            class_id = class_names.index(class_name)
            
            # 解析边界框
            xmin = int(obj.find('bndbox/xmin').text)
            ymin = int(obj.find('bndbox/ymin').text)
            xmax = int(obj.find('bndbox/xmax').text)
            ymax = int(obj.find('bndbox/ymax').text)
            
            # 转成 YOLO 归一化坐标
            x_center = (xmin + xmax) / 2.0 / width
            y_center = (ymin + ymax) / 2.0 / height
            w = (xmax - xmin) / width
            h = (ymax - ymin) / height
            
            f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}\n")

# 使用示例
class_names = ['person', 'car', 'bicycle']  # 你的类别列表
convert_xml_to_yolo('image_001.xml', './labels/', class_names)

个人经验:转换脚本一定要做边界检查。有时候标注框会超出图片边界(比如 xmin < 0),YOLO 训练时会报错。我习惯在转换时加一句 xmin = max(0, xmin) 来兜底。

从 COCO JSON 转 YOLO TXT

import json
import os

def coco_to_yolo(coco_json_path, output_dir):
    with open(coco_json_path, 'r') as f:
        coco_data = json.load(f)
    
    # 建立类别映射
    categories = {cat['id']: idx for idx, cat in enumerate(coco_data['categories'])}
    
    # 建立图片信息映射
    images = {img['id']: img for img in coco_data['images']}
    
    # 按图片分组标注
    annotations_by_image = {}
    for ann in coco_data['annotations']:
        img_id = ann['image_id']
        if img_id not in annotations_by_image:
            annotations_by_image[img_id] = []
        annotations_by_image[img_id].append(ann)
    
    # 逐图片写入
    for img_id, anns in annotations_by_image.items():
        img = images[img_id]
        base_name = os.path.splitext(img['file_name'])[0]
        
        with open(os.path.join(output_dir, base_name + '.txt'), 'w') as f:
            for ann in anns:
                class_id = categories[ann['category_id']]
                x, y, w, h = ann['bbox']
                
                # COCO bbox 是 [x, y, width, height],转成 YOLO 格式
                x_center = (x + w / 2) / img['width']
                y_center = (y + h / 2) / img['height']
                w_norm = w / img['width']
                h_norm = h / img['height']
                
                f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {w_norm:.6f} {h_norm:.6f}\n")

3.4 数据增强策略

数据增强是提升模型泛化能力最有效的手段之一。说白了,就是用有限的标注数据,生成更多样化的训练样本。

YOLO 常用的数据增强

增强方法 说明 适用场景
Mosaic 将4张图拼成1张 小目标检测,提升上下文理解
随机缩放 0.5~1.5倍随机缩放 多尺度目标
色彩抖动 调整亮度、对比度、饱和度 光照变化大的场景
随机翻转 水平/垂直翻转 对称目标(如车辆、行人)
Cutout 随机遮挡部分区域 提升遮挡鲁棒性

Mosaic 增强的原理

Mosaic 是 YOLOv4 引入的经典增强。它把四张图随机裁剪、缩放后拼成一张大图。这样做的好处是:

  • 一张图里能看到多个场景,模型学到的特征更丰富
  • 小目标被放大的概率增加,对小目标检测友好
  • 相当于变相增加了 batch size,训练更稳定

我的建议:Mosaic 增强在训练初期效果很好,但到了后期可能会让模型过拟合到拼接痕迹上。我一般会在最后 10% 的 epoch 关掉 Mosaic,只用基础增强做微调。

数据增强的注意事项

避坑指南:

  • 增强后的标注框必须同步变换,不能只增强图片不改框
  • 不要对文本类目标做水平翻转(文字会反)
  • 色彩抖动幅度不要太大,否则模型会学到颜色噪声
  • 测试集和验证集不要做任何增强,保持原始数据

3.5 数据集制作全流程

我把整个流程总结成一张图,方便你对照执行:

数据集制作全流程 原始图片采集 LabelImg标注 格式转换脚本 数据增强 数据集划分 训练集 (70%) 验证集 (20%) 测试集 (10%) 关键说明: • 训练集:用于模型训练,做数据增强 • 验证集:用于调参和早停,不做增强 • 测试集:用于最终评估,保持原始数据

数据集划分建议

我一般按 70% 训练、20% 验证、10% 测试来划分。但要注意:

  • 如果数据量少(比如不到 1000 张),可以改成 80% 训练、10% 验证、10% 测试
  • 确保每个类别在三个子集中都有分布,不要出现某个类别只在训练集里出现的情况
  • 视频抽帧的数据要按视频分组划分,避免同一视频的帧同时出现在训练和验证集里

一个小工具:我写了个简单的 Python 脚本,用 sklearn.model_selection.train_test_split 做分层采样,保证每个类别的比例在三个子集中一致。需要的话可以直接复用。

3.6 本章小结

数据集制作看起来是体力活,但细节决定成败。从标注工具的选择,到格式转换的校验,再到数据增强的策略,每一步都值得认真对待。我个人觉得,花在数据上的时间永远不会白费——好的数据集,能让模型训练事半功倍。