4. 数据集准备
在AI分拣机器人项目中,模型能否准确识别目标物体并驱动机械臂完成抓取,其根基完全取决于训练数据集的质量与结构。一个仅有数百张图片、标注不一致、类别分布失衡的数据集,即便使用最先进的YOLOv8或RT-DETR模型,也难以在真实产线上达到99%以上的抓取准确率。本章将系统拆解数据集准备的四个核心环节:数据采集策略、标注工具选型与实操、数据增强方法、以及数据集划分与格式转换,帮助读者建立一套可复用的工业级数据工程流程。
4.1 数据采集:从产线到数据集的工程化方法
核心原理:数据采集不是简单地用相机拍几张照片,而是需要模拟模型部署后的真实推理环境。对于分拣机器人,相机通常固定安装在传送带上方或机械臂末端,因此采集视角、光照条件、物体姿态必须与运行时保持一致。采集策略的核心原则是“覆盖部署域”——即训练数据的分布应尽可能覆盖推理时可能遇到的所有情况。
- 相机参数:分辨率(建议≥1280×720)、帧率(≥30fps)、曝光时间、增益、白平衡固定为与部署一致。
- 光照条件:至少采集3种光照(强光、正常、暗光),若产线有频闪灯需同步采集。
- 物体姿态:每个物体至少包含5种旋转角度(0°/45°/90°/135°/180°)和3种倾斜状态。
- 遮挡情况:模拟物体相互遮挡、部分出界、被传送带边缘切割等场景。
- 负样本:采集不含目标物体的背景图,用于降低误检率。
操作流程:首先,使用工业相机(如海康MV-CA050-10GC)或高分辨率USB相机采集原始视频流;其次,通过视频抽帧工具(如FFmpeg)按每秒1-2帧抽取关键帧,避免连续帧高度相似导致的数据冗余;最后,人工筛选剔除模糊、过曝、失焦的图片。对于分拣项目,建议每类物体采集500-1000张有效图片,总数据集规模控制在3000-8000张之间,既保证模型收敛,又不至于标注成本过高。
4.2 标注工具选型:LabelImg与CVAT的对比与实操
核心原理:标注是将原始图像转化为模型可学习的监督信号。对于目标检测任务,标注的核心是确定每个目标物体的边界框(Bounding Box)和类别标签。标注质量直接决定模型的上限——如果标注框偏移超过5个像素,或类别标签混淆,模型将学到错误特征。
| 对比维度 | LabelImg | CVAT |
|---|---|---|
| 安装方式 | 本地Python包(pip install labelImg) | Docker容器或在线服务 |
| 标注类型 | 仅矩形框 | 矩形框、多边形、关键点、标签属性 |
| 多人协作 | 不支持,单机操作 | 支持多用户同时标注、任务分配 |
| 输出格式 | Pascal VOC(XML) | VOC、YOLO、COCO、TFRecord |
| 适用场景 | 个人或小团队(<500张) | 团队协作、复杂标注、大规模数据 |
操作流程(以LabelImg为例):
- 安装并启动:
pip install labelImg && labelImg,打开图片目录。 - 使用快捷键
W创建矩形框,拖动鼠标框选目标物体。 - 在弹出的对话框中输入类别名称(如“bottle”“can”“box”),点击OK。
- 按
D切换到下一张图片,重复标注。 - 保存后生成同名XML文件,包含
<object>标签及<bndbox>坐标信息。
4.3 数据增强:提升模型泛化能力的关键手段
核心原理:数据增强通过对原始图像进行随机变换,生成更多样化的训练样本,防止模型过拟合。对于分拣场景,增强策略必须保持语义不变——即变换后的图像仍能代表真实世界中的物体。例如,水平翻转对分拣任务有效,但垂直翻转可能不符合物理规律(物体不会倒置在传送带上)。
常用增强方法分类:
- 几何变换:随机旋转(±15°)、平移(±10%)、缩放(0.8-1.2倍)、水平翻转。
- 颜色扰动:亮度调整(±20%)、对比度(±15%)、饱和度(±10%)、HSV通道偏移。
- 噪声注入:高斯噪声、椒盐噪声、运动模糊(模拟传送带运动)。
- 高级增强:Mosaic(将4张图拼接为1张)、MixUp(两张图像素混合)、CutOut(随机遮挡区域)。
实战案例:某3C电子元件分拣项目,原始数据集仅1200张图片,模型在验证集上mAP@0.5仅为78.5%。通过引入Mosaic增强(每批次50%概率使用)和随机亮度扰动后,训练数据等效扩充至5000张,mAP提升至91.2%。但需注意,增强强度过大(如旋转超过30°)会导致物体形状失真,反而降低精度。
# 使用Albumentations库实现增强(推荐)
import albumentations as A
transform = A.Compose([
A.RandomRotate90(p=0.5),
A.HorizontalFlip(p=0.5),
A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.15, p=0.7),
A.MotionBlur(blur_limit=5, p=0.3),
A.Mosaic(p=0.5) # 需配合bbox_params使用
], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))
4.4 数据集划分与格式转换:从标注到训练的无缝衔接
核心原理:数据集划分遵循分层采样原则,确保训练集、验证集、测试集中的类别分布与整体数据集一致。划分比例通常为70%训练、15%验证、15%测试。验证集用于训练过程中的模型选择与超参数调优,测试集仅在最终评估时使用一次,避免信息泄露。
格式转换流程:不同标注工具和训练框架使用不同格式,常见的有Pascal VOC(XML)、YOLO(TXT)、COCO(JSON)。YOLO格式将边界框归一化为class_id x_center y_center width height,坐标值在0-1之间。
| 格式类型 | 文件后缀 | 坐标表示 | 适用框架 |
|---|---|---|---|
| Pascal VOC | .xml | 像素值(xmin, ymin, xmax, ymax) | Faster R-CNN、SSD |
| YOLO | .txt | 归一化(x_center, y_center, w, h) | YOLOv5/v8、RT-DETR |
| COCO | .json | 像素值(x, y, w, h) | Detectron2、MMDetection |
操作流程:首先,使用脚本将LabelImg生成的XML文件转换为YOLO格式TXT文件;其次,按类别比例随机划分数据集,生成train.txt、val.txt、test.txt文件,每行记录图片路径;最后,整理目录结构如下:
dataset/
├── images/
│ ├── train/ (2100张)
│ ├── val/ (450张)
│ └── test/ (450张)
├── labels/
│ ├── train/ (对应TXT文件)
│ ├── val/
│ └── test/
├── train.txt
├── val.txt
└── test.txt
sklearn.model_selection.train_test_split的stratify参数按类别比例划分后,三类物体mAP均提升至88%以上。此外,通过编写转换脚本将VOC格式转为YOLO格式,训练时间减少了20%(YOLO格式读取更快)。