第4章 数据集准备：从标注到划分，手把手搭好训练地基

数据集准备这件事，说简单也简单，说复杂也复杂。我见过太多同学模型跑不起来，最后发现是数据集格式不对、标注文件缺漏、或者训练集和验证集混在一起。嗯，这一章咱们就把这些坑一个个填平。

4.1 LabelImg 标注工具：从安装到实战

LabelImg 是我个人最常用的图像标注工具。它轻量、跨平台、支持 YOLO 和 VOC 两种输出格式。说白了，你画个框，它帮你生成标签文件。

4.1.1 安装

我建议直接用 pip 安装，省事：

pip install labelImg
labelImg

如果遇到 Qt 相关的报错，多半是 Python 版本问题。我曾在 Ubuntu 18.04 上折腾了半小时，最后发现用 Python 3.8 就一切正常。

4.1.2 标注流程

1. 打开图片目录：点击 "Open Dir" 选择你的图片文件夹
2. 创建矩形框：按 w 键进入画框模式，拖拽鼠标框住目标
3. 输入标签名：比如 "person"、"car"、"cat"
4. 保存标注：按 Ctrl+S，默认生成同名的 XML 文件（Pascal VOC 格式）
5. 下一张：按 d 键切换到下一张图片

4.1.3 输出格式切换

LabelImg 默认输出 Pascal VOC 格式（XML），但 YOLO 需要的是 TXT 格式。你可以在菜单栏 View → Auto Save Mode 开启自动保存，然后在 PascalVOC → YOLO 切换输出格式。

4.2 COCO 数据集格式：为什么它是行业标准？

COCO（Common Objects in Context）是目前最主流的目标检测数据集格式。YOLOv5、YOLOv8 官方都支持直接读取 COCO 格式。我个人觉得，学会 COCO 格式，就等于拿到了通行证。

4.2.1 COCO 的 JSON 结构

COCO 标注文件是一个 JSON，核心包含三个字段：

{
  "images": [
    {"id": 1, "file_name": "000001.jpg", "width": 640, "height": 480}
  ],
  "annotations": [
    {
      "id": 1,
      "image_id": 1,
      "category_id": 1,
      "bbox": [100, 150, 200, 300],
      "area": 60000,
      "iscrowd": 0
    }
  ],
  "categories": [
    {"id": 1, "name": "person"}
  ]
}

这里 bbox 的格式是 [x, y, width, height]，注意是左上角坐标 + 宽高，不是中心点坐标。我刚开始做转换时就在这里栽过跟头。

4.2.2 从 LabelImg 到 COCO 的转换

LabelImg 标注出来的是 Pascal VOC 格式，我们需要转成 COCO JSON。我一般用 voc2coco 这个工具：

pip install voc2coco
voc2coco --ann_dir ./Annotations --ann_ids ./train.txt --labels ./classes.txt --output ./train.json

如果你不想装第三方库，也可以自己写个脚本。核心逻辑就是：遍历 XML 文件，提取每个目标的 bbox 和标签，然后组装成上面的 JSON 结构。

4.3 自定义数据集制作：从零到一

实际项目中，你很少能直接用公开数据集。比如我做过一个工地安全帽检测项目，网上找不到现成的数据集，只能自己拍、自己标。

4.3.1 数据采集

• 摄像头拍摄：用手机或 USB 摄像头，在不同角度、光照、距离下拍摄
• 视频抽帧：录一段视频，用 OpenCV 每隔 N 帧保存一张
• 网络爬取：注意版权问题，仅用于学术研究

我个人建议，每个类别至少收集 500 张图片。太少的话，模型容易过拟合。

4.3.2 数据集目录结构

YOLO 官方推荐的结构是这样的：

dataset/
├── images/
│   ├── train/
│   │   ├── 000001.jpg
│   │   └── ...
│   ├── val/
│   │   └── ...
│   └── test/
│       └── ...
├── labels/
│   ├── train/
│   │   ├── 000001.txt
│   │   └── ...
│   ├── val/
│   │   └── ...
│   └── test/
│       └── ...
└── dataset.yaml

注意：图片和标签的文件名要一一对应，只是后缀不同。比如 000001.jpg 对应 000001.txt。

4.3.3 编写 dataset.yaml

这个文件告诉 YOLO 你的数据集在哪里、有哪些类别：

train: ./dataset/images/train
val: ./dataset/images/val
test: ./dataset/images/test

nc: 3
names: ['person', 'car', 'cat']

4.4 数据集划分：训练集、验证集、测试集

划分比例我一般用 8:1:1 或者 7:2:1。为什么？你想想看，训练集太少模型学不到东西，验证集太少调参没参考，测试集太少评估不准确。

4.4.1 手动划分脚本

下面这个脚本我用了好几年，每次做新项目都直接拿来用：

import os
import random
import shutil

def split_dataset(img_dir, output_dir, train_ratio=0.8, val_ratio=0.1):
    images = [f for f in os.listdir(img_dir) if f.endswith('.jpg')]
    random.shuffle(images)
    
    train_count = int(len(images) * train_ratio)
    val_count = int(len(images) * val_ratio)
    
    train_images = images[:train_count]
    val_images = images[train_count:train_count+val_count]
    test_images = images[train_count+val_count:]
    
    for split, split_images in zip(['train', 'val', 'test'], 
                                   [train_images, val_images, test_images]):
        os.makedirs(f'{output_dir}/images/{split}', exist_ok=True)
        os.makedirs(f'{output_dir}/labels/{split}', exist_ok=True)
        
        for img in split_images:
            shutil.copy(f'{img_dir}/{img}', f'{output_dir}/images/{split}/{img}')
            label_file = img.replace('.jpg', '.txt')
            if os.path.exists(f'{label_dir}/{label_file}'):
                shutil.copy(f'{label_dir}/{label_file}', 
                           f'{output_dir}/labels/{split}/{label_file}')

split_dataset('./raw_images', './dataset')

4.4.2 划分时的注意事项

• 随机打乱：一定要先 random.shuffle，否则连续帧的图片会全部进训练集或验证集
• 保持分布一致：如果某个类别的样本很少，建议用分层采样，确保每个子集中都有该类别的样本
• 不要泄露：测试集在训练过程中不能碰，我见过有人用测试集调参，结果模型在真实场景中表现很差

4.5 本章知识体系总览

下面这张图把数据集准备的整个流程串起来了，你可以对照着检查自己哪一步还没做：