ARTICLE DETAIL

资讯详情

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

YOLOv5钢材表面缺陷检测实战:数据集转换与训练部署全指南

YOLOv5钢材表面缺陷检测实战:数据集转换与训练部署全指南 简介YOLOv5结合钢材表面缺陷数据集是一套面向工业质检与智能制造场景的深度学习资源适合计算机视觉入门者或缺陷检测工程师解决钢材表面缺陷的自动识别与定位问题。压缩包约501MB共2000个文件其中1971张jpg缺陷图像配有1954个txt标注和1800个xml标注txt文件保存YOLO格式检测框信息xml文件可用于VOC格式训练便于在主流检测框架间切换此外还提供yaml配置、python训练脚本、pt预训练权重和onnx模型完整覆盖数据、训练与推理链路。资源内置YOLOv5网络架构涉及CSPNet、SPP-Block、FPN和PANet等组件提升多尺度目标检测效果可针对裂纹、凹坑、氧化皮、划痕等典型表面缺陷进行训练与验证。已有2497人学习下载借助全套数据与脚本读者可快速复现缺陷检测流程并进一步通过数据增强、参数调优和模型剪枝提升性能或迁移到其他工业视觉任务中。1. 为什么钢材表面缺陷检测都绕不开 YOLOv5做钢材表面缺陷检测的朋友大概率都听过这个组合——YOLOv5 加上公开的钢材缺陷数据集尤其是 NEU 那套。它不是最炫的方案却是能在两三天内跑出结果、让甲方看到界面的最稳路子。钢材厂里常见的裂纹、夹杂、麻点、氧化皮、压入氧化皮、划伤这六类缺陷靠人工肉眼盯产线漏检率压不下来上了传统图像处理又被光照和背景折腾得够呛。目标检测模型里YOLOv5 在速度和精度上卡在一个很舒服的位置GTX 1060 级别的卡就能带起来产线级推理也能跑到几十帧。这篇就按我自己摸索过的路径从数据集格式转换、环境配置到训练参数和部署落地把能抄的作业和踩过的坑一起摊开讲。2. 钢材表面缺陷数据集从 VOC 到 YOLO 的转换与边界坑2.1 NEU-DET 数据集的构成与格式陷阱公开钢材表面缺陷数据集里东北大学的 NEU-DET 是大多数人第一步用的。它包含 1800 张热轧带钢表面灰度图每张 200×200 像素标注了六类缺陷——crazing裂纹、inclusion夹杂、patches麻点、pitted_surface氧化皮、rolled-in_scale压入氧化皮、scratches划伤每类 300 张。虽然学术界对这套数据集的标注质量有过不少讨论但作为算法验证和产线预演它仍然是可复现性最好的起点。拿到原始包之后第一个翻车点通常是标注格式。NEU-DET 的官方标注是 VOC 风格的 XML一个缺陷目标一个object块里面是bndbox标签。而 YOLOv5 训练时读的是 txt 格式的标注每行一个目标格式为class x_center y_center width height全部归一化。很多新手直接把 XML 扔进 YOLOv5 的 labels 目录报错image not found或label格式错误就是因为没做这一步转换。另一个容易被忽略的坑是类别索引。VOC 的name字段是字符串crazing而 YOLO 训练需要从 0 开始的整数类别 id。转换脚本里如果没有建立类别到 id 的映射或者映射顺序和训练 yaml 里的names列表不一致模型训练出来所有框都错位推理时类别标签会全部张冠李戴。2.2 实现一个干净的 VOC 转 YOLO 转换脚本我一般会在项目根目录放一个voc2yolo.py用 ElementTree 解析 XML同时做归一化和类别映射。这里直接给出一个经过验证的版本import os import xml.etree.ElementTree as ET from pathlib import Path # 类别字典顺序必须与训练 yaml 的 names 一致 CLASS_MAPPING { crazing: 0, inclusion: 1, patches: 2, pitted_surface: 3, rolled-in_scale: 4, scratches: 5 } def convert_voc_to_yolo(xml_path, out_dir, img_width200, img_height200): 将单张 VOC XML 转为 YOLO txt 标注并保存到 out_dir tree ET.parse(xml_path) root tree.getroot() txt_path Path(out_dir) / (Path(xml_path).stem .txt) lines [] for obj in root.findall(object): name obj.findtext(name).strip() # 跳过未在映射表里的类别 if name not in CLASS_MAPPING: print(f[warning] 未知类别: {name}, 文件: {xml_path.name}) continue class_id CLASS_MAPPING[name] bndbox obj.find(bndbox) xmin float(bndbox.findtext(xmin)) ymin float(bndbox.findtext(ymin)) xmax float(bndbox.findtext(xmax)) ymax float(bndbox.findtext(ymax)) # 计算 YOLO 需要的中心点 宽高并做归一化 x_center (xmin xmax) / 2.0 / img_width y_center (ymin ymax) / 2.0 / img_height width (xmax - xmin) / img_width height (ymax - ymin) / img_height # 过滤无效框坐标越界或宽高为 0 的目标直接舍去 if width 0 or height 0: print(f[warning] 无效框: {xml_path.name} - {name}) continue line f{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f} lines.append(line) # 如果没有有效目标不生成空 txt避免 YOLOv5 误判为背景图 if lines: with open(txt_path, w) as f: f.write(\n.join(lines) \n) def batch_convert(xml_dir, out_dir): os.makedirs(out_dir, exist_okTrue) xml_files list(Path(xml_dir).glob(*.xml)) for xml_file in xml_files: convert_voc_to_yolo(xml_file, out_dir) print(f转换完成共处理 {len(xml_files)} 个 XML 文件) if __name__ __main__: batch_convert(path/to/annotations, path/to/labels)这段脚本的核心逻辑是先建立类别映射表保证 id 从 0 开始然后从 XML 里取出左上角和右下角坐标换算成中心点加宽高再分别除以图像宽高得到归一化值。归一化是关键YOLOv5 在训练时会按照标注的相对位置来学习偏移量如果这里忘了除以宽高模型会学出一堆负坐标甚至 NaN。转换时有几个边界情况要特意处理。首先NEU-DET 的 XML 里所有图片都是 200×200但有些开源版本把每张图的尺寸写在size标签里直接用根节点下的size字段更稳妥避免硬编码。其次空标注文件不要保留。YOLOv5 会把没有 txt 的图片自动视为背景但一个仅含背景的图如果在后续训练中被强行划入训练集相当于引入一个无目标的负样本如果是小数据集会稀释正样本的浓度。最后坐标越界的框要过滤。XML 人工标注偶尔会把框出到图像边界外YOLOv5 读取时能容忍但训练时 loss 会忽高忽低。2.3 数据集划分别让随机切分毁掉你的评估指标把转换好的图片和标签对称放置后还有一个看似简单却经常翻车的问题——数据集划分。很多人直接用os.listdir随机切 8:1:1却不检查同一类缺陷是否均匀落在各个集合里。NEU-DET 每类 300 张如果随机切分时某类的训练集只剩 200 张、验证集却分了 80 张mAP 的波动会非常大小样本类别的 AP 直接崩掉。常见做法是按类别做分层抽样代码逻辑不复杂先按缺陷类别把图片路径分组再从每个组内各自做 8:1:1 切分最后合并。YOLOv5 的数据加载是读图片列表的因此还要把训练和验证的图片路径分别写进两个 txt 文件或者在 dataset.yaml 里指定train和val指向包含图片的目录。我自己的习惯是直接生成两个 txt这样换模型、换训练脚本时不用重新划分。另一个容易被忽略的问题是图片与标签的重名对应。yolov5 在训练时通过图片路径直接替换后缀找同名 txt如果你转换时把图片放在images/train、标签放在labels/train但目录名和 yaml 配置不一致会报 AssertionError: train: 找不到标签文件。两个目录结构必须完全镜像yaml 里train指向 images 目录YOLOv5 会自己推导 labels 目录。3. YOLOv5 环境配置与源码结构先跑通再谈调参3.1 从零搭一套能跑的训练环境YOLOv5 的环境配置不复杂但版本坑很多。我第一次配的时候直接pip install -r requirements.txt结果 PyTorch 装上 CPU 版本训练一个 epoch 要 40 秒根本没法调参。正确的做法是先装 PyTorch 再装其他依赖顺序反过来容易让 pip 把 torch 替换成默认的 CPU 包。常见做法是 Python 3.8 以上、PyTorch 1.8 到 2.x 均可。yolov5 官方仓库对 PyTorch 的版本要求比较宽容但如果你用的是 NVIDIA 显卡CUDA 版本和 torch 版本必须匹配。有个简单的验证方法装完后跑python -c import torch; print(torch.cuda.is_available())输出 False 就不要硬着头皮训练否则后面所有时间都是浪费。依赖安装方面pip install -r requirements.txt会拉取 numpy、opencv-python、matplotlib、pyyaml、tqdm、requests 等基础库。其中 opencv-python 装完后如果系统里本来就有其他版本建议确认下import cv2能正常执行。YOLOv5 的图像读取走 cv2opencv 版本过老会导致某些图像解码异常训练时会出现 unable to load image 的提示。3.2 数据 yaml 与模型 yaml两处必须改的配置文件跑通源码后先别急着训练有两个配置文件必须自己建。第一个是数据集 yaml内容很简单就是指定训练和验证的图片目录以及类别名称列表# steel_defect.yaml train: ./datasets/steel/images/train val: ./datasets/steel/images/val nc: 6 names: [crazing, inclusion, patches, pitted_surface, rolled-in_scale, scratches]第二个是模型 yaml。你不需要从头写网络结构直接复制官方提供的yolov5s.yaml把nc改成 6 即可。但要注意模型 yaml 里的nc和数据 yaml 里的nc必须一致否则训练时会在 head 部分报维度错误。这个报错信息通常很长新手容易蒙本质就是最后一层卷积输出的通道数和你的类别数不匹配。修改 nc 之后如果用预训练权重启动--weights yolov5s.ptYOLOv5 会自动把模型头部的卷积层通道数调整成对应类别数所以不需要手动修改 head 部分也不要怀疑是 yaml 写错了。但是如果你用的预训练权重是从别处拷贝来的、不是官方仓库提供的存在层名不一致导致加载报错的可能。常见做法是首次训练直接用官方发布的yolov5s.pt等自己的模型训出来再增量训练。3.3 YOLOv5 的后处理机制NMS 和类别过滤YOLOv5 的后处理阶段官方的做法是把三组特征图分别对应小、中、大目标的输出做多标签融合再用 NMS非极大值抑制去掉重叠框。很多做钢材缺陷检测的朋友以为后处理只是调一个 conf_thres但实际还有一个容易忽略的点——不同类别的缺陷在尺寸上差异很大裂纹和划伤往往是长条细线麻点和氧化皮偏向小而密集如果 NMS 的 IoU 阈值设太高小目标框会被大目标误杀。代码层面推理时detect.py里有一个--conf-thres参数控制置信度阈值--iou-thres控制 NMS 的 IoU 阈值。钢材表面缺陷里氧化皮这类目标边缘模糊置信度普遍偏低我把过 conf_thres 从默认的 0.25 降到 0.15确实多召回了一部分缺陷但代价是误检增加。这里没有玄学只有循环先跑一组数据看 PR 曲线再决定阈值方向。4. 用 YOLOv5 训练钢材表面缺陷超参数与数据增强实战4.1 训练命令与关键超参数选择数据准备好、yaml 写好后训练命令本身就没什么神秘的了。我通常用这样一条python train.py --img 640 --batch 16 --epochs 150 --data steel_defect.yaml --cfg yolov5s.yaml --weights yolov5s.pt --name steel_defect_exp1几个参数值得展开说。--img 640是把原始 200×200 的输入缩放到 640×640。NEU-DET 原图只有 200×200直接拉大到 640 会丢失一部分细节但 YOLOv5 的 backbone 在 640 输入下表现最稳定所以通常不直接用小图训练。如果你追求极致速度可以试试--img 320mAP 大概会掉 3 到 5 个点但推理时间能省将近一半。--batch 16是我在单张 8GB 显存卡上的选择。如果你的卡是 24GB可以尝试 32 或 64。批大小对收敛影响很大但显存不够时不要硬开优先保证 batch 不小于 8。--epochs 150对于 1800 张这种小数据集其实是偏多的通常跑到 80 个 epoch 左右验证集 mAP 已经到平台期多的 epoch 更多是让模型自己把尾巴磨干净。4.2 超参数文件里的三个关键旋钮YOLOv5 的超参数在data/hyps/hyp.scratch-low.yaml和hyp.scratch-high.yaml里官方默认的 low 配置偏保守适合在数据量不多时防止过拟合。我从不直接改文件里所有的参数只动三个最敏感的学习率lr0、Mosaic 概率mosaic和图像上下翻转fliplr。钢材表面缺陷和自然图像数据集有个明显差异——很多目标具有方向性。划伤和裂纹大多是水平方向延伸的上下翻转影响不大但左右翻转等于把语义结构保留的样本变了一个镜像如果数据里划伤框偏左镜像后会变成偏右模型反而能学得更稳。fliplr默认是 0.5对钢材缺陷数据集我一般保持默认。lr0是初始学习率小数据集上我一般从 0.01 改成 0.005让模型学得更稳。mAP 到了 40 个 epoch 还没动静先看 loss 曲线还在不在降如果 loss 已经平了再考虑手动调低学习率重启训练。YOLOv5 的调度器是余弦退火即使中途不干预它自己也会慢慢降所以我更建议对新手只调lr0剩下的交给源码里默认的 warmup 和 scheduler。4.3 数据增强对缺陷检测的实际影响Mosaic 增强是 YOLOv5 训练里默认开启的利器它对钢材缺陷这类目标数量少、背景重复度高的数据尤其有用。Mosaic 会一次性拼接 4 张图相当于把多个缺陷样本放进同一张训练图模型需要同时在多尺度上学习。但副作用也明显拼接后的边缘会出现伪目标比如一张图的裂纹和另一张图的划伤在接缝处混叠模型可能学到接缝处的假特征。官方代码里mosaic默认概率高达 1.0很多项目直接跑也不出问题但如果发现验证集上误检特别多可以把 mosaic 调到 0.5 试一试。另一个值得注意的是copy_paste增强在 yolo 5.x 后期版本中加入。它会把一张图的标注框复制粘贴到另一张图上对缺陷检测来说这相当于合成数据增强但钢材缺陷的纹理依赖背景直接抠出来贴到其他表面往往会出现不自然的边界模型在验证集上反而掉点。我的经验是保持默认关闭别追新功能。5. 训练与部署中的避坑记录四个真实翻车现场5.1 类别标签错位mAP 很高预测全是错的现象训练 100 个 epoch验证集 mAP 能到 0.7 以上但推理时每张图框的位置大致对类别却一塌糊涂——裂纹框里写着夹杂麻点框里标着氧化皮。原因这是数据集 yaml 里的names顺序和转换脚本里的类别映射表不一致导致的。比如转换脚本里crazing映射成 0但 yaml 里names列表第 0 位是inclusion模型按索引学标签本身已经张冠李戴。解决把所有 ID 映射关系统一到一张表里。我后来的习惯是转换脚本里直接读取 yaml 文件解析names列表按位置自动生成类别映射不再手动写字典。这样永远不可能错位。5.2 小目标缺陷漏检严重锚框和输入尺寸的问题现象裂纹、划伤这类细长缺陷在训练集上有标注、有检出但推理时同一张图换成 1280 分辨率推理能检出来640 分辨率就丢掉一半。原因YOLOv5 默认的锚框是根据 COCO 数据集聚类出来的COCO 里的大目标是人和车锚框偏大。NEU-DET 的缺陷框很多是 20×60、15×40 这种窄长条默认的锚框对这个尺度覆盖不足小目标层P3的分配率太低模型学不到足够的小目标特征。解决一种做法是训练时加上--multi-scale让模型在 320 到 640 之间随机缩放间接增加小目标样本比例。另一种做法是在训练前用脚本对真实标注框做 k-means 聚类生成新的锚框写进模型 yaml 的anchors字段。聚类脚本不用自己写YOLOv5 源码的utils/autoanchor.py里有现成实现训练时会自动检查锚框和标注的匹配度不适合会在终端打警告。5.3 训练 loss 正常但 mAP 一直徘徊类别不平衡和困难样本现象训练集 loss 掉得很漂亮从 0.1 掉到 0.02验证集 mAP_0.5 却一直卡在 0.5 左右不再上涨。打印每个类别的 AP发现 pitted_surface 只有 0.1其他类都接近 0.8。原因NEU-DET 虽然每类 300 张图但目标数量差异很大。比如 rolled-in_scale 这类缺陷在单张图上可能有四五个目标而 crazing 经常只是一条细线。模型对多目标类别学得更充分对少目标类别几乎放弃。加上验证集上每类的图片数不均匀mAP 被拉低。解决一种方法是在 yaml 里做类别重加权但 YOLOv5 原生不支持实际中用得多的做法是数据增强偏重困难类别——把少目标类别的图在训练循环里复制一份倍率或者用开源库做离线过采样。另一个实际有效的办法是调整置信度阈值和 IoU 阈值pitted_surface 这类目标置信度低验证时降低阈值后 AP 会好看一些但这是自欺欺人部署时还是得靠模型本身。5.4 部署时推理速度慢用了 GPU 但没吃到满现象训练时一切正常导出成 ONNX 用 ONNXRuntime 推理发现 CPU 推理一次要 300msGPU 只降到 120ms跑不满产线要求的 30fps。原因源码里的detect.py包含了图像预处理、后处理、坐标换算和画框逻辑导出 ONNX 时只导出了模型 forward 部分但前处理和后处理仍是 Python 代码在跑。Python 的 NMS 部分在批量检测时会成为主要瓶颈。解决推理时用 YOLOv5 官方的export.py导出 ONNX并开启--simplify做模型精简后处理部分要么自己用 numpy 实现 batch NMS剪掉 torch 的 dynamic 图开销要么直接用 OpenVINO 的 runtime它自带的 NMS 层能省掉大部分 Python 逻辑。如果项目允许用 TensorRTFP16 精度下推理时间可以压到原来的三分之一。6. 验证与进阶从 mAP 数字到产线可用的最后一步模型训练完别急着导出部署。先做一次老样本的回归测试——拿几张训练过程中没有见过的、现场实际采集的钢材图片跑一次推理看看框的位置和人眼判断是否一致。mAP 是统计指标但产线更在意的是单张图上漏检和误检的比例。我自己的习惯是统计 200 张图上每类缺陷的漏检数漏检率超过 3% 的那类先在测试集上单独看它的 PR 曲线确认是召回低还是置信度阈值调太高。进阶层面钢材表面缺陷数据集的原始分辨率通常远高于 YOLOv5 默认输入尺寸。如果你的产线相机是 1200 万像素直接把原图缩到 640 会丢掉大量小目标细节。常见做法是切块推理——把原图切成 640×640 的 patch给每个 patch 加 10% 的 overlapping最后合并时做 NMS 去重。这个方案缺陷检测行业里非常常见mAP 能比直接缩放高出 5 个点以上。部署载体上我在实际项目里用过两种一种是训练好后直接转 ONNX 给工控机跑 ONNXRuntime适合产线已有 Windows 上位机的情况另一种是树莓派这类边缘设备上次在 ARM 板上跑自己训练的 YOLOv5 模型量化成 int8 之后推理延迟能控制在 80ms 内。但在 ARM 板上要认真检查逐层算子支持情况有些升采样层要人工替换。说到底YOLOv5 加钢材缺陷数据集这套方案最大的优势是生态完整、资料多、遇到问题能在三分钟内搜索到同路人。这也是我不建议新手一上来就换 YOLOv8 或 v11 的原因——模型可以换代但数据格式转换、锚框分析、阈值调优这套方法论才是真正通用的。希望这篇笔记里的数据和踩坑过程帮到你尤其是那些看起来不起眼、却能让人熬夜两天的格式和版本问题。本文还有配套的精品资源点击获取
返回列表