4. 数据集准备

做工业缺陷检测，模型选得再好，算法调得再花哨，最后拼的还是数据。我见过太多项目，算法工程师花三个月调参，最后发现是训练集里缺陷样本不够，或者标注格式搞错了——那种感觉，真的很崩溃。

这一章，我们就来聊聊数据集准备的那些事。说白了，就是三件事：数据从哪来、数据怎么标、数据怎么分。

4.1 工业缺陷数据集的来源

工业场景和互联网场景最大的区别是什么？数据量少，而且缺陷样本极其稀缺。一个正常的产线，良品率可能高达99.9%，你要收集那0.1%的缺陷，得跑多少天？

我个人习惯把数据来源分成三类：

4.2 标注格式：COCO vs YOLO

标注格式这件事，看起来简单，但坑特别多。我见过有人把 COCO 格式的标注直接喂给 YOLO，结果训练出来的模型检测框全飘了——因为坐标归一化方式不一样。

先看一个直观的对比：

具体到文件内容，我们来看一个例子。假设有一张 640×480 的图片，里面有一个缺陷框，左上角在 (100, 50)，宽 200，高 150。

COCO 格式（JSON 片段）：

{
    "images": [
        {
            "id": 1,
            "file_name": "defect_001.jpg",
            "width": 640,
            "height": 480
        }
    ],
    "annotations": [
        {
            "id": 1,
            "image_id": 1,
            "category_id": 1,
            "bbox": [100, 50, 200, 150],
            "area": 30000,
            "iscrowd": 0
        }
    ],
    "categories": [
        {"id": 1, "name": "scratch"}
    ]
}

YOLO 格式（.txt 文件内容）：

0 0.3125 0.2604 0.3125 0.3125

YOLO 的坐标计算方式：

• x_center = (100 + 200/2) / 640 = 0.3125
• y_center = (50 + 150/2) / 480 ≈ 0.2604
• w = 200 / 640 = 0.3125
• h = 150 / 480 = 0.3125

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

数据划分看起来简单，不就是 7:2:1 或者 8:1:1 吗？但在工业场景里，有个很关键的问题：同一张图片的不同裁剪块，不能同时出现在训练集和验证集里。

为什么？因为工业相机拍的是连续的大图，我们通常会切成很多小图来训练。如果这些小图来自同一张大图，它们之间的背景、光照几乎一模一样，模型很容易学到“记忆”而不是“泛化”。

我建议的划分策略：

1. 按大图划分：先把原始大图分成训练、验证、测试三组，然后再切小图。
2. 保证缺陷分布均衡：统计每个类别的缺陷数量，确保验证集和测试集里每个类别都有足够的样本。
3. 时间维度考虑：如果数据是连续采集的，建议按时间顺序划分——用前几天的数据训练，后几天的数据验证。这样可以检验模型对产线变化的鲁棒性。

下面是一个简单的 Python 脚本，演示如何按大图进行划分：

import os
import random
import shutil
from sklearn.model_selection import train_test_split

# 假设所有大图都在 raw_images 目录下
raw_images = os.listdir('raw_images')
# 按 7:2:1 划分
train_imgs, temp_imgs = train_test_split(raw_images, test_size=0.3, random_state=42)
val_imgs, test_imgs = train_test_split(temp_imgs, test_size=0.333, random_state=42)

# 然后根据大图名称，找到对应的切分小图和标注文件
# 复制到对应的 train/val/test 目录
for img in train_imgs:
    # 假设小图命名规则：大图名_序号.jpg
    # 这里需要根据实际情况实现
    pass

4.4 知识结构图：数据集准备全流程

为了让你对整个流程有个直观的认识，我画了一张图：

4.5 实战中的几个小建议

最后，分享几个我在项目中积累的经验：

• 标注工具选型：我个人推荐 LabelImg（轻量级）和 LabelStudio（支持多人协作）。如果团队超过 3 个人，建议用 LabelStudio，可以省去很多合并标注文件的麻烦。
• 标注质量检查：一定要做二次质检。我见过标注员把背景上的灰尘当成缺陷标了，结果模型学了一堆“假缺陷”。
• 数据版本管理：用 DVC 或者 Git LFS 管理数据集。别问我为什么强调这个——有一次我不小心覆盖了标注文件，差点重标了 5000 张图。
• 格式转换脚本：写一个通用的 COCO 转 YOLO 脚本，放在项目仓库里。每次有新数据，跑一遍脚本就行，避免手动转换出错。