ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

YOLOv8工业工具装配检测:从数据集微调到TensorRT部署

YOLOv8工业工具装配检测:从数据集微调到TensorRT部署 简介面向工业目标检测与YOLO模型训练数据集聚焦自动化装配、智能质检和机器人视觉引导场景适合算法工程师与制造行业开发者使用。数据集提供YOLO格式标注覆盖枪型工具、钳型工具、尖头工具、扫描工具及手部共5类目标可直接接入YOLOv5/v8等主流框架便于构建工具识别与操作合规性检测模型。压缩包内共2000个文件以1873个txt标签文件为主体辅以125张jpg图片、1个yaml配置文件和1个说明文档整体大小80.98MB结构清晰适合直接导入训练流程。目前已有187人学习下载。样本来自真实装配线包含多角度操作细节可增强模型在复杂光照和遮挡条件下的鲁棒性。获取后即可获得带标注数据集、类别配置及配套说明能显著缩短工业场景数据准备周期支撑装配质量追溯与机器人自动化升级实践。1. 工具装配检测数据集通用预训练权重在产线上为什么失效产线上的工具检测和公开 COCO 数据里的杯子、椅子完全是两种难度镜头离工具很近金属表面反光枪型工具和钳型工具在侧面视角下轮廓经常叠在一起还要同时判断操作者的手是否接触工具。拿通用预训练权重直接上线漏检和误检会非常明显。这份「工具装配检测数据集」是工业装配场景实拍帧的标注集合训练集 1,873 张包含 gun-type tool、hand、pliers-type tool、pointed tool、scanning tool 五类标注以 YOLO 格式给出。因为标签文件可以直接喂给 YOLOv8 乃至更新的 YOLOv12所以适合用来训练自己的装配监控模型也适合做机器人抓取前的视觉定位和工位合规性判断。下面先从数据文件格式和类别语义开始拆。2. 五类标签与 YOLO 标注文件先从数据文件格式说起数据集的核心不是图片本身而是图片名、标签文件与类别命名之间的对应关系。工业目标检测项目和学术比赛不一样类别少但语义接近像 gun-type tool 与 pliers-type tool 在这个数据集里都带有长条形手柄如果只看局部特征很容易混。理解这五类为什么存在比直接跑训练更重要。2.1 类别设计与工业语义为什么“手”也是检测目标这个数据集的标签只有五类但每一类都有明确的产线语义。gun-type tool 指喷枪、胶枪这类以扳机触发为主的手持工具pliers-type tool 是钳子、夹持工具pointed tool 是螺丝刀、锥子、探针这类尖头工具scanning tool 是扫码枪、探伤笔这类用于扫描确认的电子工具hand 则覆盖人的手掌、手指和半握状态。category id标签名中文语义常见误配情况0gun-type tool枪型工具触发面通常有弧线与 pliers-type 在侧视时混淆1hand手部覆盖不同肤色/手套手持工具时框重叠导致一个框框两个目标2pliers-type tool钳型工具头部短粗、手柄长与 gun-type 的握把部分混淆3pointed tool尖头工具细长导向明显与小型 pliers 混淆程度低但小目标多4scanning tool扫描工具通常带指示灯或屏幕与抢型工具的握持区域混淆hand 作为独立类别作用不是“检测人”而是用来支撑合规性判断。装配线经常要求“拿取工具时必须让工具离开工件”或者“扫描前手不能遮挡条码”。只有同时检测到 hand 和 tool 各自的边界框才能进一步计算两个框之间的交并比从而判定持握、接近还是非法接触。所以训练时不能因为觉得 hand 目标太大而删掉这一类。2.2 解析 YOLO 标注一行标签的五个数值顺序YOLO 格式每个 .txt 文件对应一张同名图片每一行代表一个目标类别 id、归一化中心点 x、归一化中心点 y、归一化宽度 w、归一化高度 h。我拿到标签后习惯先写个小脚本把标注打出来确认坐标是否落在一个非空区域内避免后面训练时 loss 一开始就爆炸。import os label_file data/train/labels/1591587432198_jpg.rf.6cee4fe03f30966306dcc3e8d7360995.txt with open(label_file, r, encodingutf-8) as fp: lines [line.strip().split() for line in fp if line.strip()] for row in lines: cls_id int(row[0]) x_center, y_center, width, height map(float, row[1:]) print(fcls{cls_id} fcx{x_center:.3f} cy{y_center:.3f} fw{width:.3f} h{height:.3f})这段代码的关键点有三个第一row[0]是类别 id必须与 data YAML 里 names 的下标一致比如 0 对应 gun-type tool 而不是 hand第二后用四个浮点数都是相对图片尺寸的归一化坐标数值范围在 0 到 1 之间画框、裁剪时要重新乘以图像宽高第三坐标描述的是矩形中心点和宽高不是左上角与右下角许多刚从 COCO 转 YOLO 的人会在这里画错矩形。打印出来后如果出现w或h大于 1说明标注文件已经损坏或没归一化需要优先清洗而不是直接训练。2.3 文件名里的.rf.与多时序样本结构从文件名能读出不少采集信息例如WIN_20240613_04_24_07_Pro_mp4-0095_jpg.rf....jpg前半段多半是 Windows 设备在 2024 年 6 月 13 日凌晨 04:24 录制的视频mp4-0095表示从该视频第 95 帧附近抽帧而1591587432198_jpg这种纯时间戳命名更可能来自另一路摄像头或手机连续拍摄。两类文件名混合在一起说明这个数据集的“训练集 1,873 张”并不是同一机位的简单连续帧而是跨日期、跨时段、跨设备的样本集合。.rf.是导出工具在处理图片后追加的 id 后缀常见于 Roboflow 导出的 zip 包。它不会影响模型读取因为 YOLO 训练时只依赖图片路径和同名 txt 路径不解析文件名里的随机 hash。但你应该知道这类导出包通常按train/valid/test三个子目录组织。如果解压后只有 train 目录也完全够用只需要再用脚本从 train 里按比例切出一部分做验证集而不是直接拿全部训练集去训练否则最后只能看到过拟合后的 mAP。3. 用 YOLOv8 微调目录整理、YAML 配置与训练参数调整拿到工具装配检测数据集后最直接的方式是先用 YOLOv8 把基线跑出来。YOLOv8 对 YOLO 格式的原生支持最好安装 ultralytics 后不需要写自定义 Dataset 类。这一章从目录整理讲到训练参数每一步都是可以照抄的。3.1 统一目录结构并声明数据 YAMLRoboflow 导出的 zip 解压后路径不统一有的直接是train/images有的多套一层文件夹。为了让训练命令可重复执行我会先整理成标准结构mkdir -p dataset/{images/{train,val},labels/{train,val}} mv original/train/*.jpg dataset/images/train/ mv original/train/*.txt dataset/labels/train/ # 从训练集抽 8% 做验证集这里用简单 shell 方式正式项目可以写 sklearn 脚本 cd dataset/images/train ls *.jpg | shuf -n 150 | xargs -I {} mv {} ../val/ cd ../../labels/train ls *.txt | shuf -n 150 | xargs -I {} mv {} ../val/把图片和标签保持同名同步移动很关键因为 YOLO 训练只找同名 txt没有匹配标注的图片会被自动忽略而缺少图片的孤立 txt 不会被检查。移动完成后写一个数据描述文件tools_dataset.yamlpath: /absolute/path/to/dataset train: images/train val: images/val names: 0: gun-type tool 1: hand 2: pliers-type tool 3: pointed tool 4: scanning toolpath建议写绝对路径训练进程的工作目录经常变化相对路径会让 ultralytics 找不到数据。names的顺序就是类别 id 的顺序必须和标签 txt 里的 0-4 一一对应这里顺序错一位整个模型的输出语义就全部错位。如果你的解压包里有 valid 目录直接把val指向images/valid即可。3.2 训练命令与超参数选择在 1,873 张图上我习惯先用轻量骨架做一次快速验证yolo detect train \ datatools_dataset.yaml \ modelyolov8n.pt \ epochs120 \ batch16 \ imgsz640 \ patience20 \ projecttools_det \ namebaseline_nmodelyolov8n.pt表示加载 COCO 预训练权重而不是从随机初始化开始对工业小数据集来说这一步能省大量收敛时间。epochs120在百级批量上足够patience20是连续 20 个 epoch 验证集 mAP 不再提升就早停防止最后过拟合。batch16配合imgsz640大约需要 16GB 以上显存如果显存吃紧优先把batch降到 8再把imgsz降到 512不要同时降两个。训练结束后runs/detect/tools_det/baseline_n/weights/下会生成last.pt和best.pt前者是最后一个 epoch 的权重后者是验证指标最好的权重。后续验证和部署全部使用best.pt。3.3 统计各类别样本量提前发现不均衡很多人拿到数据集直接训练跑完发现手部类别 AP 很高工具类别 AP 很低回头一查才发现 hand 的框比所有工具加起来还多。所以在第一次训练前我会先统计一下标签文件中各类别的目标数import glob from collections import Counter total Counter() for txt_path in glob.glob(dataset/labels/train/*.txt): with open(txt_path, r, encodingutf-8) as fp: for line in fp: if line.strip(): total[int(line.split()[0])] 1 print(total)这段脚本把每个 txt 的第一列取出来做累加输出类似Counter({1: 3021, 0: 876, 2: 643, 4: 421, 3: 358})的结果。如果 hand 的样本量接近其他类别总和的 2 倍以上就不需要额外处理太多因为 hand 本来就是一个高变化度的类别多一些负样本能降低误报但如果某个工具类别只有 300 多个目标接下来就要在数据增强阶段给这个类别更保守的翻转和裁剪策略避免小类别的形态特征被过度破坏。4. 数据增强与混淆排查让模型在反光和遮挡下更稳定工业场景图片看着干净但真实推理时经常会遇到镜头抖动、工件遮住工具、金属反光让局部纹理消失等情况。这一类问题不能只靠加数据要把增强策略稳定住并且用验证集的混淆矩阵来判断模型到底在哪些类别之间打转。4.1 用 YOLOv8 内置增强参数模拟产线环境YOLOv8 允许在数据 YAML 中直接覆盖增强参数不需要修改代码。我在 baseline 跑通后会复制一份tools_aug.yaml并追加path: /absolute/path/to/dataset train: images/train val: images/val names: 0: gun-type tool 1: hand 2: pliers-type tool 3: pointed tool 4: scanning tool hsv_h: 0.015 hsv_s: 0.7 hsv_v: 0.5 degrees: 15 translate: 0.1 scale: 0.4 fliplr: 0.5 mosaic: 0.8 mixup: 0.1hsv_s和hsv_v负责模拟不同灯光下的饱和度与亮度变化speckle类增强在 ultralytics 中没有直接参数所以用 HSV 扰动替代。degrees15只做小角度旋转超过 30 度会让钳型工具看起来像枪型工具反而引入标签噪音。mosaic0.8把四张图拼在一起训练对小目标检测有提升但装配工具本身不是小目标mosaic 过高会把细长工具的语义截断。在启动训练时把增强 YAML 传给同一个命令yolo detect train \ datatools_aug.yaml \ modelruns/detect/tools_det/baseline_n/weights/best.pt \ epochs80 \ batch16 \ imgsz640 \ patience15这一步是迁移式微调从 baseline 的 best.pt 继续训练而不是重新加载 COCO 权重可以让增强策略在已经收敛的模型上发挥作用训练时间只有第一次的一半左右。4.2 验证阶段必须看混淆矩阵而不是只看 mAP命令行里显示的平均 mAP 只能反映整体水平看不出“枪型被识别成钳型”这类具体错误。训练结束后我会单独跑一次验证并强制保存所有可视化文件yolo val \ modelruns/detect/tools_det/tools_aug/weights/best.pt \ datatools_dataset.yaml \ save_jsonTrue在验证输出目录里找confusion_matrix.png。如果gun-type tool和pliers-type tool之间存在明显的交叉误检先检查增强参数中degrees是否过大如果还是混淆就在标注层面统一标准比如同一把工具的钳口在 45 度视图中是否仍然标注为 pliers而不是把不同标注习惯的样本混在一起。另一个高频问题是在hand和工具重叠时两个框的 IoU 很高推理时 NMS 很可能把其中一个删掉。遇到这种情况我会提高iou阈值到 0.6 重新验证确认模型的原始输出是被后处理抑制还是真的漏检。4.3 迁移到 YOLOv12 前的基线对比策略如果你看到 YOLOv12 或其他新结构的发布信息想换模型再刷一轮精度我的建议是先把当前 YOLOv8 的结果保存为基线不要直接迁移。标注格式和 YAML 对 YOLOv12 同样可用因为这类模型都沿用 COCO 风格的五字段标注数据端不需要改但新结构的头模块对不同类别混淆的改善程度不一样。把两个版本的results.csv放到同一张图里分别对比各类 AP 而不是总 mAP。对于这个工具装配数据集通常值得关注的不是整体精度提高多少而是pointed tool这类细长类别是不是真的受益于新注意力机制。如果换了模型只是把 hand 的 AP 从 0.96 提到 0.97而对 pliers 毫无改善那产线收益几乎为零还要额外承担 TensorRT 插件兼容成本。5. 导出 TensorRT 并实现“手是否持握工具”的判定逻辑训练完成只是第一步实际部署到工位边缘盒子时需要把 PyTorch 权重转成 ONNX 再转 TensorRT同时在推理代码里加上工具与手的空间关系判断。5.1 导出固定尺寸 ONNX 和 TensorRT 引擎yolo export \ modelruns/detect/tools_det/tools_aug/weights/best.pt \ formatonnx \ opset12 \ simplifyTrue \ imgsz640 trtexec --onnxbest.onnx --saveEnginebest.engine --fp16opset12对 TensorRT 8.x 兼容性较好simplifyTrue会删除 ONNX 中冗余 reshape 节点。如果后续部署端的输入分辨率固定为 640x640就不需要动态 batchTensorRT 引擎会按静态输入做层融合推理速度更快。5.2 用 IoU 判定“手是否接触工具”from ultralytics import YOLO model YOLO(best.engine) def box_iou(a, b): inter_w max(0, min(a[2], b[2]) - max(a[0], b[0])) inter_h max(0, min(a[3], b[3]) - max(a[1], b[1])) inter inter_w * inter_h area_a (a[2] - a[0]) * (a[3] - a[1]) area_b (b[2] - b[0]) * (b[3] - b[1]) return inter / (area_a area_b - inter 1e-6) results model(frame_0240.jpg, conf0.35, iou0.45)[0] boxes results.boxes.xyxy.cpu().numpy() classes results.boxes.cls.cpu().numpy().astype(int) hand_boxes boxes[classes 1] tool_boxes boxes[(classes 0) | (classes 2) | (classes 3) | (classes 4)] for h in hand_boxes: if any(box_iou(h, t) 0.15 for t in tool_boxes): print(hand holding tool)conf0.35是工位摄像头这类近距离动态场景的折中选择如果现场误报多再升到 0.5。iou0.45控制 NMS 合并阈值两个工具重叠时保留置信度更高的框。box_iou里的 0.15 不是随意取的手和工具只要发生接触两个框的 IoU 通常在 0.1 到 0.35 之间低于 0.1 说明手只是靠近还没拿起来高于 0.4 则大概率是标注框重叠过大。如果你想判断“手拿着枪型工具对准工件”可以进一步要求 gun-type tool 的框同时与某个工件区域重叠这个区域可以在工位投影坐标系里预先划定。本文还有配套的精品资源点击获取
返回列表