基于YOLOv5的铁路异物检测与边缘部署实践 简介面向铁路安全场景的Yolov5铁路异物检测项目源码适合目标检测学习、智慧交通研究及工业落地调试的开发者。源码以YOLOv5为框架覆盖6类轨道异物行人、动物、车辆、摩托车等提供训练、预测、数据处理等完整流程帮助读者快速复现检测模型。压缩包共12个文件以Python脚本和yaml配置为主含dataset.py、train_yolov5.py、predict_yolov5.py等核心代码以及data.yaml、requirements.txt等配置依赖说明整体仅16KB轻量易用。内容包含数据集结构划分、依赖安装、数据集类创建、模型评估与运行步骤可直接在本地环境按文档跑通实验。对于希望在铁路异物检测方向快速上手、理解YOLOv5训练全过程的读者这是份少走弯路的参考示例。目前已有106人学习使用。 去年下半年我接了一个铁路周界安防项目需求是在监控视频里实时检测轨道上出现的异物比如落石、翻越防护网的行人、闯入的工程车辆。一开始我们想走传统视觉方案毕竟铁路现场是固定机位背景相对稳定背景建模应该够用。但真正跑起来才发现轨道反光、隧道口光照突变、雨雾天气随便一个场景都能让误报率高到没法用。没办法只能转向深度学习目标检测模型选型绕了一圈最终锁定了YOLOv5。这篇文章就把整个项目从环境搭建、数据准备、训练调参到部署落地的过程完整过一遍包括那些折腾到凌晨的坑给准备做铁路异物检测或者其他固定场景目标检测项目的朋友一条可以直接下脚的路线。1. 铁路异物检测的需求拆解与YOLOv5选型理由1.1 先搞清楚“异物”到底指什么别一上来就训练做算法之前最忌讳的就是拿着模型先跑跑不动再回头想需求。铁路异物检测这个题目听起来很大但和客户对齐之后会发现真正要管的就那几类轨道和隧道里出现的落石或者大型掉落物、翻越防护网进入轨道的行人、闯入施工区域的工程车辆。目标类别定得越清晰后面的数据采集、标注、模型评估就越有抓手。这决定了整个技术方案的设计方向。这些摄像机装在接触网杆、桥隧口、防护网外机位固定、视角不变这是一个天然的优势。但同时铁路沿线的光照变化极端早中晚太阳角度不同隧道口一进一出明暗对比非常强轨道表面在雨后会形成大面积反光。这些因素叠加在一起传统背景建模和动态阈值类的方案很难稳定工作。我当时做了个简单测试用混合高斯背景建模去检测轨道上的行人晴天白天效果还可以一到阴天或者晚间直接乱套误报多到值班人员直接关掉系统。1.2 对比了Faster R-CNN、SSD和YOLOv8为什么最后还是选了YOLOv5选型阶段团队里有人倾向Faster R-CNN理由是精度高。Faster R-CNN在VOC和COCO上的mAP确实漂亮但两阶段模型在边缘设备上部署代价太高RPN加上ROI Pooling那一套即便用TensorRT优化帧率依然上不去。另一个候选是SSD实时性还行但小目标检测能力太弱铁路场景里远处侵入的行人可能只有二三十个像素SSD基本就直接丢了。YOLOv5当时并不是最新的模型YOLOv8已经出来了但它是工程成熟度最高的。社区资源不用多说训练教程、部署案例、踩坑帖子铺天盖地模型结构干净导出ONNX、量化INT8、在OpenVINO和TensorRT上都有现成方案YOLOv5s参数量只有7.2M左右INT8量化后可以跑进很低的延迟。放到今天如果业务方给的时间足够多选YOLOv8也完全没问题但如果是2023年前后做边缘端交付YOLOv5的生态成熟度是其他版本替代不了的。选成熟方案不是保守而是给项目周期留出容错空间。2. 环境搭建显卡驱动、CUDA版本和YOLOv5依赖2.1 显卡驱动匹配一个晚上才查明白的版本对应关系环境搭建这个环节网上教程一搜一大把但照着抄一定会出问题因为每个人的显卡驱动、CUDA版本、PyTorch版本组合都不一样。我当时的训练机是RTX 3090Ubuntu 20.04。第一次装的是NVIDIA驱动525配CUDA 12.0然后pip装了PyTorch 1.12。代码里跑torch.cuda.is_available()返回Falsenvidia-smi明明能看到显卡驱动也正常但PyTorch就是认不到卡。排查了一晚上才明白PyTorch安装包是预先编译好的它内部链接的CUDA runtime跟系统里装的CUDA工具包必须满足PyTorch官方支持的版本组合。PyTorch 1.12默认支持的最高CUDA版本是11.6你系统装12.0它确实不认。后来我把驱动换成了475对应CUDA 11.8系统装CUDA 11.8 toolkitPyTorch换成2.0.1cu118一次就跑通了。这个组合建议大家直接抄至少到2024年上半年它依然很稳# Ubuntu 20.04 RTX 3090 sudo apt install nvidia-driver-475 # 然后安装CUDA 11.8 toolkit并配置环境变量 export PATH/usr/local/cuda-11.8/bin:$PATH export LD_LIBRARY_PATH/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH # 创建独立conda环境 conda create -n yolov5 python3.9 conda activate yolov5 pip install torch2.0.1 torchvision0.15.2 --index-url https://download.pytorch.org/whl/cu118核心原则是先确定PyTorch版本再反推CUDA版本最后确定驱动版本。顺序反了就是给自己挖坑。2.2 YOLOv5源码和依赖安装的细节PyTorch就位之后YOLOv5环境本身就很省事了git clone https://github.com/ultralytics/yolov5.git cd yolov5 pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simplerequirements.txt里的包就那么十几个没有玄机。真正坑人的是预训练权重的下载。第一次执行detect.py时它会自动下载yolov5s.pt国内网络经常超时。我的办法是先用浏览器或者wget把权重文件下载到本地放到yolov5/weights目录然后命令改成python detect.py --weights weights/yolov5s.pt --source data/images/bus.jpg能输出一张带检测框的图片就说明环境链路通了。但这里我建议再做一步验证跑一下python train.py --data coco128.yaml --epochs 5用官方最小数据集训练几个epoch。这一步只要几分钟却能在正式训练前暴露很多隐藏问题比如显卡显存溢出、cudnn初始化失败、数据集路径配置错误。别嫌麻烦这一步能省你后面一整天的排查时间。3. 训练自己的数据集标注、划分和训练命令3.1 数据采集建议用视频抽帧并刻意维持类别均衡铁路现场的数据不可能一次性买到或者下载到现成图片最实际的办法是从监控视频里抽帧。我当时是让客户提供了最近三个月的历史监控视频每隔5到10秒抽一帧尽量覆盖白天、夜间、黎明、黄昏、晴雨天气重点保留目标出现的帧。抽完以后还要做一次数据清洗把画面几乎重复的、镜头有严重模糊的帧删掉否则模型会更容易记住重复背景而不是目标本身。抽了大概3000张有效图片之后我统计了一下类别分布行人占了60%以上落石只有15%工程车辆25%。这是典型的不均衡分布。注意这里的问题不只是数量多少而是占比最小的落石类恰恰是检测难度最高的——落石外形多变颜色与轨道碎石背景很接近光照不好时人眼都难分辨。如果你碰到类似情况不要硬着头皮训先想办法补充样本。3.2 标注规范和坐标可视化检查标注工具我用LabelImg它对YOLO格式支持很直接标注完每张图片对应生成一个txt文件每一行是类别人_center y_center width height坐标全部归一化到0到1。工具本身简单但标注规范直接决定模型上限被遮挡的目标只框可见部分不要把整个完整目标硬框进去否则模型会学到大量错误边界。铁轨上远处的小人只要人眼能辨认就尽量标出来小目标检测能力全靠这些标注学出来。类别边界要在标注前和团队对齐。比如检修工人算不算“行人”如果算是不是需要允许某个时间段出现在某个区域这些业务规则会影响数据标注的归类越早确定越好。每次新增标注之后我都会跑一个可视化脚本把labels画到原图上检查一遍。这一步对新手尤其重要因为坐标文件里只要有一行解析错误训练时就会报错或者出现不可思议的loss值。宁可多花半小时检查也不要盲目开训。3.3 目录结构和data.yaml的配置YOLOv5的数据集目录结构是约定好的建议直接按它的规矩来datasets/ rail_foreign/ images/ train/ val/ test/ labels/ train/ val/ test/划分比例我用的7:2:1脚本方式随机分配。这里有一个很容易被忽略的细节同一视频片段抽出来的帧划分的时候要保证只落在其中一个集合里不能既在train又在val。否则val里包含和train高度相似的帧mAP会虚高部署到现场立刻现原形。对应的data.yaml配置如下path: datasets/rail_foreign train: images/train val: images/val test: images/test nc: 3 names: 0: falling_rock 1: person 2: vehicle3.4 启动第一轮训练拿到baseline训练命令不长python train.py \ --img 640 \ --batch 16 \ --epochs 100 \ --data datasets/rail_foreign/data.yaml \ --weights yolov5s.pt \ --name rail_exp1这里解释一下参数选择的逻辑。--img 640是速度和精度的平衡点铁路异物目标大小不算极端640够用没必要上1280。--batch 16在RTX 3090 24G显存下比较舒适如果显卡小就调成8但不要低于4否则BN层的统计不稳定。--epochs 100是中小数据集的常见起点第一轮训练的目标是拿到一个可用的baseline不是最优结果。--weights yolov5s.pt是用官方预训练权重做迁移学习收敛速度远快于随机初始化。4. 超参数调优训练收敛速度与精度的平衡4.1 超参数文件里只有少数参数值得动手YOLOv5把训练超参数放在data目录下比如hyp.scratch-low.yaml训练时通过--hyp指定。很多新手拿到超参数文件就不知道该不该动我的建议是先跑默认再看训练曲线决定改哪里。第一轮训练我用默认的lr00.01发现loss曲线在刚开始几个epoch出现明显上升。这是学习率偏高的信号尤其是自己的数据集只有几千张数据规模远小于COCO模型很容易在初始阶段震荡。把lr0降到0.001之后loss曲线立刻平滑下来。这是我在多个项目里验证过的经验。数据增强方面的mosaic我也做了调整。默认1.0代表每张训练图都会做mosaic拼接这个操作对小目标检测非常有效但铁路场景背景本身固定拼接出来的假背景太多反而不利于模型学习真实场景特征。我把mosaic调到0.8最后10个epoch保留默认的自动关闭机制让模型在接近真实分布的数据上微调。4.2 观察哪些指标怎么判断模型在收敛还是过拟合训练过程中我主要盯results.png里的两对曲线train loss和val loss的差距以及每个类别的mAP走势。如果val loss从某个epoch开始反弹train loss还在下降基本可以确定是过拟合。这时候可以选择缩短epochs或者提前用早停。YOLOv5的patience参数默认是100训练一个小数据集时我习惯调成20让它在val不再改善时快速停下来。第二个关注点是分类别mAP。我用一个脚本单独统计每个类别的mAP发现落石类的mAP只有0.82而行人和车辆都到了0.93以上。这不是模型问题是样本量太少。处理方式是双管齐下一方面从客户的历史巡检视频里继续抽落石相关的帧补充数据另一方面对落石类做离线增强包括旋转、亮度扰动、随机裁剪这类别样本数提上来之后mAP明显往上走。4.3 精调后的部署模型选择整轮调完best.pt在测试集上的mAP0.5达到了0.89val loss没有反弹。此时我建议不要只保留一个训练好的YOLOv5s同时训练一个YOLOv5n版本作为备选。YOLOv5n的mAP大概会掉两三个点但推理速度能提升近一倍。实际部署时x86工控机用YOLOv5s边缘盒子用YOLOv5n一个模型满足高精度需求一个满足高帧率需求两个候选都提前准备好现场联调会从容很多。5. 模型导出与嵌入式部署后处理和推理优化5.1 导出ONNX时的参数选择训练完成后部署的第一步是把PyTorch权重导出为ONNXpython export.py --weights runs/train/rail_exp1/weights/best.pt --include onnx --opset 12--opset 12是我推荐的设置。opset版本太低某些新算子无法表示版本太高一些存量推理引擎兼容性又不好。opset 12在绝大多数环境下是兼容性最稳的。另外一个建议是导出时可以固定输入尺寸--imgsz 640 640如果你的部署平台对动态shape支持不友好固定尺寸能在转换成TensorRT或者NPU格式时省去大量麻烦。5.2 YOLOv5输出张量结构和后处理别直接套网上的85维代码这是部署时最容易翻车的地方。YOLOv5输入640x640的图像时输出shape是[1, 25200, 5nc]其中25200是3个尺度特征图的预测框总数5对应cx、cy、w、h和objectness置信度nc是类别数。官方COCO模型的nc80所以最后是85维。但你的自定义数据集nc3输出维度就是8如果直接套网上的通用后处理代码按85维解析类别分数会全部错位。正确后处理至少包含这几步把模型输出的中心点坐标和宽高转换成左上角、右下角坐标按objectness乘以类别分数的综合置信度过滤低分框我用0.35作为阈值对剩余候选框做NMS移除重叠度高的低分框把归一化坐标乘回原图尺寸得到像素坐标我实际写的后处理逻辑在推理类里面后面第6部分会给出简化示例。5.3 嵌入式平台部署i.MX8MP的量化与树莓派5的性能这个项目的最终端是两台设备一台x86工控机一台NXP i.MX8MP核心板。两边的部署套路完全不同。x86工控机用Intel的CPU导出ONNX后转成OpenVINO的IR格式直接调用OpenVINO Runtime推理整个过程最顺利。i.MX8MP则麻烦不少它自带NPU算力大概2.3 TOPS需要先把模型量化为INT8再通过NXP eIQ工具链转换。这里有一个典型的算子兼容性坑YOLOv5早期版本大量使用Focus层不少NPU工具链对Focus层支持很差转换时直接报错。解决方案是升级到较新版本YOLOv5代码重新导出或者手动将Focus层替换为等效的普通卷积层很多情况下问题就消失了。如果你在部署时遇到转换失败优先检查是不是Focus层的锅。树莓派5这边它没有NPU主要靠CPU跑。用ONNX Runtime CPU后端跑FP16的YOLOv5s实测10到15FPS做离线视频分析或者对实时性要求不高的周界巡检是够用的。如果追求更高帧率可以上YOLOv5n或者给树莓派加AI加速模块。INT8量化之后模型文件的mAP会有0.5到1个百分点的下降但推理速度提升非常明显。我在i.MX8MP上量化完YOLOv5s从纯CPU时的4FPS提升到NPU上的20FPS左右完全是质变。6. 项目源码结构与完整跑通流程6.1 源码目录设计为什么要对原版YOLOv5做一层封装项目源码如果只是把YOLOv5原版clone下来再放一个训练好的权重对客户来说几乎没有交付价值因为他们真正需要的是一个能持续对视频流做检测并返回告警的程序而不是一个命令行工具。我最终交付的源码目录大致是这样rail-foreign-detection/ weights/ best.pt data/ rail_foreign.yaml datasets/ rail_foreign/ scripts/ extract_frames.py split_dataset.py export_onnx.sh detector/ __init__.py inference.py tracker.py main.py requirements.txt README.md核心文件是detector/inference.py里面封装了一个RailForeignDetector类对外只暴露detect(frame)和detect_video(video_path)两个方法。上层无论是写Flask接口还是做本地视频分析都只需要调用这个类完全不用关心模型细节。6.2 推理类封装的核心代码逻辑简化版的推理类大概是这样的import cv2 import numpy as np import onnxruntime as ort class RailForeignDetector: def __init__(self, onnx_path, conf_thres0.35, iou_thres0.45): self.session ort.InferenceSession(onnx_path, providers[CPUExecutionProvider]) self.conf_thres conf_thres self.iou_thres iou_thres self.input_size (640, 640) self.class_names [falling_rock, person, vehicle] def preprocess(self, frame): img cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) img cv2.resize(img, self.input_size) img img.astype(np.float32) / 255.0 img np.transpose(img, (2, 0, 1))[None] return img def postprocess(self, outputs, orig_shape): # outputs形状为 [1, 25200, 8]8 5 3类 preds outputs[0][0] boxes [] for p in preds: cx, cy, bw, bh, obj_conf p[:5] class_scores p[5:] cls_id int(np.argmax(class_scores)) cls_conf float(class_scores[cls_id]) score obj_conf * cls_conf if score self.conf_thres: continue x1 (cx - bw / 2) * orig_shape[1] y1 (cy - bh / 2) * orig_shape[0] x2 (cx bw / 2) * orig_shape[1] y2 (cy bh / 2) * orig_shape[0] boxes.append([x1, y1, x2, y2, score, cls_id]) # 到这里之后继续按score降序、NMS去重 # 可以用cv2.dnn.NMSBoxes也可以自己写一版 return boxes def detect(self, frame): orig_h, orig_w frame.shape[:2] inp self.preprocess(frame) outputs self.session.run(None, {self.session.get_inputs()[0].name: inp}) return self.postprocess(outputs, (orig_h, orig_w))这比直接用原版detect.py稳定多了。原版detect.py的工作重点是“把检测结果画到图片上导出视频”而工程化需要的是一段能嵌入业务系统的推理模块两者目标完全不同。6.3 加一个轻量跟踪器告警才真正稳定只做单帧检测的话无论是x86还是嵌入式设备都会出现同一个目标这一帧被检测到、下一帧丢失的抖动。我加了一个简单的IOU匹配跟踪器把上一帧的目标框存下来当前帧检测框和上一帧框IOU超过阈值就认定为同一个目标同一目标连续出现5帧以上系统才触发告警并附上最近一帧的截图和坐标信息。这个逻辑看起来简单但效果出奇地好。夜间监控画面里经常出现时有时无的灯光反射、落叶飘过等干扰单帧检测很容易误报。加了连续帧确认机制后这些干扰大多被自动过滤掉了。如果是在原版detect.py上想快速达到同样的效果其实很难因为原版没有跟踪状态的概念。这也是我坚持做工程化封装的最核心原因。6.4 从零到上线的执行顺序和时间节点最后梳理一遍整个项目跑通的经验。为了保证交付顺利我建议严格按下面这个顺序执行按第2部分的配置搭建训练环境跑通官方demo再跑一个5轮的coco128训练验证流程。用视频抽帧脚本生成数据集和业务方确认类别定义用LabelImg标注按7:2:1划分注意同源帧不能跨集合。写好data.yaml启动第一轮训练拿到baseline记录loss和mAP指标。根据results.png调整超参数和类别样本均衡训练到val loss稳定上升为止。用export.py导出ONNX在目标平台上转成对应的推理引擎格式注意Focus层和opset兼容性。用自研推理类接入视频流加上跟踪和连续帧确认逻辑。到现场对着真实录像做联调重点看夜间和雨雾天的误报率必要时增加按时间段屏蔽检测区域的业务规则。我在这批项目里最终得到的收获是YOLOv5本身不是瓶颈原版模型加默认参数精度已经基本够用真正拉开体验差距的是数据质量、后处理和业务规则的配合。比如后来我们给系统加了一个“白天允许检修人员在特定区域作业、夜间禁止进入”的时间段规则告警准确率又上了一个台阶。把这些工程细节做扎实源码交付到客户手里之后维护成本才会真正低下来。本文还有配套的精品资源点击获取