YOLOv5钢轨表面缺陷检测实战:从数据集构建到模型部署全流程 钢轨表面缺陷检测的项目任务我前后做了大半年用的就是YOLOv5这套方案。标题里写的“数据集含训练好的模型”听起来很简单但真正落地的时候从数据采集、标注规范到训练参数调整、最后部署推理每一步都藏着不少坑。这篇文章就把整个过程完整拆开讲清楚我为什么选YOLOv5、数据集怎么建才够用、训练时哪些参数必须抠、以及最后模型怎么用起来。这个话题适合谁看如果你手头正好有钢轨、钢板、管道这类长条形工业表面的缺陷检测需求或者你想知道一个YOLOv5目标检测项目从零到落地的完整链路这篇可以当一份“抄作业”的参考。哪怕你对计算机视觉不太熟只要按着思路走也能把一个能用的检测模型跑起来。1. 项目背景与任务拆解1.1 钢轨表面缺陷检测到底在解决什么问题钢轨长时间承受列车碾压表面会出现擦伤、剥落、裂纹、压痕等多种缺陷。这些缺陷如果发展到一定程度直接影响行车安全所以日常巡检里必须盯着看。早年都是靠人工徒步巡检老师傅用眼睛看、用手摸效率低不说碰上雨雪天气或者夜间天窗时间漏检概率也不小。后来有了轨道巡检车架上高速相机连续拍摄一趟跑下来能拍几万张图像。问题就从“怎么拍”变成了“怎么看”海量图像如果继续靠人看眼睛很快就疲劳了漏检率反而上去了。这个项目要解决的就是让算法替人眼完成第一轮筛查——把图像里疑似缺陷的区域框出来再由人工复核这样既保住了检测质量又把人力释放了出来。这里的核心需求拆开其实是三块对钢轨表面图像中的缺陷目标进行定位输出带边界框的检测结果检测速度要足够快最好能在巡检结束后的短时间内处理完当天数据检测结果要稳定不能因为光照变化、表面油污、锈迹干扰就频繁误报。这三块需求放在一起就决定了方案选型的方向得用目标检测模型而不是单纯的图像分类或者分割。1.2 缺陷类型与标注难点钢轨缺陷不是单一形态。我实际接触过的几类典型缺陷在图像里的表现差异很大缺陷类型图像特征标注注意事项擦伤轨面沿纵向出现亮白发亮的条状区域常伴随表面金属拉伤痕迹框要包住整个发亮区域避免只框一半剥落表面出现块状凹陷颜色偏暗边缘不规则按块状区域的外接矩形标不要拆成多块裂纹细长线状方向多沿钢轨纵向或斜向这是最难标的细长目标框必须尽量贴合压痕周期性出现的凹坑间距相对均匀每个凹坑单独标注不要合并成一个框锈蚀面积较大的变色区域边界模糊容易误标需要结合纹理判断标注难点主要在两个地方。一是裂纹这类细长目标边界框稍微带歪一点IoU就降得很难看模型训练时学到的特征也跟着歪。二是剥落和锈蚀边界模糊不同标注员理解不一样框的大小能差出一大截。这块除了统一标注规范还得反复校准。1.3 为什么选YOLOv5而不是其他模型当时对比过几个主流方案不只有YOLOv5。老牌的Faster R-CNN检测精度不错但推理速度慢处理巡检车拍回来的大批量图像明显吃力。YOLOv3跑得快但小目标和密集目标的表现一般。YOLOv5在当时已经相对成熟工程生态也完善预训练权重、数据增强策略、部署导出工具链都齐全这点对实际项目很重要。选YOLOv5有几个非常现实的原因速度和精度平衡得好在工业场景里这是命根子官方仓库开箱即用训练、验证、导出、推理都有现成脚本省掉了大量造轮子的时间模型体积可控v5s、v5m、v5l几个规格可以按算力选从GPU服务器到边缘盒子都有对应方案社区活跃踩坑记录多遇到问题基本都能搜到解法。后来实际用下来v5s在钢轨数据上已经能做到满足现场需求的精度而单张推理时间在普通显卡上也就是毫秒级这套选型是站得住脚的。2. 数据集构建全流程2.1 图像采集的几个硬性要求数据集质量直接决定模型上限。模型结构再先进喂进去的数据乱七八糟最后训练出来的就是“乱糟糟的模型”。采集阶段我总结了几条硬性要求。第一覆盖场景要全。钢轨有白天拍的、夜间拍的、阴雨天拍的有干燥表面也有带水渍的。光照方向不同缺陷呈现出来的对比度完全不同如果训练数据里只有单一光照模型一上现场就“见光死”。采集时要有意识地分时段、分天气、分线路去拍。第二分辨率要保证。钢轨表面缺陷不少是非常细小的图像分辨率太低细节模糊了标注都标不准。相机拍摄距离要稳定确保目标区域在图像里足够清晰。第三原始图像尽量不压缩。现场传图时如果经过有损压缩细小裂纹的边缘会变得很糊对训练非常不利。我自己踩过这个坑最开始图方便收了一批压缩图训练出来的模型在真实图像上漏检率明显变高后来重新采集了无损图才算解决。2.2 标注规范与工具选择标注工具我用的是LabelImg操作简单直接输出YOLO格式的txt文件省去转换步骤。如果后续有做分割的打算可以换成LabelMe但从检测项目角度LabelImg够用了。标注规范必须提前定死不然几个人标出来的数据五花八门。我当时的规范是这么定的所有缺陷目标用矩形框紧贴目标外沿做到不压目标、不留大边裂纹这类细长目标框的四个边尽量贴近缺陷端点长宽比再大也要保持同一图像中同类型缺陷互相独立不能合并成一个框无法判断的模糊目标宁可放弃也不硬标避免给模型喂噪声目标只占几个像素时放大图像再判断不要凭印象标。YOLO格式的标注内容是归一化坐标一行一个目标格式是类别ID、中心点x坐标、中心点y坐标、宽度、高度。计算方式是像素坐标除以图像宽高全部归一化到0到1之间。我举个例子0 0.514 0.213 0.082 0.047这行的意思是类别ID为0的缺陷中心点在图像横向51.4%、纵向21.3%的位置框的宽度占图像宽度8.2%高度占4.7%。写脚本做数据检查时只要看这些值是否都在0到1之间就能快速判断格式有没有问题。2.3 数据增强别让模型死记硬背钢轨图像有一个特点大量背景区域长得非常像基本都是灰黑色的轨面纹理。如果模型只见过原始图像容易“背”下特定图像里的背景特征而不是真正学会缺陷长什么样这就是过拟合的典型表现。我用的数据增强以albumentations库为主主要做了以下几类亮度对比度随机调整模拟不同光照随机添加高斯噪声模拟传感器噪点轻微模糊模拟运动模糊场景随机翻转、随机裁剪增加目标位置多样性随机旋转小角度模拟拍摄角度抖动。这里要特别提醒水平翻转可以用但垂直翻转要慎重。钢轨图像里裂纹的方向是有物理意义的垂直翻转会让算法学到“倒着的裂纹”这个在实际情况里不存在的反而干扰学习。我做过对比实验加了垂直翻转后某些类别精度反而掉了1到2个点后来果断去掉。另外还有Mosaic增强YOLOv5默认在训练前几个epoch开这个。它把4张图拼成一张增强小目标检测能力对钢轨这种可能存在小缺陷的场景帮助很大。但这也意味着训练出的模型对输入尺寸更敏感推理时最好保持和训练时一致的输入尺寸。2.4 数据集划分与目录结构数据集不是简单堆在一起就行划分和目录结构都有讲究。我按8:1:1的比例划分训练集、验证集、测试集划分时注意同一条钢轨不同帧的图像不要同时进入训练集和测试集否则会造成数据泄漏测试指标虚高。目录结构按YOLO官方惯例组织datasets/ steel_defect/ images/ train/ val/ test/ labels/ train/ val/ test/images放图像labels放对应的标注txt文件名保持一致。这里有一个容易犯的错图像目录和标注目录的层级、文件名必须严格一一对应训练时报“label not found”基本都是这里出了问题。3. 训练环境与关键参数配置3.1 硬件与软件环境训练阶段我用的是一张消费级显卡显存11GB跑YOLOv5s模型完全够用YOLOv5m勉强能跑但batch size要降下来。实际上整个项目从头到尾单卡就完成了。软件环境方面PyTorch版本按官方要求来CUDA和cuDNN匹配好。我用的是官方仓库clone下来后按requirements.txt安装依赖。这里建议用虚拟环境装别直接怼到系统Python里否则后面项目多了依赖冲突会让人抓狂。git clone https://github.com/ultralytics/yolov5 cd yolov5 pip install -r requirements.txt3.2 data配置文件怎么写训练前需要定义一个数据集配置文件告诉YOLOv5三件事类别数量、类别名称、以及训练集/验证集图像路径。我这个项目的配置大概是这样的# data/steel_defect.yaml train: ../datasets/steel_defect/images/train val: ../datasets/steel_defect/images/val nc: 5 names: [scratch, spalling, crack, indentation, rust]nc是类别数量names是类别名称列表顺序必须和标注时的类别ID一一对应。如果标注时类别ID是0到4这里names的顺序不能随便换否则模型训练出来推断的类别全是乱的。3.3 训练命令与超参数选择官方YOLOv5训练命令不长关键在参数怎么调python train.py --img 640 --batch 16 --epochs 120 --data data/steel_defect.yaml --weights yolov5s.pt --cache几个参数逐个说。img是输入尺寸。我选640因为钢轨原图一般分辨率较高但缺陷尺寸不算大太大增加显存负担太小细节丢失。如果缺陷非常细小可以考虑把原图切成多块tile分别检测效果往往比直接拉大输入尺寸更好。batch是批量大小。显存16GB以下batch别硬怼32显存溢出就直接中断训练。我的显卡11GB跑640分辨率batch 16是稳定值。显存不够还想加大batch可以用--device指定多卡或者降低输入尺寸。weights是预训练权重。我从yolov5s.pt起步它是在大规模自然图像数据集上预训练过的对钢轨这种工业场景有很强的迁移学习价值训练收敛速度明显快于从零训练。epochs不是越多越好。我一开始设了200轮但到150轮左右验证集损失就已经不再下降后面的训练纯属浪费时间。后来改成120轮配合早停回调实际跑到一半就停了。这个值可以设大一点让早停机制来决定什么时候结束。训练过程要盯着几个指标训练损失、验证损失、mAP、Precision、Recall。验证集mAP达到90以上后再强行继续训练很容易出现过拟合——训练指标还在涨验证指标开始往下掉。3.4 训练过程与结果验证训练完成后weights目录下会生成best.pt和last.pt。best.pt是训练过程中验证集表现最好的权重last.pt是最后一轮的权重。部署时一律用best.pt别图省事用last.pt。验证阶段跑一下python val.py --weights run/train/exp/weights/best.pt --data data/steel_defect.yaml会输出各类别的精确率、召回率、mAP指标还会保存一批预测结果的可视化图。这些图一定要看不是只看数字。数字达标但预测框画偏了的情况很常见可视化是最直接的质检手段。4. 训练好的模型推理与部署4.1 直接用detect.py推理训练好的模型最快验证方式就是用官方detect脚本python detect.py --weights run/train/exp/weights/best.pt --img 640 --source ./test_images --conf-thres 0.25--source可以指向单张图、图片文件夹、视频文件甚至摄像头流。运行时会在runs/detect/exp下生成带框的标注图。我把一批真实巡检图扔进去跑观察哪些缺陷被检出、哪些漏掉、哪些误报这比看指标更直观。4.2 用Python脚本做自定义推理实际项目里不会只用命令行更多是要把模型嵌入到自己的处理脚本里。自定义推理的代码并不复杂import torch from pathlib import Path model torch.hub.load(ultralytics/yolov5, custom, pathrun/train/exp/weights/best.pt, force_reloadTrue) model.conf 0.25 model.iou 0.45 results model(./test_images/example.jpg) results.show() results.save(save_dir./output)torch.hub.load里path参数直接指向训练好的权重文件它会自动加载模型结构和参数。这里的conf是置信度阈值低于这个值的检测结果会被过滤掉。iou是NMS时的IoU阈值控制重叠框的去留。置信度阈值值得多讲一下。它就像一道筛子设得低检出的目标多但误报也多设得高误报少但可能漏掉真实缺陷。现场部署时一般先按0.25跑一轮看漏检率和误报率再微调。有些缺陷表面特征本来就不明显置信度很低阈值压太狠就全漏了。4.3 导出ONNX与边缘部署训练好的模型不可能永远活在训练电脑上。巡检场景经常需要把模型部署到边缘设备或者工控机上这时候ONNX是绕不开的格式。YOLOv5官方支持一键导出python export.py --weights run/train/exp/weights/best.pt --include onnx导出ONNX后可以脱离PyTorch环境运行方便用ONNX Runtime或者TensorRT加速。导出后要验证一下输出是否和PyTorch版本一致最简单的方法是用同一张图分别跑PyTorch和ONNX对比检测框差异。我试过几次导出后如果输入尺寸没对齐结果会差不少这个问题一定要查。5. 常见问题与排查经验5.1 过拟合训练集指标高真实图像表现差这是钢轨缺陷检测项目里最典型的问题之一。钢轨图像背景相对单一缺陷类别又比较固定模型太容易“背”下训练集的风格。表现就是验证集mAP很高换到另一条线路拍回来的图像漏检率明显上升。解决思路有三个方向。第一是扩充数据多样性多采集不同线路、不同光照、不同轨型的数据。第二是加强数据增强减少模型对背景纹理的依赖。第三是引入正规化手段比如调低训练轮数配合早停。我自己试下来最有效的还是第一条数据多样性是根本。5.2 细小缺陷目标检测困难钢轨上的细小裂纹在640分辨率下可能只有十几个像素宽。YOLOv5的输入图会做缩放小目标特征很容易在卷积和池化过程中丢失。实测下来小目标漏检比大目标严重得多。针对这个问题最直接的做法是把原图切成重叠的区块每个区块单独送进模型检测最后合并结果。比如一张2000像素宽的原图切成640x640的块重叠部分设128像素避免缺陷正好切在边界上。这样做的代价是推理时间变成原来的好几倍但检测效果提升非常明显。另外一个思路是调整anchor。YOLOv5有anchor自适应机制可以在训练命令里加--save-hybrid或者用automatic anchor optimization。但我实测下来对于钢轨这种细长形目标显式调整anchor的宽高比比自动优化更好使。把anchor的宽高比调向细长方向裂纹这类目标的召回率能有明显改善。5.3 光照反光和油污干扰导致误检钢轨表面有金属光泽在特定光照下会产生强反光反光区域的边缘特征和擦伤非常像算法很容易误判。油污区域颜色发暗边界也不规则和剥落、锈蚀容易混淆。这个问题靠模型本身很难彻底解决我的经验是从数据和预处理两头使劲。数据方面采集时专门收集一批反光、带油污的负样本标为“无缺陷”背景让模型见过足够多的干扰案例。预处理方面用直方图均衡化突出纹理细节或者用局部对比度增强再送进模型能够显著减少光照不均的影响。5.4 显存不足导致训练中断训练中途报CUDA out of memory几乎每个人都遇到过。排查思路不是盲目调小batch而是分步骤检查先用小batch跑通确认数据加载和模型前向没问题再逐步加大。如果某个batch大小下刚好溢出先降一半训练等模型收敛后再用更大batch微调。也可以开启--cache参数把数据预加载到内存里减少磁盘IO压力。对显存不够但内存充足的机器来说这个方法体感提升很明显。5.5 问题排查速查表现象常见原因解决建议训练损失不下降学习率过高或数据集噪声太大降低学习率检查标注是否有误验证损失反弹过拟合提前停止训练加强数据增强小缺陷漏检输入分辨率低或anchor不合理切图检测调整anchor宽高比反光区域误报训练集缺少干扰样本补充负样本预处理增强对比度推理速度慢模型规格太大或未导出换小模型导出ONNX或TensorRT标注文件报错格式或路径不匹配检查归一化坐标范围核对文件命名6. 后续还能怎么扩展这个方案当前已经能完成常规的钢轨表面缺陷检测但远不是终点。我个人的体会是这类工业视觉项目永远有迭代空间。一个方向是引入更丰富的缺陷分类体系。现在的类别划分还是粗粒度实际现场对缺陷严重程度更关心比如同是擦伤轻微擦伤和严重擦伤的处理方式完全不同。可以把检测任务从“有没有缺陷”延伸到“缺陷严不严重”这对模型的数据标注要求又高了一级。另一个方向是从检测走向分割。裂纹这类细长缺陷用边界框表示本身就有点别扭框里包含的背景区域很多。如果换成实例分割或者语义分割能拿到缺陷的精确轮廓对后续量化分析比如计算缺陷面积、长度非常有价值。现在YOLOv8系列已经有了分割模型把检测权重升级到分割模型整体改动不大收益却比较明显。部署层面如果现场有边缘计算盒子可以尝试把ONNX模型转成TensorRT引擎用FP16精度推理速度还能再快一截。我当时在测试机上对比过TensorRT比PyTorch原始推理能快2到3倍对于实时巡检场景意义很大。最后分享一个小经验模型做完不是结束得留一套完整的测试图像集定期跑一遍回归测试。什么时候数据扩了、标注改了、预训练权重换了都用同一套测试图对比指标保证模型越改越好而不是按下葫芦浮起瓢。工业项目里稳定比惊艳重要得多。