ARTICLE DETAIL

资讯详情

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

YOLOV5齿轮缺陷检测实战:数据集制作到模型部署全指南

YOLOV5齿轮缺陷检测实战:数据集制作到模型部署全指南 简介面向毕业设计、工业质检及目标检测入门场景的YOLOv5齿轮图像缺陷检测实战项目提供从数据标注、模型训练到推理部署的完整闭环。压缩包共2000个文件包含1400个txt标签文件、442张jpeg齿轮图像、51个Python脚本及50个yaml配置另有Dockerfile、tutorial.ipynb等环境与教程辅助文件整体大小约376MB。目前已有238人学习下载适合需要快速搭建缺陷检测基线或完成课程设计的学生参考。项目附有可直接使用的数据集与预训练脚本能帮助读者理解YOLOv5的训练流程、标签格式及部署细节减少从零搭建环境的时间成本。1. 为什么齿轮图像缺陷检测要选 YOLOV5齿轮加工车间的质检工位上崩角、划痕、裂纹和缺齿这些问题很多时候还要靠人工在强光下逐个翻看。人工检得快容易漏检得细又跟不上节拍一条产线往往要安排两到三个人轮班盯。传统机器视觉方案用模板匹配或阈值分割处理规则缺陷还可以一旦遇到光照变化、油污遮挡或者复杂齿面纹理特征就不稳定了。YOLOV5 作为单阶段目标检测器把“找目标”和“分类别”统一成一个回归问题在保证 mAP 不差的前提下推理速度明显优于两阶段模型。对齿轮检测这种需要部署在产线工控机或边缘设备上的场景速度直接决定能不能在节拍内完成判定。更重要的是YOLOV5 的训练链路非常成熟数据标注、训练、验证、导出 ONNX 的流程都有现成工具工程师不需要从零搭建网络结构。这篇文章面向的是视觉算法工程师、质检自动化集成商和刚接触缺陷检测的嵌入式开发者。我按照一个完整实战项目会走的路径来展开先解决数据集从哪来、怎么整理成 YOLO 格式再训练模型、调超参数然后评估和部署推理最后给出几个最容易让齿轮检测项目翻车的排查点。文章里给的命令和参数都是可复制的你可以直接套到自己的齿轮图像上跑一遍。2. 齿轮缺陷检测数据集获取与 YOLO 格式准备2.1 开源数据集与自制采集怎么选标题里说“包含完整数据集”但真实项目里“完整”指的并不是图片多而是图片、标注、类别划分和验证集都齐全。很多公开的轴承齿轮检测数据集比如在开源数据集站点上搜“轴承 齿轮”缺陷能找到带缺陷掩膜或边界框的图片集但类别命名和标注格式各不相同有的只有裂纹有的把划痕和磨损混在一起。拿到手后第一步不是直接训练而是统一类别清单。我一般会先做一次数据体检把每个类别的样本数统计出来看类别是否平衡再随机抽 10 张图叠加边界框确认标注框和齿轮缺陷的真实位置吻合。如果数据集里图片分辨率差异很大或者有大量重复帧需要先过滤。齿轮缺陷里最常见的类别是崩角、划痕、裂纹、缺齿和毛刺工业现场还会出现黑色氧化皮和磕碰伤这要结合你们的质检标准重新定义。自制采集时不要只在实验室均匀光照下拍最好覆盖产线上实际会出现的背光、强反光和油污状态。数据集完整性的另一个隐含要求是“训练集/验证集/测试集按同一分布切分”不能把同一个齿轮的连续多帧同时放进训练集和验证集否则验证指标会虚高。切分前先用文件名排序并做去重保留序列间的间隔。2.2 用 LabelImg 标注并导出 YOLO txt 格式标注工具我用得最多的是 LabelImg它支持 YOLO 格式直接导出。启动后选择 Open Dir 打开齿轮图片目录Change Save Dir 设置标注保存目录然后在 Edit 菜单里把标注格式切到 YOLO。框住一个缺陷后输入类别名保存后每张图片对应的同名 txt 文件里就会生成一行数据格式是class x_center y_center width height四个坐标值都是归一化到 0~1 的浮点数。# 安装并启动 labelimg pip install labelimg labelimg上面的命令只做了一件事通过 pip 安装 labelimg 工具包然后启动图形界面。启动后注意右下角要选择 PascalVOC 或 YOLO 格式如果选了 PascalVOC 导出的是 xml后面还要再转一步。为了省事我在标注前会先创建一个classes.txt文件里每一行写一个类别名顺序必须和后续训练用的 data.yaml 一致。label 0, score 0.xx 的类别索引到底对应什么YOLOV5 在训练时会自动读取 data.yaml 里的类别名和标注文件里的数字索引一一对应。.常见的一个坑是标注框把整张图都框住或者框有一部分在图像外面。YOLOV5 训练时会报 “all labels empty” 或者坐标出界的警告但没有强制报错。保险起见标注完成后用下面这段 Python 脚本扫一遍所有 txt把坐标小于 0 或大于 1 的文件列出来。import os label_dir labels for name in os.listdir(label_dir): if not name.endswith(.txt): continue with open(os.path.join(label_dir, name)) as f: lines f.readlines() for line in lines: parts line.strip().split() if len(parts) 5: x, y, w, h map(float, parts[1:]) if not (0 x 1 and 0 y 1 and 0 w 1 and 0 h 1): print(f越界标注: {name}: {line.strip()})这段代码遍历标注目录下的所有 txt读取每一行并根据第五个字段之前的四个坐标判断是否在 0 到 1 范围内。只要有一个值超出边界就会打印出对应文件名和原始行。齿轮缺陷的边界框通常都比较小如果坐标出现 1.0 或 0.0大概率是标注时鼠标拖出了图像边缘手动修正后再训练能省掉很多排查时间。2.3 数据划分与目录结构标注完成后目录结构我建议按 YOLOV5 默认习惯组织这样不用改代码路径。gear_defect/ ├── images/ │ ├── train/ │ ├── val/ │ └── test/ ├── labels/ │ ├── train/ │ ├── val/ │ └── test/ └── data.yaml每个训练图片必须有一个同名的 txt 文件放在 labels 对应目录里否则 YOLOV5 训练时会把没有标签的图片当作背景样本导致漏检增多。划分比例我用 8:1:1即 80% 训练、10% 验证、10% 测试。如果齿轮缺陷样本只有几百张我会把验证集固定留出 100 张左右而不是机械地按比例切。下面是一个干净的划分脚本片段。import random from pathlib import Path img_paths list(Path(images_all).glob(*.jpg)) random.shuffle(img_paths) train, val, test img_paths[:800], img_paths[800:900], img_paths[900:1000] for split, paths in [(train, train), (val, val), (test, test)]: for p in paths: # 复制图片到 images 对应目录 # 同时复制同名 txt 到 labels 对应目录 print(split, p.name)注意脚本里我只是演示划分逻辑实际使用时要把复制文件的代码补全。关键是先 shuffle 再切片避免文件按日期排序导致同一个批次全部落在训练集。复制完成后在data.yaml里指定图片目录和类别名这个文件的内容在第 3 章会给出。表格里列出齿轮缺陷检测常用的类别定义类别名缺陷定义说明chipped_teeth崩角/缺齿齿轮边缘出现小块缺失scratch划痕齿面或端面的条状损伤crack裂纹细长线状裂缝burr毛刺加工后残留的金属突起3. 用 YOLOV5 训练齿轮缺陷检测模型环境与超参数3.1 环境配置与代码准备YOLOV5 对 Ubuntu 和 Windows 都支持我一般建议用 Ubuntu 20.04 配合 Python 3.8 和 PyTorch 1.10 以上的组合。先拉取官方代码库创建虚拟环境并安装依赖。git clone https://github.com/ultralytics/yolov5.git cd yolov5 python -m venv venv source venv/bin/activate pip install -r requirements.txt这四行命令分别完成了从官方仓库克隆 YOLOV5 源码到当前目录、进入代码目录、创建并激活虚拟环境、安装 requirements.txt 里列出的所有依赖包。requirements 里包含 torch、torchvision、opencv-python 等基础库GPU 版本需要在安装 torch 前单独指定 CUDA 版本比如pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118。如果机器没有 GPU也可以跑 CPU 训练但齿轮图像分辨率高的话速度会很慢建议至少用 6GB 显存的显卡。环境配置里最容易出问题的是 PyTorch 和 CUDA 版本不匹配。装完后用python -c import torch; print(torch.cuda.is_available())检查一次输出 True 再继续。否则训练时会落到 CPU迭代速度慢很多但不容易发现。3.2 编写 data.yaml 并启动训练进入项目目录后需要手动创建一个data.yaml指向你的齿轮数据集。train: /home/user/gear_defect/images/train val: /home/user/gear_defect/images/val nc: 4 names: [chipped_teeth, scratch, crack, burr]train 和 val 两个字段是图片目录的绝对路径YOLOV5 会自动到同级的 labels 目录下寻找对应的标签文件。nc 是类别总数names 是类别名列表顺序必须和标注文件里的数字索引一致。很多项目坚持“数据完整”但训练时报错原因就是 labels 里的类别索引在 names 里找不到对应项。训练命令我推荐先跑一个快速实验python train.py --data data.yaml --weights yolov5s.pt --img 640 --batch 16 --epochs 100 --name gear_exp1这里--weights yolov5s.pt表示使用 YOLOV5s 预训练权重做迁移学习--img 640把训练图片统一缩放到 640x640--batch 16是每批训练样本数--epochs 100是训练轮数--name gear_exp1指定输出目录名训练曲线和权重会保存在runs/train/gear_exp1下。第一次训练不建议直接换大模型先用 small 版本跑通流程确认损失能下降再升级到 medium 或 large。齿轮缺陷里像划痕和裂纹这种细长目标直接缩放到 640 可能丢失细节。如果原图分辨率是 2048x1536我会把--img 1280或--img 1536作为第二组对比实验同时调小 batch 防止显存溢出。即使 YOLOV5 有自适应锚框大的输入尺寸对细小缺陷的召回率提升依然明显。3.3 超参数与锚框调整YOLOV5 默认的锚框是基于 COCO 目标分布设计的齿轮缺陷是小目标居多直接用默认锚框也能训练但收敛速度会受影响。训练前可以用--analyze参数检查数据集的锚框聚类结果也可以在训练脚本里开启--autoanchor。我实际项目中会把超参文件data/hyps/hyp.scratch-low.yaml里和损失权重相关的项单独调。下表是我在齿轮缺陷检测上常用的初始超参数参数名推荐值作用lr00.01初始学习率类别不平衡时可降到 0.005lrf0.1学习率衰减到初始值的比例imgsz1280输入尺寸影响小缺陷检出batch8batch 越大收敛越稳定但显存占用也越高epochs200小数据集容易过拟合配合早停用mosaic1.0增强训练样本多样性齿轮数据量少时建议保留fl_gamma1.5Focal loss 参数解决背景占比大导致的漏检设置--hyp可以指定自定义超参数文件。如果发现训练 loss 震荡大先把 lr0 降到 0.005再把 fl_gamma 提高到 1.5背景占比过高时这一项对召回率影响很大。同类别样本太少时还可以在超参文件里增加hsv_h、hsv_s这些颜色增强模拟车间里不同的光照和油污效果。训练日志里重点关注 val/box_loss 和 val/cls_loss 这两条曲线如果训练损失一直下降但验证损失在某个 epoch 开始上升可以提前停止并保存之前的最优权重。4. 模型评估、推理与导出4.1 验证集指标与混淆矩阵怎么看训练完成后YOLOV5 会在runs/train/gear_exp1下生成 weights 和验证结果。用val.py单独在测试集上评估一次确认模型的泛化能力。python val.py --data data.yaml --weights runs/train/gear_exp1/weights/best.pt --img 1280 --batch 8执行后会打印出 mAP0.5、mAP0.5:0.95、Precision、Recall 四个指标并画出归一化混淆矩阵。齿轮缺陷里最需要关注的不是整体 mAP而是裂纹这一类在低 IoU 下的 Recall 值。如果有大量裂纹没被召回问题往往出在输入尺寸或锚框而不是分类头。混淆矩阵图片上如果某个类别被大量预测成背景说明该类别样本太少或被遮挡严重。我一般会回到标注环节把漏检的图片单独提出来看它们的缺陷尺寸占比。齿轮齿面边缘的崩角通常很小在 640 尺寸下可能只有 5 个像素这种目标即使训练收敛推理时也容易丢。4.2 用训练好的模型做图像和视频推理推理命令和训练类似但不需要标签文件直接跑 detect.py 就可以。python detect.py --weights runs/train/gear_exp1/weights/best.pt --source /path/to/test_images --img 1280 --conf 0.25 --iou 0.45--source可以指向一张图片、一个目录或一段视频。--conf 0.25是置信度阈值低于这个值的检测框会被过滤--iou 0.45是 NMS 的 IoU 阈值用来合并重叠框。齿轮检测中如果过度遮挡导致重叠框多把 iou 提高到 0.6 可以减少误删。推理结果会保存到runs/detect/exp每张图会在原图上画了边界框和类别名。对视频流做在线检测时建议把--source改成视频设备编号比如 0同时开启--view-img实时显示检测画面。如果推理帧率达不到产线要求优先降低输入尺寸比如从 1280 降到 960再考虑是否换更轻的模型。齿轮缺陷虽然小但很多产线机构会先引导相机到指定位置停稳再拍照静态检测对推理时间的要求并不苛刻。4.3 导出 ONNX/TensorRT 并部署到边缘设备项目落地时往往需要把模型从 PyTorch 导出成 ONNX再转成 TensorRT 引擎部署到 Jetson Nano 或工控机。python export.py --weights runs/train/gear_exp1/weights/best.pt --include onnx --opset 12导出后得到best.onnx可以用 ONNX Runtime 或 TensorRT 加载。如果部署设备是 Jetson Nano推荐直接导出 TensorRT 引擎YOLOV5 官方代码支持--include engine但需要在 Jetson 上先安装对应版本的 TensorRT。ONNX 导出的关键参数是--opset太高的算子版本在旧版 TensorRT 上可能不支持12 是兼容性比较好的选择。导出的 ONNX 模型在推理前最好用onnxruntime跑一遍同一个输入和 PyTorch 的输出对比一下差异。齿轮缺陷检测对边界框坐标精度要求没那么严格一般中心点偏差在 2 个像素内就可以接受。5. 齿轮缺陷检测项目里 5 个容易被忽略的排查点5.1 细小裂纹漏检先做切片检测裂纹在 640 缩放下只有几个像素直接检测很难。常见的做法是把原图切成 4 个或 9 个 tiles每个 tile 单独送进模型最后把检测框坐标映射回原图。YOLOV5 的 detect.py 本身不支持切图需要写一个小脚本把图片切块并合并结果。切块会带来重复检测对同一个位置的不同 tile 用 NMS 合并一次即可。5.2 标注框越界导致训练不收敛第二点在第 2 章提过这里再强调一次。YOLOV5 对越界标注只是警告不会中断训练。但越界框会影响 anchor 匹配和目标损失计算让 mAP 在某个值上卡住。用我前面提供的扫描脚本跑一遍把所有越界 txt 修正或删掉再重新训练往往能直接提升 3 到 5 个点。5.3 类别不平衡集中在毛刺很多齿轮数据集里毛刺样本可能是裂纹的 10 倍。YOLOV5 的 cls loss 默认对每个类别等权样本多的类别会压制样本少的类别。解决方式有两种一是用采样策略让每个 epoch 里每个类别出现的次数接近二是把fl_gamma调高让模型更关注难分的裂纹样本。5.4 推理时光照和油污干扰训练集如果全是干净光照推理时遇到反光就会误检。我一般会在训练阶段把图像保存到train/之前加一层随机亮度增强或者在超参里调大hsv_v。最终部署时在相机镜头前加偏光片能显著减少金属表面的反光这也是成本最低的优化。5.5 验证集里出现了训练集同一齿轮的不同帧齿轮缺陷项目的数据集经常来自视频抽帧同一个齿轮的连续帧很相似。如果切分时没有按文件名前缀分组验证集会变相变成训练集的一部分看到的 mAP 会比实际高很多。推荐按齿轮 ID 或批次号分组切分而不是按单帧随机切。调试时在数据目录里放一个sample_check.py遍历验证集文件名并检查原始序列分组能快速发现这类泄露问题。最后一个实用技巧是打开验证集预测结果图把置信度阈值从 0.25 调到 0.5观察漏检目标在 0.5 附近的数量。如果大量漏检目标集中在 0.3 到 0.5 之间说明模型定位正确但分类置信度不足优先优化输入尺寸和焦点损失参数而不是简单降低阈值兜底。本文还有配套的精品资源点击获取
返回列表