轨道缺陷检测这个方向最近两年在工业视觉里挺火但真正能把模型跑通、跑稳、还能落地的人并不多。我接触这个需求最早是帮朋友处理钢轨高速巡检图像客户给了一批包含裂纹、剥落、光带异常的路面照片要求用算法把缺陷区域逐像素标记出来。当时第一个想到的就是UNet原因是这个任务本质上是图像分割而UNet在医学影像分割里已经验证得很充分再加上轨道缺陷样本通常很小几十张图就能微调出不错的效果。今天这篇东西就是把完整的一套流程摊开来讲从数据集怎么搞、模型怎么写、训练怎么收敛到最后怎么评估效果全部按实际操作的顺序来附上可直接用的代码适合正在做检测、分割方向的学生和初级算法工程师。RSDDs数据集是学术界公开的轨道表面缺陷数据集网上能搜到不少分享包。整份数据包含快速道和普通道两个大类每类下面又细分了裂纹、剥落、光带异常等多种缺陷。整个分割任务要做的事情非常简单直白给一张轨道表面灰度图输出一张和原图等尺寸的mask图缺陷区域标为白色背景标为黑色。就这么个看起来朴素的任务实际做起来坑非常多比如裂纹这种目标细长而且窄占整张图的像素可能不到1%模型一不小心就把所有像素都预测成背景因为整体准确率照样99%。这也是我为什么不建议直接套通用分割代码而是要专门调整损失函数和训练策略的原因。做这个项目前我把网上能搜到的相关笔记翻了一遍发现很多文章只停留在“下载了数据集、训练了UNet、发了两张预测图”的阶段缺少完整可复现的代码和调参细节。所以说这篇更像一份工作记录把我实际跑通的流程和踩过的坑一并写出来。阅读的时候建议找个能跑深度学习的机器跟着代码走一遍比干看文字印象深得多。1. 为什么用UNet做轨道缺陷检测——从任务本质说起1.1 轨道缺陷检测到底难在哪轨道表面的缺陷和图里其他内容的区别很多人第一反应是“缺陷不就是表面上的裂纹破损吗用模板匹配或者边缘检测不就行了”。这个想法在理想情况下成立但真实工况下非常脆弱。首先是光照户外拍回来的图钢轨表面反光特别严重裂纹在强反光下对比度极低其次是背景干扰轨面的划痕、轮轨接触产生的黑色碾光带、枕木和道砟的纹理这些都会在特征上和真实缺陷混淆再次是缺陷形态变化大裂纹可能只有几个像素宽剥落区则呈现不规则块状边缘模糊程度完全不同。传统的图像处理方式比如Canny边缘检测加形态学操作对付固定环境、单一类型的纹理异常还行一旦换成不同区段、不同磨损状态的钢轨图像阈值根本没法统一调整。某次我在一张低对比度裂纹图上调好了Canny阈值换到另一张强反光的图直接检测出一大片假边缘效果没法用。这也是我转向深度学习分割的直接原因——目标检测框是矩形框得住剥落却框不住细长裂纹而分割能够逐像素判断类别天然适合这种细长形目标。1.2 选UNet而不选其他分割网络的核心理由图像分割领域的网络一抓一大把比如SegNet、DeepLabV3、PSPNet各有所长。但在轨道缺陷这个场景下我更推荐UNet原因有三个。第一样本量小。轨道缺陷数据集的公开部分规模很小RSDDs加起来大概就几百张图如果直接上DeepLab这种带空洞卷积的大网络或者更重的Transformer结构参数量大小数据很难训好泛化也容易出问题。UNet的参数量相对可控即使不加载预训练权重从零开始训也能在几十上百张图上收敛到可用的效果。第二细节边界好。UNet的编码器逐层提取高层语义特征同时保留详细空间信息的跳跃连接把低层特征直接拼到解码器对应层。对于轨道表面的细小裂纹低层纹理信息是救命稻草有了跳跃连接网络在解码阶段能同时兼顾整体位置和局部细节分割出来的缺陷边缘更连续。这一点对比过实际效果不加跳跃连接的版本裂纹中间经常断成几截加了跳跃连接后完整性明显改善。第三结构简单改起来方便。UNet的U型架构足够直白训练、部署、调试都不费劲。我在做项目时需要把中间特征层拿出来分析UNet可以很轻松地在解码路径上加分支输出换成其他复杂网络就没这么顺手。团队里两三天就能上手后续做缺陷面积统计、缺陷类别再分直接在UNet的mask基础上扩展就行。2. RSDDs数据集来了——先搞懂数据长什么样2.1 RSDDs数据集的几组关键数字RSDDs全称是Rail Surface Defect Dataset最早由国内某高校实验室公开是轨道缺陷检测领域用得比较多的基准数据之一。整个数据集分成两个子集RSDDs Type-I对应快速道Fast railRSDDs Type-II对应普通道Normal rail。快轨数据主要包含高速铁路区段的钢轨图片普通道数据包含普通列车轨道的图片。从缺陷类别上看两张子集各有侧重但总体上涵盖了这几类常见缺陷轨面裂纹包括横裂纹、纵裂纹、剥落掉块、轨面擦伤、光带异常等。不过需要说明的是RSDDs的官方标注并非像有些语义分割数据集那样做像素级精细标注它提供的是缺陷位置的矩形框标注也就是目标检测的标注格式。所以如果我们想直接拿来跑UNet这样的分割模型就得把矩形框转成像素级mask做法是遍历每张图对应的标注框把落在框内的像素标记为前景。这种方式不算像素级精标准确但胜在自动化处理成本低实践中完全够用。从数量上说Type-I大约有67张带缺陷的图Type-II大约有128张带缺陷的图加上一些无缺陷图像总体规模在两三百张这个量级。训练分割模型时如果直接整张图丢进去训练样本数会非常少所以常规做法是切片Patch。把一张512×512的图切分成若干128×128或256×256的小块从中筛选出包含缺陷的小块用于训练这样样本量一下子就扩充了十几倍。具体的参数和目录整理方式下面一节直接给操作。2.2 下载与目录整理RSDDs数据集因为发布有一段时间了有些原始官网链接可能失效但现在很多研究者在GitHub仓库里做了镜像整理。你可以直接搜索“RSDDs dataset download”找到包含图像和标注文件的仓库或学术资源页下载下来。如果对方提供了百度网盘或Google Drive的分享链接按说明提取即可。无论如何下载完之后我们需要把文件整理成统一的结构便于后续脚本读取。我的习惯是根目录建一个data文件夹内部按如下方式组织data/ ├── Type_I/ │ ├── IMG/ │ │ ├── 01.bmp │ │ ├── 02.bmp │ │ └── ... │ └── GT/ │ ├── 01.bmp │ ├── 02.bmp │ └── ... ├── Type_II/ │ ├── IMG/ │ │ ├── 01.bmp │ │ ├── 02.bmp │ │ └── ... │ └── GT/ │ ├── 01.bmp │ ├── 02.bmp │ └── ... └── split/ ├── train.txt └── val.txt这里我习惯把图像和标注分开放在IMG和GT两个目录文件名保持一一对应。GT可以是标注人员画好的mask图也可以是程序转出来的mask图。如果下载到的GT是xml或txt格式的矩形框标注可以参考下面的Python转换脚本通过PIL和numpy生成对应的maskimport os import cv2 import numpy as np from PIL import Image, ImageDraw def rects_to_mask(img_path, gt_path, mask_save_path): # 以Original的图像尺寸为基准 img cv2.imread(img_path) h, w img.shape[:2] mask Image.new(L, (w, h), 0) draw ImageDraw.Draw(mask) # 伪代码实际需要解析txt/xml的每个矩形框然后按坐标画矩形 # for box in boxes: # draw.rectangle([x1, y1, x2, y2], fill255) mask.save(mask_save_path)需要注意的是不同来源提供的标注文件单位可能不同有的给的是归一化坐标有的是绝对像素坐标有的是左上角和右下角有的是中心点加宽高。读取之前务必手工打开一张看看避免后期全部mask错位。2.3 训练数据增强把有限的样本“变”多两三百张图直接训练模型很容易过拟合在训练集上表现很好换一批真实图片就崩。我的做法是离线做数据增强也就是在训练前先把图像和mask一起做变换生成一批增强后的样本保存到新目录而不是在训练循环里在线增强。这样做的优点是方便检查增强效果也能避免训练过程中CPU和GPU速度不匹配影响训练进度。常用的增强方式包括随机水平翻转轨道图像对称性还算明显翻转不会破坏语义、随机旋转90度、随机亮度对比度调整、添加高斯噪声、随机裁剪以及整体缩放。但轨道缺陷检测里要特别注意裂纹这类细长形目标经不起任意角度旋转和拉伸。比如随机旋转30度裂纹的几何形态和真实情况严重不符非等比例缩放会让裂纹变胖或变细破坏物理特征。所以我实操时只用90度的整数倍旋转、水平翻转、小范围缩放和亮度抖动。具体的增强代码不依赖外部库也能写比如用OpenCV的warpAffine配合旋转矩阵或者在PIL里直接使用transpose。但如果追求效率可以直接用albumentations库它对图像和mask同步增强做得很好。我常用的增强管道像下面这样import albumentations as A train_transform A.Compose([ A.HorizontalFlip(p0.5), A.RandomRotate90(p0.5), A.RandomBrightnessContrast(brightness_limit0.2, contrast_limit0.2, p0.5), A.GaussNoise(var_limit(10.0, 50.0), p0.3), A.Resize(height256, width256, always_applyTrue) ])在增强时需要保证图像和mask使用完全相同的随机参数这样才能保证位置对齐。albumentations这类库已经封装好这个同步逻辑直接调用即可。如果自己写代码千万记住在随机数生成上保持同一个状态。3. 手把手搭建UNet训练环境与代码3.1 训练环境与硬件配置开始写代码之前先把环境准备好省得到时候装包装到怀疑人生。我的训练环境配置如下仅供参考操作系统Ubuntu 20.04 / Windows 11 都可以Python版本3.8 或 3.9深度学习框架PyTorch 1.10下文代码基于PyTorch 2.0写CUDA11.x或12.x均可取决于你的显卡驱动显卡建议至少4GB显存我这边用了一张16GB的卡训练不存在显存压力安装依赖可以直接用pippip install torch torchvision opencv-python pillow numpy albumentations tqdm如果你用的是GPU版PyTorch记得去官网选择匹配你本机CUDA版本的那条安装命令不要直接pip install torch装上CPU版本否则训练速度会慢得让人崩溃。CPU版也不是不能跑只是跑一个epoch的时间足够你打两把游戏再回来。3.2 UNet模型结构实现PyTorch完整代码UNet的核心结构是编码器-解码器结构编码器由4到5个block组成每经过一个block空间尺寸减半通道数翻倍解码器则进行上采样空间尺寸加倍通道数减半。在编码器和解码器之间有跳跃连接把对应尺度的特征拼接起来。以下是简洁可用的UNet实现输入输出都是单通道适合二值分割。import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): 两次卷积 BN ReLU def __init__(self, in_ch, out_ch): super(DoubleConv, self).__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, in_channels1, out_channels1, features[64, 128, 256, 512]): super(UNet, self).__init__() self.encoder nn.ModuleList() self.pool nn.MaxPool2d(kernel_size2, stride2) self.decoder nn.ModuleList() # 编码器 for f in features: self.encoder.append(DoubleConv(in_channels, f)) in_channels f # 瓶颈层 self.bottleneck DoubleConv(features[-1], features[-1] * 2) # 解码器反转特征通道缩回去 for idx in range(len(features) - 1, -1, -1): self.decoder.append( nn.ConvTranspose2d(features[idx] * 2, features[idx], kernel_size2, stride2) ) self.decoder.append(DoubleConv(features[idx] * 2, features[idx])) self.final_conv nn.Conv2d(features[0], out_channels, kernel_size1) def forward(self, x): skip_connections [] for enc in self.encoder: x enc(x) skip_connections.append(x) x self.pool(x) x self.bottleneck(x) skip_connections skip_connections[::-1] for idx in range(0, len(self.decoder), 2): x self.decoder[idx](x) skip skip_connections[idx // 2] if x.shape ! skip.shape: x F.interpolate(x, sizeskip.shape[2:], modebilinear, align_cornersTrue) x torch.cat((skip, x), dim1) x self.decoder[idx 1](x) return torch.sigmoid(self.final_conv(x))这段代码里编码器层层提取特征每个DoubleConv内部做的都是卷积、批归一化、ReLU的标准组合。批归一化非常重要它能加快收敛也能减少对初始权重和输入分布的敏感度。解码器的转置卷积负责把特征图逐步放大每次放大后和对应的编码器特征进行通道拼接再经过DoubleConv融合。需要提醒的是原始UNet用两个3×3卷积叠加相当于一个5×5感受野同时减少了参数量。在轨道缺陷分割里缺陷目标较小所以没必要加深太多层四层特征已经足够。如果你希望进一步提升精度可以把第一层的初始通道数设为32或64上面代码默认64。3.3 数据读取与训练循环数据读取阶段需要同时读入原图和mask并且把它们变成相同的尺寸。训练时用小patch256×256左右比较均衡。下面是数据集的PyTorch封装import os import cv2 import torch from torch.utils.data import Dataset class RailDataset(Dataset): def __init__(self, image_dir, mask_dir, image_size256, transformNone): self.image_dir image_dir self.mask_dir mask_dir self.image_size image_size self.transform transform self.images sorted(os.listdir(image_dir)) def __len__(self): return len(self.images) def __getitem__(self, idx): img_name self.images[idx] img_path os.path.join(self.image_dir, img_name) mask_path os.path.join(self.mask_dir, img_name) image cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) # 统一尺寸 image cv2.resize(image, (self.image_size, self.image_size), interpolationcv2.INTER_LINEAR) mask cv2.resize(mask, (self.image_size, self.image_size), interpolationcv2.INTER_NEAREST) # mask只取0/1 mask cv2.threshold(mask, 127, 1, cv2.THRESH_BINARY)[1] if self.transform is not None: transformed self.transform(imageimage, maskmask) image transformed[image] mask transformed[mask] image torch.from_numpy(image).float().unsqueeze(0) / 255.0 mask torch.from_numpy(mask).float().unsqueeze(0) return image, mask这里有一个很容易忽略的细节resize mask时插值方式必须用INTER_NEAREST也就是最近邻插值。如果用线性插值去缩小mask原先是255的像素可能会变成100多甚至0造成边界重影那么loss计算时会产生大量错误梯度导致模型学出模糊的边缘。训练循环部分我更倾向于把BCE和Dice结合作为总损失。BCE负责像素级的收敛Dice损失则直接优化分割区域的重叠程度对小目标更友好。下面这个交叉Dice的实现正是为了应对缺陷小、正负样本不平衡问题def dice_loss(pred, target, smooth1.0): pred pred.contiguous().view(-1) target target.contiguous().view(-1) intersection (pred * target).sum() return 1 - (2.0 * intersection smooth) / (pred.sum() target.sum() smooth) def combined_loss(pred, target): bce F.binary_cross_entropy(pred, target) dice dice_loss(pred, target) return bce dice训练循环就常规操作遍历数据集前向传播算损失反向传播更新权重。以我自己的配置批量大小设为16学习率0.001Adam优化器训练60个epoch整个过程大约需要十到二十分钟。如果你用的是CPU时间会成倍增加建议有条件的还是用GPU。3.4 模型评估与结果可视化训练结束后不能只看训练集loss还需要在验证集上量化指标。语义分割常用的指标包括Dice系数F1、IoU、Pixel Accuracy。由于轨道缺陷区域很小Pixel Accuracy往往虚高没有参考价值。我主要看IoU和Dice这两个指标对预测区域和真实区域的重合程度更敏感。评估代码可以用下面这段def calculate_metrics(pred_mask, true_mask): pred_mask (pred_mask 0.5).int() true_mask (true_mask 0.5).int() intersection (pred_mask * true_mask).sum().item() union (pred_mask | true_mask).sum().item() iou intersection / (union 1e-6) dice (2 * intersection) / (pred_mask.sum().item() true_mask.sum().item() 1e-6) return iou, dice预测时把验证集的每张图和对应mask读入经过模型推理得到概率图阈值0.5转成二值mask。可视化时推荐将原图、标注、预测三张图横向拼在一起这样一眼就能看出哪些地方漏检了、哪些地方是误检。建议保存成图片到result/目录方便后续做报告或者排查。4. 踩坑实录与效果优化4.1 Loss不降或震荡怎么办训练时最糟心的就是loss纹丝不动。我遇到这个情况先从三个方向排查第一检查数据预处理看看mask是否真的和图像对齐很多标注文件存在裁切偏移肉眼看不到但loss会一直很高第二学习率是不是太大或太小Adam默认的0.001已经比较激进配合小数据集时甚至可能震荡我后来降到0.0005就稳定了第三正负样本极度不平衡如果一张图缺陷只占0.1%像素BCE loss会被大量背景像素淹没所以必须配合Dice Loss或者对正样本提权。遇到过最诡异的一次是loss在第20个epoch突然从0.3跳到0.9再也没降回来。排查发现是BN层的均值和方差在batch噪声大的时候波动严重解决办法是增大batch size或者把BN改成GroupNorm。因为GPU显存有限我选择用了GroupNorm版本的UNet明显稳定不少。4.2 裂纹预测断成一截一截的裂纹这种长条形目标有个普遍问题模型预测出的结果经常中间有断裂。分析原因一个是原图上裂纹本身对比度就低中间某段跟背景融为一体网络难以判断另一个原因是下采样太多次小裂纹细节在深层特征里已经丢掉了仅靠高级语义信息补不回来。解决办法我实践下来有三招。第一用切片训练而不是整图训练让模型看到更多细节同时避免resize把裂纹挤没第二把UNet的第一层通道数从64改成32减少下采样信息压缩但这会轻微增加计算量第三后处理环节使用形态学闭运算把距离很近的分段裂纹连接起来比如用OpenCV的cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)我用的kernel是5×5的椭圆核连接效果不错但也容易把间隙很小的正常纹理连成缺陷需要根据实际图片调kernel大小。4.3 从60分到90分我用过的3个提高精度手段如果基础UNet跑通了精度却还差一口气我强烈建议按下面顺序迭代。第一损失函数优化。单独BCE很容易让网络偏向于学习背景单独Dice Loss则在训练初期过于不平滑。我用的组合是BCE Dice在一些小裂纹数据集上比单独BCE能提升3到5个点的IoU。如果你的缺陷特别小也可以尝试Focal Loss DiceFocal Loss能抑制易分类的大背景样本让模型更关注难的缺陷区域效果如何也要结合数据实测。第二预训练编码器。因为RSDDs样本量不够大从头训卷积核容易欠拟合。如果不想引入外部数据可以先用ImageNet预训练的ResNet作为UNet编码器只随机初始化解码器。这个做法在很多分割任务里是常规操作能够显著提升收敛速度和最终精度。在PyTorch中加载方式很简单把UNet的encoder部分替换为torchvision.models.resnet34(pretrainedTrue)并去掉最后的全连接层就好。第三后处理加连通域筛选。模型输出的二值mask里通常会有一些零散的小块误检这些往往是背景中的亮点或锈斑。我们可以借助连通域分析删除面积小于某个阈值的区域。例如设定最小缺陷面积为25像素如果某个连通域面积小于这个阈值就把它置为0。这个后处理十分粗暴有效能够很好地抑制零星噪声但要注意阈值不能设太大否则真正的小裂纹也会被误删。我的做法是先统计验证集标注中所有连通域的面积分布取最小的1%作为阈值下限。4.4 顺手放一套完整的训练主程序从数据加载到推理代码分散在各个小节里可能看起来不够连贯。下面我整合了一段主训练脚本涵盖数据路径配置、训练循环、模型保存和推理可视化你可以直接复制过去改动路径使用。import os import torch import torch.optim as optim from torch.utils.data import DataLoader from tqdm import tqdm import cv2 import numpy as np # 假设你已经把UNet类、RailDataset类、combined_loss函数定义在上方 device torch.device(cuda if torch.cuda.is_available() else cpu) # 配置 train_img_dir data/train/images train_mask_dir data/train/masks val_img_dir data/val/images val_mask_dir data/val/masks batch_size 16 lr 0.0005 epochs 60 save_path checkpoints/unet_rail.pth # 数据与加载 train_ds RailDataset(train_img_dir, train_mask_dir, image_size256) val_ds RailDataset(val_img_dir, val_mask_dir, image_size256) train_dl DataLoader(train_ds, batch_sizebatch_size, shuffleTrue, num_workers4) val_dl DataLoader(val_ds, batch_sizebatch_size, shuffleFalse, num_workers4) # 模型与优化器 model UNet(in_channels1, out_channels1, features[64, 128, 256, 512]).to(device) optimizer optim.Adam(model.parameters(), lrlr) best_iou 0.0 for epoch in range(epochs): model.train() total_loss 0 for images, masks in tqdm(train_dl, descfEpoch {epoch1}/{epochs}): images, masks images.to(device), masks.to(device) preds model(images) loss combined_loss(preds, masks) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() avg_loss total_loss / len(train_dl) # 验证 model.eval() epoch_iou 0 with torch.no_grad(): for images, masks in val_dl: images, masks images.to(device), masks.to(device) preds model(images) pred_mask (preds 0.5).float() iou, dice calculate_metrics(pred_mask, masks) epoch_iou iou epoch_iou / len(val_dl) print(fEpoch {epoch1} Loss: {avg_loss:.4f} Val IoU: {epoch_iou:.4f}) if epoch_iou best_iou: best_iou epoch_iou torch.save(model.state_dict(), save_path)主程序里我没有加学习率调度器实际使用中如果验证集IoU连续多个epoch不提升可以手动把学习率降到原来的0.5倍整个训练过程五六次调整就差不多了。模型保存用的state_dict文件比较小加载时也方便。5. 一个完整的隐藏流程把检测结果导出成缺陷报表很多人做完分割就停了但在实际项目里往往还需要统计每个缺陷的面积、位置甚至生成一份可视化报表。简单分享一个我一直在用的后处理流程虽然和训练本身无关但它是从算法到落地的最后一公里。拿到模型输出的二值mask之后先用cv2.findContours找到所有缺陷轮廓然后对每个轮廓计算面积和最小外接矩形。轨道巡检场景中比较关心的是单个裂纹的长度或面积所以我会额外计算轮廓的外接矩形长宽比长宽比明显大于3的基本就是裂纹。这个统计结果可以直接输出为CSV方便工务段生成台账。def defect_report(mask_path, csv_path): mask cv2.imread(mask_path, 0) _, binary cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY) contours, _ cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) results [] for cnt in contours: area cv2.contourArea(cnt) if area 25: continue x, y, w, h cv2.boundingRect(cnt) ratio max(w, h) / (min(w, h) 1e-6) defect_type crack if ratio 3 else spalling results.append({area: area, width: w, height: h, type: defect_type}) with open(csv_path, w) as f: f.write(defect_type,area,width,height\n) for r in results: f.write(f{r[defect_type]},{r[area]},{r[width]},{r[height]}\n)这只是一个简单的示例真实项目里可能还要结合相机标定把像素面积转换为实际物理面积平方毫米这个就要根据轨面宽度和图像分辨率来算了。但整体思路就是这样分割得到mask后处理提取目标统计参数输出报表。6. 最后聊几点训练效率与复现心得项目做到后面更多时间其实花在了调参和修复诡异问题上。这里按我个人的体会提三个建议。第一每次训练前固定随机种子。PyTorch里不同轮之间虽然希望初始化不一样但如果完全一样的实验配置跑两次结果差异很大说明环境随机性太强不利于排错。所以在程序开头固定Python、NumPy、PyTorch、CUDA的随机种子至少在同一个环境下能稳定复现结果。import random def seed_everything(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic True第二不要盲目增加网络深度。轨道缺陷图像分辨率不会特别大缺陷面积一般不大四层UNet已经足够。再增加层数只会让模型更笨重训练更慢却不一定带来精度提升。如果一定要尝试更大的容量优先加宽卷积核的通道数而不是增加下采样层数。第三把脚本写成一个可配置的工程。训练参数用命令行参数或yaml管理不要每次改一个值都要在代码里找。举个实际例子有人把batch_size写在训练函数内部调参的时候要改八九处最后漏改一处跑出来的结果和预期完全不同。工程上的规范决定了算法迭代的速度这一点在多人协作时尤其明显。轨道缺陷检测这条路从数据集整理到UNet落地中间没有太多玄学更多是对细节的理解和调试经验的积累。上面写的所有代码均基于实际项目中验证过的方式整理你可以直接拿去跑有问题欢迎在评论区留言。最后再分享一个小技巧训练完成后一定要在完全没参与训练的几张图上做一次“随机巡检”看看模型的泛化能力而不是只盯着验证集指标。有时候指标不错实际拍摄环境下就是会翻车提前发现总比现场部署时再被用户打回来强。 
长河君 · 2026/9/17 0:10:55 · 阅读 1.2k 轨道缺陷检测实战：基于UNet的图像分割与RSDDs数据集训练全流程
轨道缺陷检测这个方向最近两年在工业视觉里挺火但真正能把模型跑通、跑稳、还能落地的人并不多。我接触这个需求最早是帮朋友处理钢轨高速巡检图像客户给了一批包含裂纹、剥落、光带异常的路面照片要求用算法把缺陷区域逐像素标记出来。当时第一个想到的就是UNet原因是这个任务本质上是图像分割而UNet在医学影像分割里已经验证得很充分再加上轨道缺陷样本通常很小几十张图就能微调出不错的效果。今天这篇东西就是把完整的一套流程摊开来讲从数据集怎么搞、模型怎么写、训练怎么收敛到最后怎么评估效果全部按实际操作的顺序来附上可直接用的代码适合正在做检测、分割方向的学生和初级算法工程师。RSDDs数据集是学术界公开的轨道表面缺陷数据集网上能搜到不少分享包。整份数据包含快速道和普通道两个大类每类下面又细分了裂纹、剥落、光带异常等多种缺陷。整个分割任务要做的事情非常简单直白给一张轨道表面灰度图输出一张和原图等尺寸的mask图缺陷区域标为白色背景标为黑色。就这么个看起来朴素的任务实际做起来坑非常多比如裂纹这种目标细长而且窄占整张图的像素可能不到1%模型一不小心就把所有像素都预测成背景因为整体准确率照样99%。这也是我为什么不建议直接套通用分割代码而是要专门调整损失函数和训练策略的原因。做这个项目前我把网上能搜到的相关笔记翻了一遍发现很多文章只停留在“下载了数据集、训练了UNet、发了两张预测图”的阶段缺少完整可复现的代码和调参细节。所以说这篇更像一份工作记录把我实际跑通的流程和踩过的坑一并写出来。阅读的时候建议找个能跑深度学习的机器跟着代码走一遍比干看文字印象深得多。1. 为什么用UNet做轨道缺陷检测——从任务本质说起1.1 轨道缺陷检测到底难在哪轨道表面的缺陷和图里其他内容的区别很多人第一反应是“缺陷不就是表面上的裂纹破损吗用模板匹配或者边缘检测不就行了”。这个想法在理想情况下成立但真实工况下非常脆弱。首先是光照户外拍回来的图钢轨表面反光特别严重裂纹在强反光下对比度极低其次是背景干扰轨面的划痕、轮轨接触产生的黑色碾光带、枕木和道砟的纹理这些都会在特征上和真实缺陷混淆再次是缺陷形态变化大裂纹可能只有几个像素宽剥落区则呈现不规则块状边缘模糊程度完全不同。传统的图像处理方式比如Canny边缘检测加形态学操作对付固定环境、单一类型的纹理异常还行一旦换成不同区段、不同磨损状态的钢轨图像阈值根本没法统一调整。某次我在一张低对比度裂纹图上调好了Canny阈值换到另一张强反光的图直接检测出一大片假边缘效果没法用。这也是我转向深度学习分割的直接原因——目标检测框是矩形框得住剥落却框不住细长裂纹而分割能够逐像素判断类别天然适合这种细长形目标。1.2 选UNet而不选其他分割网络的核心理由图像分割领域的网络一抓一大把比如SegNet、DeepLabV3、PSPNet各有所长。但在轨道缺陷这个场景下我更推荐UNet原因有三个。第一样本量小。轨道缺陷数据集的公开部分规模很小RSDDs加起来大概就几百张图如果直接上DeepLab这种带空洞卷积的大网络或者更重的Transformer结构参数量大小数据很难训好泛化也容易出问题。UNet的参数量相对可控即使不加载预训练权重从零开始训也能在几十上百张图上收敛到可用的效果。第二细节边界好。UNet的编码器逐层提取高层语义特征同时保留详细空间信息的跳跃连接把低层特征直接拼到解码器对应层。对于轨道表面的细小裂纹低层纹理信息是救命稻草有了跳跃连接网络在解码阶段能同时兼顾整体位置和局部细节分割出来的缺陷边缘更连续。这一点对比过实际效果不加跳跃连接的版本裂纹中间经常断成几截加了跳跃连接后完整性明显改善。第三结构简单改起来方便。UNet的U型架构足够直白训练、部署、调试都不费劲。我在做项目时需要把中间特征层拿出来分析UNet可以很轻松地在解码路径上加分支输出换成其他复杂网络就没这么顺手。团队里两三天就能上手后续做缺陷面积统计、缺陷类别再分直接在UNet的mask基础上扩展就行。2. RSDDs数据集来了——先搞懂数据长什么样2.1 RSDDs数据集的几组关键数字RSDDs全称是Rail Surface Defect Dataset最早由国内某高校实验室公开是轨道缺陷检测领域用得比较多的基准数据之一。整个数据集分成两个子集RSDDs Type-I对应快速道Fast railRSDDs Type-II对应普通道Normal rail。快轨数据主要包含高速铁路区段的钢轨图片普通道数据包含普通列车轨道的图片。从缺陷类别上看两张子集各有侧重但总体上涵盖了这几类常见缺陷轨面裂纹包括横裂纹、纵裂纹、剥落掉块、轨面擦伤、光带异常等。不过需要说明的是RSDDs的官方标注并非像有些语义分割数据集那样做像素级精细标注它提供的是缺陷位置的矩形框标注也就是目标检测的标注格式。所以如果我们想直接拿来跑UNet这样的分割模型就得把矩形框转成像素级mask做法是遍历每张图对应的标注框把落在框内的像素标记为前景。这种方式不算像素级精标准确但胜在自动化处理成本低实践中完全够用。从数量上说Type-I大约有67张带缺陷的图Type-II大约有128张带缺陷的图加上一些无缺陷图像总体规模在两三百张这个量级。训练分割模型时如果直接整张图丢进去训练样本数会非常少所以常规做法是切片Patch。把一张512×512的图切分成若干128×128或256×256的小块从中筛选出包含缺陷的小块用于训练这样样本量一下子就扩充了十几倍。具体的参数和目录整理方式下面一节直接给操作。2.2 下载与目录整理RSDDs数据集因为发布有一段时间了有些原始官网链接可能失效但现在很多研究者在GitHub仓库里做了镜像整理。你可以直接搜索“RSDDs dataset download”找到包含图像和标注文件的仓库或学术资源页下载下来。如果对方提供了百度网盘或Google Drive的分享链接按说明提取即可。无论如何下载完之后我们需要把文件整理成统一的结构便于后续脚本读取。我的习惯是根目录建一个data文件夹内部按如下方式组织data/ ├── Type_I/ │ ├── IMG/ │ │ ├── 01.bmp │ │ ├── 02.bmp │ │ └── ... │ └── GT/ │ ├── 01.bmp │ ├── 02.bmp │ └── ... ├── Type_II/ │ ├── IMG/ │ │ ├── 01.bmp │ │ ├── 02.bmp │ │ └── ... │ └── GT/ │ ├── 01.bmp │ ├── 02.bmp │ └── ... └── split/ ├── train.txt └── val.txt这里我习惯把图像和标注分开放在IMG和GT两个目录文件名保持一一对应。GT可以是标注人员画好的mask图也可以是程序转出来的mask图。如果下载到的GT是xml或txt格式的矩形框标注可以参考下面的Python转换脚本通过PIL和numpy生成对应的maskimport os import cv2 import numpy as np from PIL import Image, ImageDraw def rects_to_mask(img_path, gt_path, mask_save_path): # 以Original的图像尺寸为基准 img cv2.imread(img_path) h, w img.shape[:2] mask Image.new(L, (w, h), 0) draw ImageDraw.Draw(mask) # 伪代码实际需要解析txt/xml的每个矩形框然后按坐标画矩形 # for box in boxes: # draw.rectangle([x1, y1, x2, y2], fill255) mask.save(mask_save_path)需要注意的是不同来源提供的标注文件单位可能不同有的给的是归一化坐标有的是绝对像素坐标有的是左上角和右下角有的是中心点加宽高。读取之前务必手工打开一张看看避免后期全部mask错位。2.3 训练数据增强把有限的样本“变”多两三百张图直接训练模型很容易过拟合在训练集上表现很好换一批真实图片就崩。我的做法是离线做数据增强也就是在训练前先把图像和mask一起做变换生成一批增强后的样本保存到新目录而不是在训练循环里在线增强。这样做的优点是方便检查增强效果也能避免训练过程中CPU和GPU速度不匹配影响训练进度。常用的增强方式包括随机水平翻转轨道图像对称性还算明显翻转不会破坏语义、随机旋转90度、随机亮度对比度调整、添加高斯噪声、随机裁剪以及整体缩放。但轨道缺陷检测里要特别注意裂纹这类细长形目标经不起任意角度旋转和拉伸。比如随机旋转30度裂纹的几何形态和真实情况严重不符非等比例缩放会让裂纹变胖或变细破坏物理特征。所以我实操时只用90度的整数倍旋转、水平翻转、小范围缩放和亮度抖动。具体的增强代码不依赖外部库也能写比如用OpenCV的warpAffine配合旋转矩阵或者在PIL里直接使用transpose。但如果追求效率可以直接用albumentations库它对图像和mask同步增强做得很好。我常用的增强管道像下面这样import albumentations as A train_transform A.Compose([ A.HorizontalFlip(p0.5), A.RandomRotate90(p0.5), A.RandomBrightnessContrast(brightness_limit0.2, contrast_limit0.2, p0.5), A.GaussNoise(var_limit(10.0, 50.0), p0.3), A.Resize(height256, width256, always_applyTrue) ])在增强时需要保证图像和mask使用完全相同的随机参数这样才能保证位置对齐。albumentations这类库已经封装好这个同步逻辑直接调用即可。如果自己写代码千万记住在随机数生成上保持同一个状态。3. 手把手搭建UNet训练环境与代码3.1 训练环境与硬件配置开始写代码之前先把环境准备好省得到时候装包装到怀疑人生。我的训练环境配置如下仅供参考操作系统Ubuntu 20.04 / Windows 11 都可以Python版本3.8 或 3.9深度学习框架PyTorch 1.10下文代码基于PyTorch 2.0写CUDA11.x或12.x均可取决于你的显卡驱动显卡建议至少4GB显存我这边用了一张16GB的卡训练不存在显存压力安装依赖可以直接用pippip install torch torchvision opencv-python pillow numpy albumentations tqdm如果你用的是GPU版PyTorch记得去官网选择匹配你本机CUDA版本的那条安装命令不要直接pip install torch装上CPU版本否则训练速度会慢得让人崩溃。CPU版也不是不能跑只是跑一个epoch的时间足够你打两把游戏再回来。3.2 UNet模型结构实现PyTorch完整代码UNet的核心结构是编码器-解码器结构编码器由4到5个block组成每经过一个block空间尺寸减半通道数翻倍解码器则进行上采样空间尺寸加倍通道数减半。在编码器和解码器之间有跳跃连接把对应尺度的特征拼接起来。以下是简洁可用的UNet实现输入输出都是单通道适合二值分割。import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): 两次卷积 BN ReLU def __init__(self, in_ch, out_ch): super(DoubleConv, self).__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, in_channels1, out_channels1, features[64, 128, 256, 512]): super(UNet, self).__init__() self.encoder nn.ModuleList() self.pool nn.MaxPool2d(kernel_size2, stride2) self.decoder nn.ModuleList() # 编码器 for f in features: self.encoder.append(DoubleConv(in_channels, f)) in_channels f # 瓶颈层 self.bottleneck DoubleConv(features[-1], features[-1] * 2) # 解码器反转特征通道缩回去 for idx in range(len(features) - 1, -1, -1): self.decoder.append( nn.ConvTranspose2d(features[idx] * 2, features[idx], kernel_size2, stride2) ) self.decoder.append(DoubleConv(features[idx] * 2, features[idx])) self.final_conv nn.Conv2d(features[0], out_channels, kernel_size1) def forward(self, x): skip_connections [] for enc in self.encoder: x enc(x) skip_connections.append(x) x self.pool(x) x self.bottleneck(x) skip_connections skip_connections[::-1] for idx in range(0, len(self.decoder), 2): x self.decoder[idx](x) skip skip_connections[idx // 2] if x.shape ! skip.shape: x F.interpolate(x, sizeskip.shape[2:], modebilinear, align_cornersTrue) x torch.cat((skip, x), dim1) x self.decoder[idx 1](x) return torch.sigmoid(self.final_conv(x))这段代码里编码器层层提取特征每个DoubleConv内部做的都是卷积、批归一化、ReLU的标准组合。批归一化非常重要它能加快收敛也能减少对初始权重和输入分布的敏感度。解码器的转置卷积负责把特征图逐步放大每次放大后和对应的编码器特征进行通道拼接再经过DoubleConv融合。需要提醒的是原始UNet用两个3×3卷积叠加相当于一个5×5感受野同时减少了参数量。在轨道缺陷分割里缺陷目标较小所以没必要加深太多层四层特征已经足够。如果你希望进一步提升精度可以把第一层的初始通道数设为32或64上面代码默认64。3.3 数据读取与训练循环数据读取阶段需要同时读入原图和mask并且把它们变成相同的尺寸。训练时用小patch256×256左右比较均衡。下面是数据集的PyTorch封装import os import cv2 import torch from torch.utils.data import Dataset class RailDataset(Dataset): def __init__(self, image_dir, mask_dir, image_size256, transformNone): self.image_dir image_dir self.mask_dir mask_dir self.image_size image_size self.transform transform self.images sorted(os.listdir(image_dir)) def __len__(self): return len(self.images) def __getitem__(self, idx): img_name self.images[idx] img_path os.path.join(self.image_dir, img_name) mask_path os.path.join(self.mask_dir, img_name) image cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) # 统一尺寸 image cv2.resize(image, (self.image_size, self.image_size), interpolationcv2.INTER_LINEAR) mask cv2.resize(mask, (self.image_size, self.image_size), interpolationcv2.INTER_NEAREST) # mask只取0/1 mask cv2.threshold(mask, 127, 1, cv2.THRESH_BINARY)[1] if self.transform is not None: transformed self.transform(imageimage, maskmask) image transformed[image] mask transformed[mask] image torch.from_numpy(image).float().unsqueeze(0) / 255.0 mask torch.from_numpy(mask).float().unsqueeze(0) return image, mask这里有一个很容易忽略的细节resize mask时插值方式必须用INTER_NEAREST也就是最近邻插值。如果用线性插值去缩小mask原先是255的像素可能会变成100多甚至0造成边界重影那么loss计算时会产生大量错误梯度导致模型学出模糊的边缘。训练循环部分我更倾向于把BCE和Dice结合作为总损失。BCE负责像素级的收敛Dice损失则直接优化分割区域的重叠程度对小目标更友好。下面这个交叉Dice的实现正是为了应对缺陷小、正负样本不平衡问题def dice_loss(pred, target, smooth1.0): pred pred.contiguous().view(-1) target target.contiguous().view(-1) intersection (pred * target).sum() return 1 - (2.0 * intersection smooth) / (pred.sum() target.sum() smooth) def combined_loss(pred, target): bce F.binary_cross_entropy(pred, target) dice dice_loss(pred, target) return bce dice训练循环就常规操作遍历数据集前向传播算损失反向传播更新权重。以我自己的配置批量大小设为16学习率0.001Adam优化器训练60个epoch整个过程大约需要十到二十分钟。如果你用的是CPU时间会成倍增加建议有条件的还是用GPU。3.4 模型评估与结果可视化训练结束后不能只看训练集loss还需要在验证集上量化指标。语义分割常用的指标包括Dice系数F1、IoU、Pixel Accuracy。由于轨道缺陷区域很小Pixel Accuracy往往虚高没有参考价值。我主要看IoU和Dice这两个指标对预测区域和真实区域的重合程度更敏感。评估代码可以用下面这段def calculate_metrics(pred_mask, true_mask): pred_mask (pred_mask 0.5).int() true_mask (true_mask 0.5).int() intersection (pred_mask * true_mask).sum().item() union (pred_mask | true_mask).sum().item() iou intersection / (union 1e-6) dice (2 * intersection) / (pred_mask.sum().item() true_mask.sum().item() 1e-6) return iou, dice预测时把验证集的每张图和对应mask读入经过模型推理得到概率图阈值0.5转成二值mask。可视化时推荐将原图、标注、预测三张图横向拼在一起这样一眼就能看出哪些地方漏检了、哪些地方是误检。建议保存成图片到result/目录方便后续做报告或者排查。4. 踩坑实录与效果优化4.1 Loss不降或震荡怎么办训练时最糟心的就是loss纹丝不动。我遇到这个情况先从三个方向排查第一检查数据预处理看看mask是否真的和图像对齐很多标注文件存在裁切偏移肉眼看不到但loss会一直很高第二学习率是不是太大或太小Adam默认的0.001已经比较激进配合小数据集时甚至可能震荡我后来降到0.0005就稳定了第三正负样本极度不平衡如果一张图缺陷只占0.1%像素BCE loss会被大量背景像素淹没所以必须配合Dice Loss或者对正样本提权。遇到过最诡异的一次是loss在第20个epoch突然从0.3跳到0.9再也没降回来。排查发现是BN层的均值和方差在batch噪声大的时候波动严重解决办法是增大batch size或者把BN改成GroupNorm。因为GPU显存有限我选择用了GroupNorm版本的UNet明显稳定不少。4.2 裂纹预测断成一截一截的裂纹这种长条形目标有个普遍问题模型预测出的结果经常中间有断裂。分析原因一个是原图上裂纹本身对比度就低中间某段跟背景融为一体网络难以判断另一个原因是下采样太多次小裂纹细节在深层特征里已经丢掉了仅靠高级语义信息补不回来。解决办法我实践下来有三招。第一用切片训练而不是整图训练让模型看到更多细节同时避免resize把裂纹挤没第二把UNet的第一层通道数从64改成32减少下采样信息压缩但这会轻微增加计算量第三后处理环节使用形态学闭运算把距离很近的分段裂纹连接起来比如用OpenCV的cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)我用的kernel是5×5的椭圆核连接效果不错但也容易把间隙很小的正常纹理连成缺陷需要根据实际图片调kernel大小。4.3 从60分到90分我用过的3个提高精度手段如果基础UNet跑通了精度却还差一口气我强烈建议按下面顺序迭代。第一损失函数优化。单独BCE很容易让网络偏向于学习背景单独Dice Loss则在训练初期过于不平滑。我用的组合是BCE Dice在一些小裂纹数据集上比单独BCE能提升3到5个点的IoU。如果你的缺陷特别小也可以尝试Focal Loss DiceFocal Loss能抑制易分类的大背景样本让模型更关注难的缺陷区域效果如何也要结合数据实测。第二预训练编码器。因为RSDDs样本量不够大从头训卷积核容易欠拟合。如果不想引入外部数据可以先用ImageNet预训练的ResNet作为UNet编码器只随机初始化解码器。这个做法在很多分割任务里是常规操作能够显著提升收敛速度和最终精度。在PyTorch中加载方式很简单把UNet的encoder部分替换为torchvision.models.resnet34(pretrainedTrue)并去掉最后的全连接层就好。第三后处理加连通域筛选。模型输出的二值mask里通常会有一些零散的小块误检这些往往是背景中的亮点或锈斑。我们可以借助连通域分析删除面积小于某个阈值的区域。例如设定最小缺陷面积为25像素如果某个连通域面积小于这个阈值就把它置为0。这个后处理十分粗暴有效能够很好地抑制零星噪声但要注意阈值不能设太大否则真正的小裂纹也会被误删。我的做法是先统计验证集标注中所有连通域的面积分布取最小的1%作为阈值下限。4.4 顺手放一套完整的训练主程序从数据加载到推理代码分散在各个小节里可能看起来不够连贯。下面我整合了一段主训练脚本涵盖数据路径配置、训练循环、模型保存和推理可视化你可以直接复制过去改动路径使用。import os import torch import torch.optim as optim from torch.utils.data import DataLoader from tqdm import tqdm import cv2 import numpy as np # 假设你已经把UNet类、RailDataset类、combined_loss函数定义在上方 device torch.device(cuda if torch.cuda.is_available() else cpu) # 配置 train_img_dir data/train/images train_mask_dir data/train/masks val_img_dir data/val/images val_mask_dir data/val/masks batch_size 16 lr 0.0005 epochs 60 save_path checkpoints/unet_rail.pth # 数据与加载 train_ds RailDataset(train_img_dir, train_mask_dir, image_size256) val_ds RailDataset(val_img_dir, val_mask_dir, image_size256) train_dl DataLoader(train_ds, batch_sizebatch_size, shuffleTrue, num_workers4) val_dl DataLoader(val_ds, batch_sizebatch_size, shuffleFalse, num_workers4) # 模型与优化器 model UNet(in_channels1, out_channels1, features[64, 128, 256, 512]).to(device) optimizer optim.Adam(model.parameters(), lrlr) best_iou 0.0 for epoch in range(epochs): model.train() total_loss 0 for images, masks in tqdm(train_dl, descfEpoch {epoch1}/{epochs}): images, masks images.to(device), masks.to(device) preds model(images) loss combined_loss(preds, masks) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() avg_loss total_loss / len(train_dl) # 验证 model.eval() epoch_iou 0 with torch.no_grad(): for images, masks in val_dl: images, masks images.to(device), masks.to(device) preds model(images) pred_mask (preds 0.5).float() iou, dice calculate_metrics(pred_mask, masks) epoch_iou iou epoch_iou / len(val_dl) print(fEpoch {epoch1} Loss: {avg_loss:.4f} Val IoU: {epoch_iou:.4f}) if epoch_iou best_iou: best_iou epoch_iou torch.save(model.state_dict(), save_path)主程序里我没有加学习率调度器实际使用中如果验证集IoU连续多个epoch不提升可以手动把学习率降到原来的0.5倍整个训练过程五六次调整就差不多了。模型保存用的state_dict文件比较小加载时也方便。5. 一个完整的隐藏流程把检测结果导出成缺陷报表很多人做完分割就停了但在实际项目里往往还需要统计每个缺陷的面积、位置甚至生成一份可视化报表。简单分享一个我一直在用的后处理流程虽然和训练本身无关但它是从算法到落地的最后一公里。拿到模型输出的二值mask之后先用cv2.findContours找到所有缺陷轮廓然后对每个轮廓计算面积和最小外接矩形。轨道巡检场景中比较关心的是单个裂纹的长度或面积所以我会额外计算轮廓的外接矩形长宽比长宽比明显大于3的基本就是裂纹。这个统计结果可以直接输出为CSV方便工务段生成台账。def defect_report(mask_path, csv_path): mask cv2.imread(mask_path, 0) _, binary cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY) contours, _ cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) results [] for cnt in contours: area cv2.contourArea(cnt) if area 25: continue x, y, w, h cv2.boundingRect(cnt) ratio max(w, h) / (min(w, h) 1e-6) defect_type crack if ratio 3 else spalling results.append({area: area, width: w, height: h, type: defect_type}) with open(csv_path, w) as f: f.write(defect_type,area,width,height\n) for r in results: f.write(f{r[defect_type]},{r[area]},{r[width]},{r[height]}\n)这只是一个简单的示例真实项目里可能还要结合相机标定把像素面积转换为实际物理面积平方毫米这个就要根据轨面宽度和图像分辨率来算了。但整体思路就是这样分割得到mask后处理提取目标统计参数输出报表。6. 最后聊几点训练效率与复现心得项目做到后面更多时间其实花在了调参和修复诡异问题上。这里按我个人的体会提三个建议。第一每次训练前固定随机种子。PyTorch里不同轮之间虽然希望初始化不一样但如果完全一样的实验配置跑两次结果差异很大说明环境随机性太强不利于排错。所以在程序开头固定Python、NumPy、PyTorch、CUDA的随机种子至少在同一个环境下能稳定复现结果。import random def seed_everything(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic True第二不要盲目增加网络深度。轨道缺陷图像分辨率不会特别大缺陷面积一般不大四层UNet已经足够。再增加层数只会让模型更笨重训练更慢却不一定带来精度提升。如果一定要尝试更大的容量优先加宽卷积核的通道数而不是增加下采样层数。第三把脚本写成一个可配置的工程。训练参数用命令行参数或yaml管理不要每次改一个值都要在代码里找。举个实际例子有人把batch_size写在训练函数内部调参的时候要改八九处最后漏改一处跑出来的结果和预期完全不同。工程上的规范决定了算法迭代的速度这一点在多人协作时尤其明显。轨道缺陷检测这条路从数据集整理到UNet落地中间没有太多玄学更多是对细节的理解和调试经验的积累。上面写的所有代码均基于实际项目中验证过的方式整理你可以直接拿去跑有问题欢迎在评论区留言。最后再分享一个小技巧训练完成后一定要在完全没参与训练的几张图上做一次“随机巡检”看看模型的泛化能力而不是只盯着验证集指标。有时候指标不错实际拍摄环境下就是会翻车提前发现总比现场部署时再被用户打回来强。