ARTICLE DETAIL

资讯详情

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

YOLOv5钢轨缺陷检测实战:从数据集标注到部署全流程解析

YOLOv5钢轨缺陷检测实战:从数据集标注到部署全流程解析 简介基于YOLOv5的钢轨缺陷检测项目面向铁路维护、计算机视觉相关研发人员及深度学习入门者用于实现对钢轨表面裂纹、磨损、剥离等缺陷的自动识别与分类。项目以YOLOv5为核心同时涉及YOLOv7模型应用并包含数据预处理、标签生成、训练验证等关键环节可直接用于缺陷检测场景的算法验证与二次开发。资源包共2000个文件约28.77MB以txt标注与配置类文件为主辅以Python脚本、Markdown说明文档及yaml配置涵盖数据集标签、训练脚本和使用文档结构清晰便于按流程调用。已有167人学习下载。通过该包可获取完整的钢轨缺陷检测工程框架包括split.py等数据处理脚本、voc_labelhrsc.py等标注转换工具以及模型训练与验证的yaml配置有助于快速复现检测流程并理解YOLO系列模型在工业质检场景中的应用思路。1. 一张轨道照片为什么值得跑一个yolov5工程铁路工务段每天要做的事里有一件特别单调又特别考眼力沿钢轨查看轨面有没有裂纹、掉块、擦伤和波纹磨耗。过去靠人工巡道两个人一天走十几公里眼睛看花了也难免漏掉细小的横向裂纹。现在很多研究组和现场试点项目把相机装到巡检小车或轨旁设备上拍回来的图像交给目标检测模型去筛。这个标题里挂的 zip就是把“yolov5训练自己的数据集轨面伤损识别”打包成的一个典型工程训练脚本、标注数据、推理脚本放在一起解压后照理就能跑通一个钢轨缺陷检测的最小闭环。它对两类人最有用一类是刚接触yolov5、手里有一批轨面图片但不知道从哪下手的研究生另一类是工务系统的技术人员想评估视觉巡检能不能把漏检率压下来。先说一个反直觉的结论这类项目最终卡脖子的往往不是yolov5的网络结构而是你对轨面缺陷的认知和标注质量。2. 钢轨缺陷检测的数据基础伤损分类与标注体系2.1 轨面伤损到底有哪几类标注时怎么分视觉检测能看见的是钢轨表面的伤损内部伤损如核伤、螺孔裂纹二维图像基本无能为力那是超声探伤车的活儿。做yolov5钢轨缺陷检测先要把“缺陷”这个词落到一个可标注的类别清单上。我整理过一套比较稳妥的起步方案分为五到六类轨头踏面横向裂纹垂直于走行方向早期是一条细黑线放大后能看到从轨面延伸出的暗纹。轨面擦伤车轮打滑或制动造成表现为一块亮白的金属疤痕边界不规则常伴有一定深度的凹坑。剥离掉块轨面局部剥落形成坑洼比擦伤更深、轮廓更清晰。波纹磨耗周期性波浪状磨损在侧光下能看到明暗交替的横纹。轨侧磨损与肥边轨距角处金属被挤压翻边形成一道亮边。鱼鳞纹踏面浅层接触疲劳伤损密集的小裂纹群远看像一片鱼鳞。类别不是越多越好。初期建议控制在六类以内类别太多且样本不足时yolov5会出现“谁都不像”的误判。宁可把“剥离掉块”和“严重擦伤”合并成一个“表面剥落类”先把召回率做上去再细分。2.2 从 VOC 标注转成 YOLO 的 txt转换脚本与坐标坑大多数标注工具导出的是 Pascal VOC 格式的 XML而yolov5源码里默认读的是每个图像对应一个同名 txt每行内容为“类别ID x_center y_center width height”全部归一化到 0~1。转换这一步不算难但坐标计算和边界数值经常出问题。下面是一个我常用的转换脚本骨架import os import xml.etree.ElementTree as ET from pathlib import Path def voc_to_yolo(xml_path, out_dir, class_names): tree ET.parse(xml_path) root tree.getroot() img_w int(root.find(size/width).text) img_h int(root.find(size/height).text) out_lines [] for obj in root.findall(object): name obj.find(name).text if name not in class_names: continue class_id class_names.index(name) box obj.find(bndbox) xmin float(box.find(xmin).text) ymin float(box.find(ymin).text) xmax float(box.find(xmax).text) ymax float(box.find(ymax).text) x_center (xmin xmax) / 2.0 / img_w y_center (ymin ymax) / 2.0 / img_h w (xmax - xmin) / img_w h (ymax - ymin) / img_h # 边界裁剪某些标注框略微超出图像范围 x_center min(max(x_center, 0.0), 1.0) y_center min(max(y_center, 0.0), 1.0) w min(max(w, 0.0), 1.0) h min(max(h, 0.0), 1.0) if w 0 or h 0: continue out_lines.append(f{class_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}) if out_lines: out_file Path(out_dir) / (Path(xml_path).stem .txt) out_file.write_text(\n.join(out_lines), encodingutf-8) # 用法class_names 顺序必须与后续 data.yaml 中的类别顺序完全一致 class_names [transverse_crack, scratch, spalling, corrugation, side_wear] for xml_file in Path(path/to/xml).glob(*.xml): voc_to_yolo(str(xml_file), path/to/labels, class_names)这段脚本有几个容易被忽略的点。第一是 class_names 的顺序训练时 data.yaml 里的类别列表必须和这里一致否则类别 ID 全错位loss 会很低但预测结果完全不对。第二是宽高为 0 的框要跳过我见过标注软件导出的空框导致训练直接报错。第三是归一化后的坐标要裁到 0~1 区间因为有些标注框右边线正好画在图像边缘不裁剪会出现大于 1 的数值。转换完成后花十分钟做一次可视化检查用 OpenCV 把 txt 里的框画回原图跟 XML 原标注做个目测比对。这一步能拦下大多数坐标错位问题比训练到一半再回头查数据高效得多。2.3 类别不平衡钢轨缺陷数据里最现实的坑钢轨图片里“擦伤”和“波纹磨耗”往往占了大头“横向裂纹”可能只有几十个样本。yolov5默认的类别损失是等权重的样本少的类会被多数类“带偏”表现为训练集 mAP 尚可但横向裂纹这类关键伤损的召回率极低。针对这个我一般按顺序试三招第一招是过采样把少样本类的图片在训练集中多复制几份注意复制时要配合随机增强否则模型容易过拟合到固定的几张图上。第二招是给少样本类加标注权重在训练脚本里修改损失计算yolov5的损失函数在 loss.py 中可以把每个类别的 loss 乘上一个权重系数这个做法比简单过采样更可控。第三招是 mosaic 增强里多掺少量样本图yolov5的 mosaic 会把四张图拼成一张少样本图参与拼接的次数多了模型看到它的机会自然上升。对于“横向裂纹”这种细长小目标还有一个更朴素但有效的办法把原图按 2~4 倍放大切成小块来训练。钢轨原图可能是 5184×3456 的yolov5默认输入是 640×640直接缩下去裂纹只剩几个像素什么网络都学不到。先把轨面裁成若干 640×640 的子图相当于让模型在原始分辨率下看缺陷比调任何超参数都管用。3. 把 yolov5 工程跑起来环境配置与最小训练流程3.1 conda 环境与依赖版本怎么配不翻车yolov5项目对 Python 和 PyTorch 版本有一定要求但也没那么玄学。我习惯用 conda 单独建一个环境避免把系统 Python 环境搞乱。一个典型的创建命令是conda create -n rail_yolov5 python3.9 -y conda activate rail_yolov5 pip install torch1.13.1 torchvision0.14.1 --index-url https://download.pytorch.org/whl/cu117 pip install -r requirements.txt这里把 torch 和 torchvision 单独装是因为 requirements.txt 里的 torch 默认版本可能与你机器的 CUDA 不匹配。yolov5环境配置里最常见的翻车点有两处一是 CUDA 版本和 torch 不匹配装完跑起来提示“CUDA not available”二是 OpenCV 版本冲突报错信息往往是导入 cv2 时出现 GLIBC 错误。遇到第一处用python -c import torch; print(torch.cuda.is_available())先验证一下再往下走。遇到第二处可以考虑pip install opencv-python-headless替代带 GUI 的 opencv-python在服务器上跑推理完全够用。如果手里只有 CPU 机器也可以训练但要把 batch size 调小、输入分辨率调小训练时间会慢一个量级。对于钢轨缺陷检测这种现场数据量通常在几千张的项目CPU 训练勉强能接受但迭代周期长建议还是找一张显存 8G 以上的显卡。3.2 数据目录结构与 data.yaml 的写法yolov5源码的资料里数据目录结构有约定训练时通过 data.yaml 指定路径。我常用的做法是建一个 datasets 目录里面分 images 和 labels再按 train/val 划分。具体布局如下datasets/rail_defect/ images/ train/ val/ labels/ train/ val/ rail_defect.yaml对应的 rail_defect.yaml 内容train: datasets/rail_defect/images/train val: datasets/rail_defect/images/val nc: 5 names: [transverse_crack, scratch, spalling, corrugation, side_wear]这里有个容易踩的细节train 和 val 路径最好写相对路径或者写绝对路径两者都不要写错。yolov5有时候对路径的解析比较敏感如果报AssertionError: train: ... does not exist先把 yaml 里的路径打印出来检查一遍。另外nc必须和names列表长度一致不一致时训练会在加载数据阶段直接报错。3.3 第一次训练命令、日志怎么看环境配好、数据就位后第一次训练我建议先别急着调一堆参数跑一个最朴素的训练流程验证全链路是通的。命令如下python train.py \ --data datasets/rail_defect/rail_defect.yaml \ --weights yolov5s.pt \ --img 640 \ --batch 16 \ --epochs 100 \ --device 0这里用 yolov5s.pt 作为预训练权重适合第一次验证如果钢轨缺陷和 COCO 的类别差异大后面可以考虑从头训练或换用钢材表面缺陷数据集预训练。--img 640是训练输入分辨率钢轨裂纹属于细长小目标后面要试 960 或 1280但一开始先用 640 跑通流程否则显存不够会一直在报错调参。训练开始后日志里每两行就有一个 metrics 表格重点看三列box_loss、cls_loss、mAP0.5。前几个 epoch mAP 可能是 0不用慌这是正常现象。正常情况下 box_loss 和 cls_loss 应该缓慢下降mAP 从第 20 个 epoch 左右开始抬升。如果 loss 完全不降或者前 20 个 epoch 的 mAP 一直是 0大概率是 label 和 data.yaml 的类别对齐出了问题回头检查标注。第一次训练跑完去看runs/train/exp/weights/best.pt这个文件是否存在。yolov5会同时存last.pt和best.ptbest 是验证集 mAP 最高的那个实际部署时优先用best.pt。4. yolov5超参数与训练调优让模型记住钢轨而不是背景4.1 超参数文件里哪些值得动yolov5源码里有一个data/hyps/hyp.scratch-low.yaml和hyp.scratch-high.yaml定义了所有训练超参数。低增强版本适合数据量不大、缺陷形态相对固定的场景高增强版本适合数据量充足、想增强泛化能力的场景。钢轨缺陷检测的数据量一般不算大我建议从hyp.scratch-low.yaml起步。几个值得关注的超参数lr0初始学习率默认 0.01对钢轨这种类别数少的任务我经常降到 0.005。mosaic默认 1.0开启 mosaic 增强。钢轨缺陷如果尺寸小mosaic 拼接后缺陷可能被缩得更小可以降到 0.5 试试。hsv_h/hsv_s/hsv_v色彩增强幅度。钢轨表面是金属色光照变化本来就大这三个值建议保守一点比如hsv_s: 0.2否则增强出来的图像颜色失真模型学到的是“颜色异常”而不是“结构异常”。fliplr水平翻转默认 0.5轨面左右对称可以保持默认。修改超参数文件后训练时通过--hyp指定新文件不必改 train.py 源码。我自己习惯对每组超参数实验单独复印一份 hyp 文件方便回头对比避免改来改去忘了哪份是基准。4.2 预训练权重怎么选COCO 权重不一定最优yolov5官方提供的预训练权重是在 COCO 上训练的COCO 有 80 类包含人、车、猫狗等跟钢轨表面缺陷的视觉特征差别极大。用 COCO 权重做迁移学习主要迁移的是浅层特征——边缘、纹理、颜色梯度这些对钢轨缺陷依然有效。但如果你的标注数据量足够大比如超过 2000 张缺陷图可以考虑另外一个路径先用公开的钢材表面缺陷数据集做预训练再在你的轨面数据上微调。钢材表面缺陷数据集里的“划痕”“夹杂”“氧化皮”与轨面“擦伤”“剥离”在视觉形态上更接近迁移效果通常比 COCO 好。具体做法是用钢材表面的数据跑一遍 yolov5训练得到best.pt然后把这个权重作为你钢轨训练的--weights传入。没有公开数据集的情况下也可以用 COCO 权重先训练 50 个 epoch再在全部数据上从头微调相当于自己制造一个中间权重。需要提醒的是yolov5源码里加载预训练权重时会要求类别数一致或做部分匹配。如果预训练模型 nc5你的数据集 nc6torch.load之后就要手动删掉最后一层检测头的权重否则会报形状不匹配的错误。yolov5在这里有个表现就是载入权重时只加载相同形状的参数其他层随机初始化日志里会提示Transferred xxx/yyy layers看到这个不用紧张属于正常机制。4.3 后处理参数与置信度阈值先看 mAP 再谈阈值很多人训练完直接拿detect.py跑测试图发现漏检一堆于是回去改模型。其实在改网络之前先调整后处理参数往往就能找回几个百分点的召回率。yolov5后处理时有两个关键参数--conf-thres和--iou-thres。钢轨缺陷检测场景里横向裂纹的置信度天然偏低因为它细、短、对比度弱。我在验证时会把conf-thres从默认 0.25 降到 0.1看看模型是不是“其实检测到了但没显示出来”。如果 0.1 时框出来了说明模型能力没问题只是置信度阈值卡得高部署时可以用 0.15~0.2。如果 0.1 时依然没有框那才是真正的漏检。iou-thres控制 NMS 的区间去重力度默认 0.45。轨面擦伤和剥离这类相邻缺陷容易交叠如果同一个目标上叠了多个框将 iou-thres 调到 0.5 或 0.55 通常能去得更干净。4.4 从热力图和误报图找模型的“视野盲区”调参到一定程度后mAP 不再涨就要看具体失败样本。yolov5的val.py跑完会保存预测结果图我一般把验证集里漏检的图单独挑出来按类别统计。钢轨缺陷检测常见的失败模式有三种一是细长裂纹被背景纹理淹没二是擦伤与油污、水渍混淆因为金属反射光照下油污的亮白色块和擦伤很像三是暗光下缺陷与道砟阴影叠在一起。针对第一种失败解决方案是在训练时加入“侧光增强”——模拟低角度光照下的阴影变化让模型学到裂纹的线状结构而不是依赖固定的灰度值。针对第二种失败需要加入大量负样本也就是没有缺陷但有油污、水渍、道砟、扣件阴影的背景图类别标注为空即可。负样本的作用是让模型学会“不像缺陷”的特征这个效果往往比增加几千张正样本更明显。5. 钢轨缺陷检测常见的坑与排查从数据到训练再到推理5.1 现象训练 loss 下降但验证集 mAP 始终为 0原因标注类别 ID 和 data.yaml 中的 names 对不上或者 val 集中的标签文件缺失。我见过一种情况images/val 里有图labels/val 里却是空的yolov5不会报错只会把所有验证图片当成背景预测mAP 自然为 0。解决先用python val.py --data rail_defect.yaml --weights yolov5s.pt跑一小段看输出里是否提示WARNING: 没有找到 label之类的信息再写一个脚本统计 labels/train 和 labels/val 里的 txt 文件数量必须与图像数量一一对应。5.2 现象钢轨裂纹检测出来了但框的位置偏到旁边原因标注框未严格贴合缺陷边缘尤其是细长目标。标注软件里画框习惯性留了边距对普通目标无所谓对横向裂纹这种长宽比达到 10:1 的目标框的偏移直接导致 IoU 计算偏低。解决重新审视这类缺陷的标注规范严格沿裂纹两端和上下边缘画框宁可小一点也不能大出太多。另外可以针对细长目标开启 yolov5 的多种尺度训练也就是在--img参数之外把数据的保持宽高比缩放到多种尺寸让模型适应不同长宽比的目标。5.3 现象训练前中期 loss 突然反弹甚至变为 NaN原因学习率过大或 batch size 太小导致梯度不稳定。钢轨数据集中如果存在异常亮斑图片损坏的像素也可能给 loss 注入异常大的值。解决在超参文件里把warmup_epochs从默认值略微调大让学习率缓慢爬升同时检查数据集中是否有全黑、全白或者文件损坏的图片并剔除。遇到 NaN直接降低lr0到 0.003 再训练一次比调其他参数更有效。5.4 现象验证集 mAP 高到 0.9现场跑却频繁误报原因数据分布偏移。训练图像可能是白天顺光拍的现场是阴天侧光拍的金属反光的视觉特征完全不同。yolov5对光照的泛化能力有限训练集里没见过的光照模式模型只能靠猜。解决训练集中刻意混入不同季节、不同时段、不同角度光照下的轨面图像。其中一个有效手段是提升训练图像的灰度多样性把图像随机转成灰度图再参与训练让模型不把“颜色”当作核心特征。钢轨本身是灰暗金属色缺陷也以结构特征为主灰度增强对轨面检测很安全。5.5 现象小目标检测结果忽好忽坏同一段轨面不同帧一会检出一会丢失原因视频流中钢轨表面反光会随角度变化yolov5的预测对输入噪声敏感。另外如果部署时用视频流逐帧检测没有做帧间平滑单帧误报和漏检会被直接暴露。解决推理时对连续帧的检测框做时间维度的滤波——连续三帧中至少两帧都检测到同一个位置的同一个类别才输出这个框。这个策略能显著减少单帧闪烁检测问题实现起来也简单用卡尔曼跟踪或简单的“最近距离关联”都可以。yolov5基础笔记里常被忽略的一点就是检测模型输出的原始结果和业务可用的结果之间还隔着一层“时间稳定性处理”。6. 从 best.pt 到现场部署ONNX 导出、滑窗去重与低成本设备推理6.1 把权重导出成 ONNX摆脱 PyTorch 环境现场部署不可能每台机器都装一套 conda yolov5环境。常见做法是把训练好的best.pt导出成 ONNX 格式用 OpenCV DNN 模块或 ONNX Runtime 做推理。导出命令为python export.py --weights runs/train/exp/weights/best.pt --include onnx --opset 12导出后检查两个东西一是输出的 ONNX 文件尺寸是否正常二是用 ONNX Runtime 跑一次验证对比与 PyTorch 推理的 mAP 差异。yolov5部署时最容易出问题的点在于 NMS 要不要一起导出。export.py默认导出的 ONNX 只包含网络前向计算NMS 需要自己在部署代码里实现。yolov5python 里有一个models/experimental.py里的attempt_load可以加载 ONNX但更常见的做法是在部署端用 ONNX Runtime 跑出原始预测向量再自己写一份 NMS。这样最灵活现场调置信度阈值不用重新导出。ONNX 动态输入分辨率需要设置--dynamic参数如果现场输入尺寸固定建议别开动态性能和稳定性都更好。6.2 大图滑窗检测与去重从 640 到现场原图现场采集到的轨面图通常是几千万像素不能直接缩到 640 输入。我的做法是滑窗裁图每个窗口 640×640步长取 320 或 480相邻窗口有重叠。滑窗检测带来的一个副作用是同一个缺陷可能在多个窗口内被重复检出需要做跨窗口去重。去重的思路是把滑窗产生的检测框坐标映射回原图坐标系然后对这些框再跑一次全局 NMS。这里注意一个坑滑窗在边缘时缺陷可能被窗口边界切了一半导致召回率下降。解决方法是窗口重叠率不要低于 0.3并且允许窗口在原图边缘向外 paddingpadding 区域以灰度 0 填充检测后再把框坐标裁回原图范围。import numpy as np import cv2 import onnxruntime as ort session ort.InferenceSession(best.onnx) image cv2.imread(rail_scan.jpg) H, W image.shape[:2] window_size 640 stride 320 boxes_global [] for y in range(0, H, stride): for x in range(0, W, stride): x1, y1 x, y x2, y2 min(x window_size, W), min(y window_size, H) patch image[y1:y2, x1:x2] input_blob cv2.dnn.blobFromImage(patch, 1/255.0, (640, 640), (0, 0, 0), swapRBTrue) outputs session.run(None, {session.get_inputs()[0].name: input_blob}) # outputs 经过解码后得到 boxes坐标相对滑窗需加上偏移量 (x1, y1) boxes_global.append(decode_boxes(outputs, offset(x1, y1)))这段代码的逻辑是遍历整张大图每个窗口独立推理再把窗口内的检测框坐标叠加窗口偏移量得到全图坐标。参数说明stride 越小重复检测越多检测精度越高但耗时成倍增加stride 640 时没有重叠速度最快但容易切断目标。三个值之间选一个平衡点我用 320 比较多。6.3 树莓派 5 这类低算力设备的部署经验有人想在树莓派 5 上跑自己训练的yolov5模型可行但要注意模型规模和量化方式。yolov5n 或 yolov5s 是唯一选择yolov5m 以上在树莓派 CPU 上单帧推理时间会到 5 秒以上没有实用价值。把 ONNX 进一步转成低成本设备上的格式时我会用 INT8 量化量化后 mAP 可能会下降 1~3 个百分点用验证集评估后决定是否可接受。树莓派 5 上跑推理时CPU 频率要锁定到最高档散热器必须装好否则长时间推理会触发温降频速度反而更慢。关于线程数ONNX Runtime 的 intra_op_num_threads 设置为 4 时速度最快设为 8 反而会因为争抢资源而下降。还有一个经验树莓派做滑窗检测时缓存策略很关键。把轨道图片切块后不用的窗口及时释放内存不然十几张图跑下来内存就爆了。多级滑窗后处理加上帧间滤波能明显提升现场可用性。我在现场验证项目时通常会用一段 30 秒的实拍视频做验收计算“每 100 张图误报几个”和“同一个缺陷连续被检出多少帧”这两个指标。若多帧命中率达到 60% 以上就认为这个模型具备走进生产环境的潜力。部署前的最后一道工序是把检测结果与行位置绑定也就是给每个检测框打上“里程桩号”或“图像帧号”。这一步不做现场人员拿到一堆坐标也定位不到钢轨缺陷的实际位置项目价值会打折扣。我自己在部署这类工程时习惯把检测输出写成一个 CSV字段包括帧号、类别、置信度、x、y、w、h后续无论是人工复核还是接入工务系统都用得上。钢轨缺陷检测这个方向真正落地时拼的不是网络结构的炫技而是把数据分布、阈值策略、部署约束一层层磨合适配的过程。希望这份基于yolov5的落地笔记能帮你减少一些重复试错的成本。本文还有配套的精品资源点击获取
返回列表