ARTICLE DETAIL

资讯详情

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

基于YOLOv5s改进的铁路信号灯小目标检测与部署实践

基于YOLOv5s改进的铁路信号灯小目标检测与部署实践 简介面向铁路安全运输场景这套深度学习实践资料围绕卷积神经网络CNN的铁路信号灯识别方法展开适合图像识别入门者、计算机视觉方向学生及铁路智能监测相关研究人员。资源以普通铁路信号灯为研究对象从CNN基础模型出发逐步改进结构并优化学习率解决信号灯位于地面与高处、多灯光混杂等条件下的颜色识别难题为列车司机预判提供辅助思路。包内共141个文件其中127张jpg图片构成红、黄灯等典型样本集7个py脚本覆盖手写数字识别、颜色识别与信号灯实验代码还附有tfrecords数据文件、png示意图、xls结果表及说明文档整体压缩包约21.55MB结构便于对照学习。目前已有289人浏览学习适合需要从零搭建CNN颜色识别模型、开展实验对比与结果分析的读者可据此理解图像特征提取与分类辨识的完整流程。1. 项目概述与需求拆解1.1 为什么铁路信号灯识别值得单独做一个项目铁路信号灯这个东西乍一看不就是红绿黄三个颜色嘛能有多难真正上手做过才知道它跟马路上那种交通信号灯压根不是一个难度量级。列车运行速度高、制动距离长信号灯没看清或者看晚了后果非常严重。人眼在极端天气、强逆光、长时间驾驶疲劳的情况下误判漏判的概率并不低所以用视觉算法做辅助识别在这个场景里有实实在在的需求。更麻烦的是铁路信号灯在画面里的占比通常非常小。举个例子我用1080p的相机去拍几百米外的信号灯灯体在图像里可能只有十几个像素加上镜头畸变、振动模糊、晨雾夕阳这些干扰传统图像处理那一套——颜色阈值分割、形状匹配——很容易就崩了。这正是我决定用深度学习来做识别而不是继续抠传统算法的主要原因。除了技术层面这个项目的工程价值也很明显。铁路巡检、车载辅助驾驶、道口监控这些都是信号灯识别的落地场景。它能做到实时检测、持续运行不依赖人工盯屏一旦发现异常可以联动告警。对于做工业视觉、轨道交通相关的从业者来说这类项目既是CV入门的一套经典练手内容也是一条能写进简历的完整工程链路。1.2 项目技术路线与选型逻辑这个项目涉及的核心技术点集中在图像分类和目标检测两个方向。早先有一些方案把信号灯识别做成纯分类任务——先通过传统方法把灯的位置裁剪出来再用CNN判断颜色和状态。这种方式的问题在于定位一旦不准确分类再准也没用。所以我这版方案采用的是端到端的检测思路把“灯在哪里”和“灯是什么颜色”这两个问题一起解决掉。模型选型上我对比过两个方向。两阶段检测器比如Faster R-CNN精度确实高但推理速度在嵌入式设备上吃紧。单阶段检测器比如YOLO系列速度快小目标召回率也能通过技巧补上来。综合权衡检测实时性和部署成本之后我选了YOLOv5s作为基线模型然后再针对信号灯这个小目标场景做专门的优化。具体到网络层面我会在后面的章节把改动细节拆开讲。数据层面公共数据集里几乎没有专门针对铁路信号灯的高质量标注集所以我一开始就把自建数据集当成核心工作来对待。这其实也是实际工程和学校作业最大的区别——模型结构可以复现但数据才是决定项目上限的地方。2. 数据集构建与预处理2.1 数据采集的坑与经验铁路信号灯的公开数据少这是做这个项目首先要面对的现实。能拿到的原始数据来源大概有三类一是铁路沿线的固定监控摄像头二是车载前视相机录制三是从一些公开的视频片段里抽帧。我自己主要用的是前两种加起来采集了大约12000张图像。采集过程有几个特别容易踩的坑。第一个坑是画面比例问题——监控摄像头往往是俯视角度车载相机是平视角度二者信号灯的形状特征差异很大。如果混着训练不处理模型容易学到“角度特征”而不是“灯的特征”。我的做法是尽量让数据集里某一角度的图片占主导然后用少量另一角度的图片做微调验证。第二个坑是样本不平衡。绿灯和黄灯的样本数量远少于红灯——毕竟铁路信号里绿灯常见、黄灯相对少。如果不做处理模型对少样本类别的召回率会非常难看。我的解决思路是尽量特意去多录一些黄灯出现的时间段同时在后续做数据增强时对黄灯样本采用更高的增强倍率。第三个坑涉及到标注一致性。信号灯目标小标注框稍微画大一点画小一点对训练结果影响都不小。我检查标完的框时就发现不少人习惯把灯外面的黑色背板也框进去导致模型学了一堆背景信息。后来我重新统一了标注规范框必须紧贴灯体发光区域不包含背板和灯柱。2.2 数据标注流程与格式转换标注工具我用的是LabelImg虽然老一些但胜在稳定支持Pascal VOC格式的XML输出。我的流程是先把所有抽帧图像按7:2:1划分成训练集、验证集、测试集划分完之后再做标注——先划分再标注的好处是能避免同一条视频里前后帧的相似图像同时出现在训练集和测试集里否则验证结果会虚高。标注完成之后要把VOC格式转换成YOLO训练需要的txt格式。每一行对应一个目标格式是“类别id 中心点x 中心点y 框宽 框高”所有坐标都归一化到0到1之间。这个转换本身不复杂但写脚本时要注意坐标换算别出bug尤其是从XML的左上右下坐标转成YOLO的中心宽高格式容易搞混。import os import xml.etree.ElementTree as ET def voc_to_yolo(xml_file, target_dir, class_names): tree ET.parse(xml_file) root tree.getroot() img_w int(root.find(size/width).text) img_h int(root.find(size/height).text) lines [] for obj in root.iter(object): cls obj.find(name).text if cls not in class_names: continue cls_id class_names.index(cls) 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 lines.append(f{cls_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}) txt_name os.path.splitext(os.path.basename(xml_file))[0] .txt with open(os.path.join(target_dir, txt_name), w) as f: f.write(\n.join(lines))类别这块我用的是red、green、yellow三个类别先不做红色双灯、绿色双灯这种细分。灯芯数量、排列方式的分辨属于更细粒度的任务放到后续迭代再处理。首批先保证三个基础颜色能稳下来项目跑通了再叠加细节分类。2.3 数据增强策略的针对性设计信号灯目标小的特点决定了通用数据增强策略不能直接套用得做针对性的调整。Mosaic增强在YOLOv5里是默认开的它把四张图拼在一起训练对小目标检测特别有效因为拼图相当于把整体分辨率摊薄了目标在画面中的相对尺寸变得更小模型被迫去适应更严苛的检测条件。这个增强我保留了。颜色抖动这块我做了专门强化。铁路场景的光照变化非常剧烈大雾、雨天、逆光、夜晚都会影响灯体颜色表现。我在HSV空间做了较大范围的随机扰动H通道扰动幅度控制在±0.02以内——因为色相一旦扰动过头红灯的色相漂到橙色甚至黄色就会给模型注入错误标注。S和V通道可以放宽到±0.3。有一个增强操作我是刻意不用或者少用的水平翻转。这个要结合信号灯的物理位置逻辑来考虑——铁路信号机是有固定排列规则的比如有些地方红灯固定在左侧或上方。如果做水平翻转等于强行制造了一批“红灯在右侧”的样本跟真实物理场景分布不一致反而会干扰模型学习位置先验。这种思路上我当时也查了不少资料整体上倾向于保持原始空间关系。3. CNN模型结构设计与改进3.1 基线模型YOLOv5s的结构拆解YOLOv5s的结构分三块Backbone骨干网络、Neck特征融合层、Head检测头。Backbone用的是CSPDarknet核心模块是C3结构——它把输入特征分成两路一路经过若干Bottleneck模块提取深层特征另一路直接跳过最后在通道维度上拼接这样在控制计算量的同时增强了梯度流动。Neck部分采用FPNPAN结构FPN自顶向下传语义信息PAN自底向上传空间信息两者结合确保不同尺度的特征图都能拿到足够信息。Head负责在三个不同尺寸的特征图上输出预测框和类别概率。这个结构用在通用目标检测上很成熟但直接拿来检测铁路信号灯有两个问题。第一特征图下采样倍数太大。YOLOv5s的Backbone对输入图像做了32倍下采样最后一层特征图上一个小目标的特征几乎被压没了。信号灯目标小这一个问题直接决定了它容易漏检。第二三个尺寸的检测头中最大的特征图是80x80在640输入下本来应该负责检测小目标但它的感受野和语义信息都是最浅层的对小目标的表征能力有限——这是需要改进的地方。3.2 针对信号灯小目标的三处修改第一处修改是增加一个更高分辨率的检测头。原版YOLOv5s有三个检测头对应8倍、16倍、32倍下采样的特征图。我额外加了一个4倍下采样的检测头也就是输入640时输出160x160的特征图。这层特征图保留的空间信息最多适合捕捉小尺寸信号灯。对应的Neck部分也做了调整让这条新检测头所在的支路能同时获得深层的语义信息避免只靠浅层特征导致误检率升高。第二处修改是在Backbone里引入注意力模块。我在C3模块后面接了一个轻量级的SE注意力模块。SE模块的思路是先对特征图做全局平均池化得到每个通道的全局统计信息再通过两个全连接层先压缩再扩展学习每个通道的重要性权重最后把权重乘回原特征图。相当于告诉模型“什么特征重要就强调什么”。对信号灯这种目标小、但颜色和形状特征明显的场景通道注意力能有效增强语义响应。import torch.nn as nn class SEBlock(nn.Module): def __init__(self, channels, reduction16): super().__init__() self.squeeze nn.AdaptiveAvgPool2d(1) self.excitation nn.Sequential( nn.Linear(channels, channels // reduction), nn.ReLU(inplaceTrue), nn.Linear(channels // reduction, channels), nn.Sigmoid() ) def forward(self, x): b, c, _, _ x.size() y self.squeeze(x).view(b, c) y self.excitation(y).view(b, c, 1, 1) return x * y.expand_as(x)第三处修改是替换激活函数。原版YOLOv5s在C3模块里用的是SiLU激活函数我在深层网络部分尝试替换成了ReLU并结合权重初始化调整做了对比实验。这里说下结论ReLU和SiLU在信号灯检测任务上精度差异不大约0.3%mAP但ReLU在推理时对硬件更友好某些嵌入式设备的加速库对ReLU的支持更完善。如果只做实验不搞部署保持原版SiLU完全没问题。3.3 模型轻量化探索记录我还尝试过用MobileNetV3替换Backbone做轻量化。MobileNetV3用深度可分离卷积替代普通卷积参数量和计算量大幅下降但实验效果显示在信号灯数据集上替换后的mAP下降了约4.5个点主要是小目标召回率掉了。后续我有做过分析小目标检测对浅层空间分辨率敏感MobileNetV3为了轻量把浅层通道数压得比较低不利于保留小目标的细节特征。有一个方向值得继续试FPN结构加双向融合BiFPN它在相邻层之间增加了跨尺度连接。原版YOLOv5s的PAN虽然也是双向但BiFPN会给不同输入特征学习权重理论上对小目标的特征融合更友好。不过BiFPN带来的额外计算量跟精度收益之间的平衡在当前数据集上还需要更多实验验证。这个我在文末会再提一句。4. 训练过程与关键参数4.1 训练环境与超参数配置训练环境这块我用的是一块RTX 3090显卡显存24GB。软件栈是Python 3.9、PyTorch 1.12、CUDA 11.6。这里有个经验可以分享PyTorch版本和CUDA版本一定要匹配否则编译某些自定义算子的时候会报莫名其妙的错排查起来浪费时间。核心超参配置如下参数数值说明输入尺寸640x640兼顾精度与速度更大尺寸对小目标更友好但显存压力大Batch Size323090可承受再大收益不明显Epochs120前60轮冻结Backbone训练后60轮解冻全部训练初始学习率0.01采用余弦退火策略动态调整权重衰减5e-4防止过拟合优化器SGDMomentum设为0.937比Adam更适合检测任务类别权重red:green:yellow1:1:1.5给数量较少的黄灯更高的损失权重这里特别注意SGD和Adam的选择。很多人习惯直接用Adam但YOLOv5的默认配置SGD在检测任务上往往更稳训练到后期不容易出现loss震荡的问题。Adam收敛快但最终精度通常不如调好的SGD。这个差异在信号灯数据集上实测差了接近1个mAP点。4.2 训练曲线怎么判断是否正常训练要盯的不只是loss降没降更关键的是验证集mAP有没有跟着涨。我之前有过一次训练loss一路降到很低但mAP始终上不去——后来发现是训练集和验证集之间数据分布出现了明显差异训练集里大量是同源视频抽帧特征太相似了。正常训练的loss曲线特征是前10轮下降非常快从初始的0.1以上迅速降到0.05以下之后下降变缓在60轮左右进入平台期。如果loss降到平台期之后还在缓慢下降说明模型还在学有效特征不必着急提前停止。mAP曲线和loss曲线大致呈镜像关系但会有小幅滞后这是正常的。学习率的调整策略我直接用了YOLOv5自带的余弦退火它在每个epoch都会调整学习率让学习率从初始值逐渐降到接近0。相比固定学习率这种方式在训练后期能有效避免参数在最优解附近震荡稳定性和最终精度都有提升。4.3 训练后模型表现数据120轮训练完成后我在测试集上的模型表现如下类别精确率Precision召回率RecallmAP0.5red96.2%93.8%95.1%green94.7%91.2%92.8%yellow90.6%84.5%88.2%全部--92.1%可以看到yellow的召回率显著低于red和green这跟样本量少直接相关。虽然做了类别加权和增强策略但数据量本身不够这个先天因素不是完全能靠技术弥补的。后续要提高黄灯召回率最有效的还是继续采集黄灯样本。5. 部署实践与实时推理优化5.1 从PyTorch模型到TensorRT部署训练好的模型不能直接拿到现场跑PyTorch模型在GPU上推理一张640x640图像需要几十毫秒听起来还行但如果跑在车载嵌入式设备上这张模型动辄几百毫秒而且对显存占用太高。所以部署这一步必须做模型转换和推理加速。我的部署流程是PyTorch模型转ONNX再做TensorRT下的FP16量化。ONNX转换这一步要特别注意YOLOv5的模型里有不少操作是PyTorch自定义的转ONNX时可能会产生多余的reshape节点影响后面TensorRT的执行效率。如果遇到这个问题可以尝试把模型输入输出的attribute固定住减少动态维度。转换成TensorRT之后我在同一张3090上做了推理速度对比PyTorch原生推理约22ms/帧ONNX Runtime约15ms/帧TensorRT FP16量化后约6ms/帧。提升幅度相当可观。FP16量化带来的精度损失很小实测mAP只掉了约0.3个百分点对信号灯识别这个任务来说完全可接受。5.2 若干实用推理优化技巧推理阶段还有一个容易被忽视的点图像预处理的方式。检测网络通常用letterbox方式将输入图像等比缩放不足部分补灰边避免目标变形。但letterbox会改变目标在图像中的实际位置坐标所以推理时要把预测框坐标反向映射回原图坐标系这个映射公式一定要写对否则部署出去框位置会偏移。另外滑动窗口做高分辨率图像推理也是一种方案。铁路监控摄像头输出的图像往往不止1080p直接resize到640会丢失大量小目标信息。这种情况下可以把原图切块每块独立推理后再合并NMS。代价是推理时间成倍增加只能用在准实时场景。电力或硬件资源受限的场景下还可以考虑用OpenVINO做CPU推理。Intel的核显CPU跑YOLOv5s 640输入可以做到30ms左右一帧虽然比TensorRT慢但胜在不需要独立显卡适合低成本改造既有监控系统。这个方向我在项目里做了一部分验证后面有时间可以单独展开写一篇。6. 常见问题与排查记录6.1 漏检红灯问题训练完成之后测试发现一个非常要命的场景问题在远距离、大逆光的条件下红灯经常漏检。排查思路分两步走。第一步我先看是不是训练数据里缺少这种“逆光远距离红灯”的样本。统计结果证实了猜测这类样本占比不到2%。解决办法是专门针对这种场景做数据扩充——手动收集逆光时段的数据并补充到训练集里。第二步在模型层面调低类别置信度阈值。YOLO系列默认的conf_thres是0.25对于信号灯检测这种宁可误检也不能漏检的场景我把阈值降到了0.1。代价是误检数量确实变多了但配合后处理的“同一位置连续多帧确认逻辑”可以把误检率压到可接受水平。6.2 夜间灯光泛光导致误检夜间的信号灯存在一个特殊现象——发光灯体会出现泛光光线在镜头里晕开一大片跟灯体实际物理尺寸严重不符。模型一开始会把整个泛光区域都识别成灯预测框比真实灯体大好几倍。解决这个问题的最直接方案是在数据增强阶段增加模拟泛光的数据变换。我在HSV扰动之外额外加了一个操作对部分训练样本做高斯模糊叠加同时随机放大亮度通道的值模拟夜间灯光扩散效果。经过这样处理后模型对泛光区域的预测框收敛显著明显更贴合真实灯体位置。另外一个排查技巧是在推理阶段对预测框做宽高比过滤。铁路信号灯基本上是圆形或接近圆形宽高比接近1。如果预测框宽高比超过2.0以上基本可以断定是误检直接丢掉即可。6.3 雨雾天气识别率下降的实测体验雨雾天气是信号灯识别最难对付的场景。雾气导致灯体颜色衰减严重红灯看起来偏灰粉黄灯看起来偏白。我在雨雾样本上的实测mAP掉到了70%左右跟晴天场景差了20多个点。这个问题的根治方案是采集大量雨雾天气的样本加入训练但实际条件不可能等下雨再去采集。退而求其次的替代方案是做“雾化增强”合成数据——用暗通道先验原理对晴天样本做合成雾化模拟按不同浓度增强后送入训练。我的实验数据显示加入雾化增强后雨雾测试集mAP从70%提升到了76%有效但有限。真正想做得更好还是需要真实的恶劣天气数据。7. 常见问题速查表问题现象可能原因解决方案Loss下降但mAP不涨训练集loss很低验证集mAP停滞训练集与验证集分布差异大检查数据划分是否有同源泄漏重新划分数据集远距离红灯漏检100米外红灯检测不到训练样本中远距离小目标占比低扩充远距离样本降低conf阈值增加高分辨率检测头夜间泛光导致框过大预测框远大于灯体训练集中缺乏泛光样本数据增强加入高斯模糊亮度扰动推理时做宽高比过滤黄灯召回率低黄灯类别mAP显著低于其他黄灯样本过少做类别加权、提高增强倍率重点采集黄灯时段数据TensorRT推理报错转换时出现不支持的算子ONNX动态维度或冗余节点固定输入输出维度修复ONNX图结构雨雾天气精度暴跌雨雾测试集mAP大幅下降训练集缺少雨雾样本用暗通道算法合成雾化增强数据数据标注不一致验证时部分预测框偏大误把背板也标进框重新统一标注规范框只能贴紧发光区域8. 写在最后的实操建议这个项目从数据采集到模型部署完整走下来花了三个多月。如果让我做一次复盘有一个建议最值得分享给想复现这个项目的朋友第一把时间花在数据上永远比花在调参上划算先认真分析和补齐数据分布再考虑改网络结构第二小目标检测的关键突破口往往在网络输入分辨率和检测头的设计上与其在Attention上花大量功夫不如先把这两个基础维度的效果做到位。最后再补充一个容易被人忽视的细节。训练和推理时对图像做的预处理必须保持一致。比如训练时letterbox填充用的是灰色像素值114推理时如果用了0或者255去填充模型效果会显著变差。这种细节问题很难通过看日志发现但排查起来又特别简单——对比训练和推理代码中的预处理函数最好抽出来封装成同一个模块复用。这个习惯帮我省了很多排查问题的时间也一并分享给各位。本文还有配套的精品资源点击获取
返回列表