4. 数据集准备:工业缺陷检测的“弹药库”

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

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

4.1 工业缺陷数据集的来源

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

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

来源类型 说明 优缺点
公开数据集 如 MVTec AD、NEU-DET、PCB Defect 等 ✅ 拿来就能用,适合算法验证
❌ 场景单一,和实际产线差距大
产线采集 在真实产线上架设工业相机,连续采集 ✅ 数据最真实,缺陷分布符合实际
❌ 良品太多,缺陷太少,需要长周期
数据增强/合成 用 GAN、扩散模型或传统图像处理生成缺陷 ✅ 快速扩充缺陷样本
❌ 合成数据与真实数据存在域偏移
我的建议: 别只依赖公开数据集。我做过一个项目,用 MVTec AD 训练出来的模型,在客户产线上直接掉点 15 个点。后来老老实实去产线蹲了一周,重新采集数据,效果才上来。公开数据集只适合做 baseline,真正的工业项目,必须有自己的产线数据。

4.2 标注格式:COCO vs YOLO

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

先看一个直观的对比:

对比维度 COCO 格式 YOLO 格式
坐标表示 左上角 (x, y) + 宽高 (w, h) 中心点 (x_center, y_center) + 宽高 (w, h)
归一化 像素值 相对图像尺寸(0~1)
存储方式 一个 JSON 文件包含所有标注 每张图片对应一个 .txt 文件
类别定义 通过 categories 字段映射 直接使用类别索引(从0开始)
适用场景 通用目标检测、实例分割 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
我曾经踩过的坑: 有一次用 LabelImg 标注,默认输出的是 YOLO 格式,但我没注意类别索引从 0 开始。我的缺陷类别有 3 种,结果第 3 类被标成了索引 2,但训练时类别总数设成了 3,导致模型把第 2 类和第 3 类搞混了。排查了两天才发现。所以,标注完成后,一定要检查类别索引是否正确。

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 知识结构图:数据集准备全流程

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

工业缺陷检测数据集准备流程 数据来源 公开数据集 产线采集 数据增强/合成 标注格式 COCO 格式 YOLO 格式 格式转换工具 数据划分 训练集 70% 验证集 20% 测试集 10% 关键注意事项 按大图划分 避免同源小图 出现在不同集合 缺陷分布均衡 每个类别在验证集 测试集中都有样本 时间维度划分 按采集时间顺序 检验模型泛化能力 最终产出 结构清晰的 train/val/test 目录 + 统一格式的标注文件

4.5 实战中的几个小建议

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

  • 标注工具选型:我个人推荐 LabelImg(轻量级)和 LabelStudio(支持多人协作)。如果团队超过 3 个人,建议用 LabelStudio,可以省去很多合并标注文件的麻烦。
  • 标注质量检查:一定要做二次质检。我见过标注员把背景上的灰尘当成缺陷标了,结果模型学了一堆“假缺陷”。
  • 数据版本管理:用 DVC 或者 Git LFS 管理数据集。别问我为什么强调这个——有一次我不小心覆盖了标注文件,差点重标了 5000 张图。
  • 格式转换脚本:写一个通用的 COCO 转 YOLO 脚本,放在项目仓库里。每次有新数据,跑一遍脚本就行,避免手动转换出错。
一个小技巧: 在划分数据之前,先跑一个简单的统计脚本,看看每个类别的样本数量。如果某个缺陷类别只有几十个样本,建议先做数据增强,再划分。否则验证集里可能一个该类别的样本都没有,评估结果会虚高。