基于U-Net与MobileNetV3的无人机铁路障碍物图像分割系统

简介本资源是一套面向人工智能与计算机视觉初学者及铁路智能运维从业者的无人机铁路障碍物图像分割实践系统聚焦于利用深度学习技术实现复杂场景下的障碍物精准识别与定位。压缩包共25个文件含19张标注/测试用PNG图像、4个核心Python脚本train.py、val.py、predict.py、ui.py、1份README.md说明文档及1份详细设计说明DOCX文档整体大小为4.19MB结构清晰便于快速复现实验流程与模型推理。资源已获37人学习下载涵盖数据预处理、CNN模型训练、分割结果可视化及简易交互界面等完整环节特别适合作为课程设计、毕业项目或行业轻量化部署的参考范例。读者可直接运行代码完成端到端图像分割任务并基于提供的图像样本拓展训练集、调优网络参数或适配不同障碍物类别。1. 项目概述与整体思路1.1 这个系统解决什么问题先说说我为什么要做这个项目。无人机在铁路场景的巡检其实是一个非常典型的“看起来简单、做起来全是坑”的活儿。很多人觉得无人机飞上去拍一圈、回来看看视频就行但真正落地的时候你会发现几百公里的铁路沿线靠人眼从几个小时的高清视频里找障碍物比如落石、倾倒的树木、非法入侵的施工机械、甚至横跨线路的废弃横幅效率和准确率都低得吓人。我看到太多巡检团队带着大疆的飞机出去飞一天回来三个人对着视频看三天最后眼睛都快瞎了还容易漏掉关键隐患。所以我做了这套“无人机铁路障碍物图像分割系统”。它的核心任务不是做目标检测框那是另一套思路而是做像素级的图像分割——简单说就是把画面里每一个属于“障碍物”的像素点都标出来而不是只画一个框。为什么要这么干因为铁路巡检的障碍物形状极不规则一棵被风吹倒的树它的树枝是发散的一块从山坡滚下来的石头轮廓是带棱角的一个横在轨道上的工程机械形态和背景的区分度极低。检测框在这种场景下很容易把背景也框进去或者框不住完整的障碍物边缘导致后续的距离测算和尺寸估算全都不准。分割的好处是能拿到精确的轮廓配合无人机的GPS坐标和高度信息能推算障碍物的实际大小和侵占轨道的范围这对调度中心的应急处置决策至关重要。这套系统适合谁用呢一是铁路工务段、供电段的技术人员他们需要把无人机巡检从“录像回看”升级成“自动识别”二是做无人机行业应用开发的技术团队想在自己的巡检平台里集成障碍物识别能力三是搞计算机视觉的同学想找一个真实工业场景来练手图像分割模型。我打包的那个zip里有完整的代码、训练好的权重、标注好的数据集示例按本文的步骤走基本上一天之内能跑通整个流程。1.2 技术方案整体架构整个系统的技术栈是这样的图像分割部分用的是改良版U-Net作为基础网络主干backbone换成了轻量化的MobileNetV3而不是原始的VGG或者ResNet34。为什么这么换因为无人机是嵌入式平台算力极其有限。你去看看大疆的M300或M350的机载算力哪怕是加了妙算或者PSDK外挂跟服务器上的显卡完全不是一个量级。我要的是模型在低算力下也能实时推理同时保留分割精度。MobileNetV3的深度可分离卷积能大幅减少参数量配合U-Net的跳跃连接结构能在精度和速度之间取得一个工程上可接受的平衡。训练阶段用服务器我实测用一块RTX 3090显存24GB推理阶段会导出成ONNX然后转TensorRT在机载端的推理延迟能压到30毫秒以内。整套系统还包括数据预处理、标注工具链、模型训练脚本、后处理模块比如连通域分析、轮廓提取、面积计算和与无人机地面站通信的接口。下文会把每个部分的核心细节拆开讲包括我踩过的坑和最后的解决方案。整个流程可以概括为无人机航拍采集视频流 - 逐帧抽帧 - 图像分割模型推理 - 后处理得到障碍物掩码和像素面积 - 结合无人机遥测数据计算实际尺寸 - 回传地面站并告警。这个流程里每一步都有非常多的细节比如抽帧频率怎么定、分割结果怎么去噪、尺寸怎么换算下面逐个展开。2. 图像分割模型选型与原理2.1 为什么是U-Net而不是YOLO系列很多朋友一上来就问我现在YOLOv8这么火你为啥不用它用YOLOv8-seg也能做实例分割啊。这里其实有一个非常关键的工程场景差异。YOLO系列做的是实例分割它的设计目标是把图像里多个独立的个体分别分割出来比如图上有三辆车它给你三个不同的实例mask这种场景下模型会关注“哪块像素属于哪个个体”。但铁路障碍物分割这个场景我更关心的是“哪些像素不是铁轨和正常路肩”也就是语义层面的异常区域提取。你说轨道上有一堆碎石那到底是一个障碍物还是十一个障碍物对于调度决策来说它就是一个需要清障的隐患区域我只需要整片的掩码mask就够了。U-Net天生就是做语义分割的它输出的是一个跟原图同尺寸的概率图每个像素点一个类别概率这种表示方式更适合障碍物侵占区域的识别。另一个原因是样本量。U-Net是出了名的“小数据友好型”网络它的跳跃连接结构能让浅层的空间细节直接传递到深层哪怕训练样本只有几百张也能训练出可以用的模型——这在工业场景太重要了因为铁路障碍物的真实样本本来就稀少你不可能像互联网公司那样攒几百万张图。YOLO系虽然也强但想让它在小样本下稳定收敛调参成本会高很多而且一旦遇到没见过的障碍物形态边界很容易飘。2.2 网络结构拆解与改动点我用的改良版U-Net结构是这样的这里只讲关键改动完整代码在zip的models/unet_mobilenetv3.py里编码器部分Encoder把MobileNetV3-Large的前四个stage作为特征提取器分别输出四个不同分辨率的特征图尺寸从H/2到H/16。这四个特征图会通过跳跃连接送入解码器。原来的U-Net整个编码器都是用卷积堆出来的我换成了MobileNetV3之后参数量从30M降到了4.2MFLOPs也降低了将近7倍。在3090上训练一个epoch从原来的6分钟缩到了1分半效果还更稳了。解码器部分Decoder保持U-Net经典的4层上采样结构但每一层上采样之后接的不是简单的两个3x3卷积而是用了一个改进的深度可分离卷积 SE注意力模块的组合。SE注意力模块的引入是我做对比实验以后决定保留的它能在通道维度上自动学习哪些特征更重要。比如树枝这种细长形物体的边缘特征在第3层、第4层解码的语义特征中会不断增强而不容易被上采样过程稀释掉。还有个细节原始的U-Net输出层用的激活函数是Sigmoid配合BCE损失函数做二分类分割。我这里改成了Sigmoid Lovász-Softmax损失函数的组合。Lovász损失主要是针对IoU指标做直接优化它对小目标比如远处的一小块落石特别友好。纯BCE损失会把压倒性的背景像素学得特别好但对小障碍物区域却容易漏掉。两者加权组合以后我在测试集上的mIoU从0.71提升到了0.79这个提升在铁路场景里非常值钱。2.3 输入尺寸与推理分辨率的取舍无人机拍回来的画面一般是3840x21604K或者2720x15302.7K你不可能把整张4K图直接塞进网络里训练显存不够推理速度也扛不住。我最终把训练输入尺寸定为512x512推理时用滑窗 拼接的方式处理大图。为什么不是768或者1024我实测了三个分辨率512x512的最佳mIoU是0.79768x768能到0.82但推理时间直接从28ms涨到了75ms1024x1024能到0.84但已经要105ms了而且在机载端的算力上根本跑不动。对铁路巡检来说我的判断是5%的精度提升不值得牺牲3倍的速度。尤其无人机是在运动的检测延迟太高意味着发出告警时飞行器已经飞过障碍物几十米了。滑窗策略也要讲一下。我用的是overlap64像素的窗口滑动也就是相邻两个窗口有64像素的重叠区域。这样做不是为了好看而是为了消除边缘拼接伪影。如果不设重叠分割结果在窗口交界处经常会出现明显的断层线因为网络对图像边缘区域的感受野是不完整的预测置信度会降低。我会把重叠区域的预测结果做一个高斯加权融合中心像素权重大、边缘权重小拼接出来就自然多了。3. 数据集采集、标注与预处理3.1 无人机航拍数据的采集要点图像分割模型的效果60%以上取决于训练数据的质量。我在这里踩过很多坑现在把关键经验分享出来。首先是采集高度。铁路障碍物的尺度差异极大落石可能只有30cm见方倾倒的树可能有十几米长。如果你只在固定的50米高度拍模型只能学到一种尺度的特征。我的做法是分三个高度段采集低空15-25米用来捕捉小块障碍物、中空40-60米覆盖一般巡检场景、高空80-120米获取全局路况。三个高度段的数据混合训练后模型对尺度的鲁棒性明显提升。其次是角度。无人机正射视角镜头垂直向下方便做尺寸换算但很多障碍物从正上方看并不明显比如斜靠在护栏上的树枝正射看它可能只是一根线。所以我会在巡检路线中加一部分45度斜视角的采集数据让模型认识障碍物的“侧面形态”。斜视角数据在分割精度上会略低一些因为形变大但它能显著降低漏检率。最后是光照。铁路巡检经常会遇到清晨、黄昏、阴影遮挡这些光照条件。我在采数据的时候同一段线路会安排早中晚各飞一次保证数据集里有足够的弱光样本。模型训练时我还做了颜色抖动和随机亮度的数据增强这样即使在光线条件差的场景下分割结果也不会崩掉。3.2 标注方案与工具链标注工具我用的是LabelMe配合自写的半自动标注辅助脚本。LabelMe输出的是JSON格式的多边形标注我会写脚本把JSON转换成COCO格式的RLE编码Run-Length Encoding作为训练标签。RLE编码在存储上非常紧凑一张4K图的分割掩码存成RLE可能只有几KB比PNG掩码图小了两个数量级。标注的类别设计也值得细说。我没有只设一个“障碍物”类别而是分了四类rock落石、tree倾倒的树木/树枝、vehicle非法入侵的工程机械、debris其他杂物比如塑料布、横幅、废弃路料。这样模型能学到底层特征的差异石头是硬边缘、纹理粗糙树枝是细长形、纹理复杂比笼统地归为“障碍物”要容易收敛得多。最后融合的时候我会把四类的概率图做逐像素取最大值合成一张最终的障碍物区域图。我给初始版本准备的训练集是1200张标注好的图像其中800张来自实测飞行采集400张来自公开无人机数据集主要是城市和公路场景做预训练迁移。预训练之后再用实测数据微调30个epoch效果比直接拿1200张实测图从头训练好了不少。3.3 数据预处理的几个细节数据增强除了常规的随机翻转、旋转、缩放之外我还加了一个铁路场景特有的增强——轨道区域遮挡模拟。因为实际巡检中铁轨本身在画面里是强结构特征模型容易把“轨道区域”当成先验知识来推断障碍物的位置这会导致泛化能力下降。比如它学会了“铁轨中间出现物体才是障碍物”遇到曲线轨道或者道岔区域就抓瞎了。所以我会随机在标注掩码附近加一些黑块遮挡强迫模型去学更鲁棒的纹理特征而不是位置先验。另外无人机图传画面偶尔会出现运动模糊尤其是快速变焦或转向时。我用了随机高斯模糊 随机运动模糊作为数据增强项kernel大小在3到7之间随机。这能让模型对模糊图像的鲁棒性提升实测在模糊验证集上的mIoU从0.64提升到0.72。4. 训练流程与关键参数4.1 环境准备与zip包解压先讲最基础的从我发布的“无人机铁路障碍物图像分割系统.zip”开始你要在Linux服务器上把代码跑起来。很多同学第一步就卡在解压上我特意把解压命令写清楚# 解压到指定目录 unzip 无人机铁路障碍物图像分割系统.zip -d /data/projects/railway_seg # 或直接原地解压 cd /data/projects unzip 无人机铁路障碍物图像分割系统.zip # 查看压缩包内容确认文件结构完整 unzip -l 无人机铁路障碍物图像分割系统.zip如果你在Windows上解压后放到Linux上偶尔会遇到文件权限问题可以用chmod -R 755把整个目录的权限处理一下避免启动脚本时报Permission denied。代码依赖我用的是conda环境Python版本锁定3.9PyTorch 1.13.1 CUDA 11.7。为什么不用更新的PyTorch 2.x因为TensorRT转换工具链对2.x的支持当时还不完全稳定我项目落地的时候用1.13.1最省心。完整的环境创建命令在zip里的requirements.txt和environment.yml里都有直接执行conda env create -f environment.yml conda activate railway_seg4.2 训练参数的核心配置训练脚本的核心参数我写在configs/train_config.yaml里有几个关键参数直接决定模型效果我逐个解释。学习率策略初始学习率0.0008batch_size16时采用CosineAnnealingWarmRestarts调度器T_010、T_mult2。为什么用带热重启的余弦退火因为铁路障碍物数据集的类别不均衡严重用余弦退火的周期性重启能帮模型跳出局部最优我实测比固定学习率或者step decay的收敛效果好3%-5%。损失函数权重Lovász损失和BCE损失按7:3加权。这个比例是我试出来的。Lovász比例太高训练早期loss会震荡得很厉害不太好收敛太低则小目标分割效果提升有限。你也可以从5:5开始试然后根据验证集的IoU微调。batch_size这个直接取决于你的显存。309024GB在输入512x512、batch_size16的情况下会吃满显存。如果显存只有12GB比如3080把batch_size降到8同时把学习率同步降到0.0004。训练轮数我的经验是从头训练的话150-200个epoch能达到稳定收敛如果你用我提供的预训练权重继续微调40个epoch就够了。早停策略设的patience20也就是说验证集IoU连续20个epoch不涨就自动停。4.3 训练过程中需要盯的几个曲线训练不能光把脚本扔进去就跑。我会在训练过程中同时盯这几条曲线train_loss和val_loss如果val_loss先降后升说明过拟合了赶紧加正则或者减少epoch。val_miou这是最重要的指标低于0.7说明模型还没学好优先检查数据标注有没有问题。val_rock_iou和val_tree_iou分类别IoU如果某一类的IoU明显低于其他类大概率是该类样本太少去补充数据比调参有效得多。有一次我训了一个晚上起来发现val_miou停在0.65死活上不去后来查了一下是标注多边形在转换RLE的时候有个1的像素偏移bug导致mask和原图对不上。这种低级错误在图像分割项目里很常见排查方式是随机抽几张训练样本把原图和掩码重叠可视化出来看一眼就能发现问题。5. 部署与推理链路5.1 导出ONNX与TensorRT加速训练好的PyTorch模型不能直接扔到无人机机载端去跑主要有两个原因一是PyTorch的推理框架太重嵌入式环境下装起来麻烦二是推理速度不够PyTorch的Eager模式在边缘设备上跑MobileNetV3U-Net实测要80-100毫秒一帧没法用。我的做法是先转ONNX再转TensorRT FP16精度的engine。# 导出ONNX import torch import onnx from models.unet_mobilenetv3 import UNetMobileNetV3 model UNetMobileNetV3(num_classes4) model.load_state_dict(torch.load(checkpoints/best_model.pt, map_locationcpu)) model.eval() dummy_input torch.randn(1, 3, 512, 512) torch.onnx.export( model, dummy_input, model.onnx, input_names[input], output_names[output], opset_version13, dynamic_axes{input: {0: batch}, output: {0: batch}} )然后转TensorRT在NVIDIA Jetson设备上用# 使用trtexec工具转换 trtexec --onnxmodel.onnx --saveEnginemodel.engine --fp16 --workspace8192转完以后FP16精度下的推理延迟在Jetson Orin NX上实测大约28毫秒512x512输入完全满足实时的需求。需要提醒的是TensorRT的engine是硬件绑定的你在Orin上转出来的engine文件不能搬到Xavier上直接用必须在目标设备上重新做一次转换。5.2 后处理从掩码到障碍物信息分割模型输出的是一张4通道的概率图这不直接是最终结果还需要一系列后处理。第一步是argmax获取每个像素的类别索引得到一张单通道掩码图。然后做一次中值滤波kernel5去除零星噪点——因为无人机高空中看到的细小纹理容易被误判为障碍物。这一步很便宜但是效果显著能过滤掉大量假阳性区域。第二步是连通域分析把属于同一块障碍物的像素聚合起来。我用的是OpenCV的connectedComponentsWithStats函数它能输出每个连通域的边界框、面积和中心点坐标。这一步输出的是“一个障碍物实例”的级信息而不是逐像素的掩码。第三步是尺寸换算这是把像素面积转换成实际物理尺寸的关键环节。核心公式如下实际尺寸 像素尺寸 / (焦距 / 飞行高度)展开来说对于一个焦距为f毫米、传感器像素尺寸为s微米/像素的相机每个像素对应的地面物理尺寸为地面采样距离 (飞行高度 × s) / f。比如大疆M300的XT2相机焦距12mm像素尺寸3.45μm飞行高度50米那每个像素对应的地面距离大约是1.44厘米。如果你算出来的掩码像素总面积是5000像素那么实际面积大约是5000 × (0.0144m)² 1.04平方米。这样调度人员就能直观地评估障碍物的规模。5.3 与无人机地面站的联动分割结果不只是存在飞机SD卡里的我会通过RTMP推流到地面站同时把告警信息通过MAVLink协议的STATUSTEXT消息发送出来。地面站收到以后可以自动弹窗告警并附带截屏图片和GPS坐标。这套接口的代码在zip的communication/mavlink_bridge.py里你只需要改成你实际用的地面站地址和端口即可。另外一个细节也可能是很多人忽略的在做实时推理的时候不要把每一帧都处理一般3-5帧取一帧做检测就足够了。因为无人机是连续运动的连续帧之间的画面重叠度非常高每帧都推理纯属浪费算力还容易造成检测结果抖动——同一块障碍物刚告警完下一帧又告警一次调度中心那边会被骚扰死。我的策略是“低速巡检2帧抽1帧、高速巡检5帧抽1帧”这样既能保障不漏检又能减少重复告警。6. 常见问题与排查实录6.1 高频问题速查表我在开发这套系统的过程中遇到并解决了大量问题。下面整理几个最典型的大家遇到可以直接对照排查。问题现象可能原因解决方案解压失败报file is not a zip file下载文件不完整或者被浏览器/网盘截断重新下载用zip -FF damaged.zip --out fix.zip尝试修复invalid zip archive: could not find eocdzip包结构损坏或磁盘满先检查磁盘空间换用7z x重试解压训练时显存溢出CUDA OOMbatch_size过大调小batch_size到8或4同步调低学习率推理结果全是背景没有障碍物输入图像的归一化参数不对检查预处理代码要用训练时的mean[0.485,0.456,0.406]和std[0.229,0.224,0.225]分割掩码有大量细小噪点没有做中值滤波或置信度阈值过低后处理加滤波把置信度阈值从0.3提到0.5同一障碍物反复告警没有做去重逻辑加上时间窗口去重同一位置5秒内只告警一次Jetson上推理速度只有8FPSTensorRT未生效还在用PyTorch推理确认加载的是.engine文件而非.pt文件6.2 我踩过的三个大坑第一个坑是zip压缩包内部路径包含中文。在Linux上解压后某些版本的PyTorch在读取中文路径下的数据时会报编码错误。后来我把整个项目根目录改成纯英文路径同时内部数据集结构也是英文命名问题就消失了。所以zip里的代码路径默认都是英文这个不是偷懒是实战教训。第二个坑是无人机相机参数标定。我第一次做尺寸换算的时候直接参考网上别人给的焦距数值结果算出来的障碍物面积比实际大了整整50%。后来才意识到每台相机的传感器参数不同不能想当然。正确做法是去大疆开发者文档查你这个机型的具体参数最好还能通过已知尺寸的地面参照物比如轨枕间距标准是0.6米来反推地面采样距离这个校准方法最靠谱。第三个坑是模型在夜间/逆光下的误判率飙升。一开始我只在白天的数据上训练结果傍晚试飞的时候长阴影区域的铁路扣件和碎石被大量误报成障碍物。解决办法是我又补了一批逆光和阴影的数据同时在后处理中增加了一个基于HSV颜色空间的“阴影区域概率衰减”操作对低亮度区域的分割置信度乘以一个0.8的折扣系数误报率立刻下降了不少。6.3 模型调优的两个独家技巧分享两个我验证过非常管用的调优技巧。第一个技巧是“伪标注迭代”用当前模型对大量未标注的巡检视频做推理把置信度较高的分割结果自动生成伪标签然后人工抽检修正再混入训练集继续训练。这个技术能非常有效地扩充数据集。我第二版模型用的是这个方案把训练集从1200张扩充到了接近5000张mIoU又涨了0.04。你需要控制伪标注的置信度阈值不要低于0.95否则会把模型的错误当成老师教回来。第二个技巧是“混合精度训练配合EMA”训练时开AMP混合精度模型参数维护一份指数移动平均EMA推理时用EMA权重而不是最后一步的权重。这小改动看着不起眼但实测能把最终的mIoU再提升1.5%左右而且几乎不增加训练时间。EMA的好处是相当于对模型权重的历史做平均能有效抑制训练后期权重震荡带来的性能波动尤其适合小数据集场景。代码实现可以看zip里的utils/ema.py。7. 系统扩展与后续规划这套系统目前已经在我负责的巡检项目里稳定运行了几个月处理了上百架次的飞行数据。不过说实话它现在只能算“1.0版本”还有不少可以演进的地方。目前我正在做的一个方向是把检测和分割统一到同一个网络里也就是多任务框架——既输出障碍物的类别和位置框又输出像素级分割掩码。这样在告警时可以更明确地告诉调度人员“前方200米线路右侧有一棵倒树侵限范围约3.2平方米”决策信息量比单纯的掩码大得多。另一个方向是引入时序信息用视频前后帧的光流场来辅助判断障碍物是静止的还是移动的比如闯入线路的动物这对报警的优先级排序非常有用。如果你也想在类似场景里落地图像分割我的建议是先别急着追新模型把数据和后处理做扎实了效果远比你换一个大模型来得好。U-Net这个级别的网络在工业场景里把精度做到80%以上mIoU完全够用。核心还是你能不能把“从图像到决策”的链路走通——分割只是中间一环后面的连通域分析、尺寸换算、告警联动、数据管理每一步都在为最终的用户价值服务。另外提醒一句项目发布出来的zip包由于包含代码和训练好的权重体积大概有2.4GB下载的时候注意别被网盘中转服务器截断。如果你在运行时遇到什么奇怪的问题建议先打开logs目录下的日志文件看看报错信息大多数问题都能从报错堆栈里定位到具体模块。本文还有配套的精品资源点击获取