做轨道缺陷检测这个方向很多人一上来就想着换新模型、上大网络其实往往忽略了一个关键事实在缺陷这种小目标、细长结构、背景噪声还特别严重的场景里UNet这种老伙计反而格外能打。加上RSDDs这个公开的轨道表面缺陷数据集整个项目完全可以在一张消费级显卡上跑起来非常适合作为接轨工业视觉的入门实战。这篇内容我一次性把思路、数据、代码、踩坑记录全部捋清楚手把手带你把一个完整的轨道缺陷检测流程跑通。适合正在学习图像分割、准备做工业检测毕设或者刚入职做算法验证的同学参考你不需要一开始就追求SOTA先把一个可靠的pipeline跑起来后面再谈改进。1. 项目概述与整体方案设计1.1 轨道缺陷检测到底难在哪轨道表面缺陷检测跟一般的语义分割任务有个非常大的不同——它本质上是在极高分辨率的大图上找极小面积的异常区域。钢轨表面的裂纹可能只有几个像素宽剥落区域的形状又不规则再加上光照反射、油污、锈迹这些干扰传统视觉的阈值分割和边缘检测根本扛不住。RSDDs数据集里的缺陷形态也验证了这个痛点。Type-I子集的缺陷以裂纹和剥落为主Type-II子集则在复杂背景下包含更多快速形成的损伤两类缺陷的尺寸、长宽比、灰度分布在同一个数据集里差异非常大。如果直接用分类网络去判断有没有缺陷只能输出一个全局结论根本定位不了缺陷的像素位置这在工业场景里是没有落地价值的。你需要的是像素级的分割结果。1.2 为什么选择UNet作为分割主模型UNet最早是为医学图像分割设计的但它的结构特性跟轨道缺陷检测的需求高度匹配。第一是编码器-解码器的对称结构。编码器逐层下采样抓特征解码器逐层恢复空间分辨率这个压缩-还原的漏斗结构天然适合全局理解局部定位的任务组合。轨道图像里的缺陷虽然小但它依赖周围的灰度分布和纹理上下文才能判断纯局部看很容易漏检。第二是跳跃连接。这是UNet的灵魂。下采样会丢失边缘细节而跳跃连接把编码器各层的特征直接拼接到解码器对应层相当于给小目标检测开了一条信息高速公路。分割任务里边缘的精确程度直接影响质量评估这一点上没有多少模型能比UNet更直接。第三是训练成本可控。不是说非得选UNet而是对一个刚上手的项目来说UNet参数量适中两三百张图的RSDDs数据集不会把它喂撑也不会欠拟合。同样是分割模型DeepLabV3需要ASPP模块的精细调参SegFormer又要额外处理Transformer的收敛问题两张图测下来UNet的性价比已经很高了。先把UNet跑透再去试改进版UNet和TransUNet路径清晰得多。1.3 整体技术路线设计这套方案的技术流程可以拆成四段数据准备、模型搭建、训练调优、推理可视化。数据部分负责把RSDDs原始图像和掩码处理成模型能吃的标准格式并对有限的数据做增强以降低过拟合模型部分直接实现UNet的编码器、解码器和跳跃连接训练部分重点设计损失函数单一的二值交叉熵在缺陷像素占比极低的时候会失效必须配合Dice Loss使用推理部分则要做阈值分割和后处理把模型的概率图变成可以指导检修的缺陷区域。2. RSDDs数据集获取与预处理实战2.1 RSDDs数据集构成与获取方式RSDDs是暨南大学团队公开的轨道表面缺陷数据集分为Type-I和Type-II两个子集。Type-I包含67张轨道表面灰度图缺陷以典型裂纹、擦伤为主Type-II包含128张灰度图缺陷类型更复杂包括快速形成的压痕和波浪状损伤。两张子集均为像素级标注缺陷区域为白、背景为黑这决定了模型训练时用二类分割即可。数据集本身不通过外部网盘发布建议优先在论文作者机构公开页面或数据集官方GitHub仓库获取。搜索RSDDs dataset时注意核对路径下是否同时包含IMAGES和MASKS文件夹并确认掩码是PNG或BMP整图标注而不是多边形标注。重要提示拿到数据先看一眼掩码的灰度值。RSDDs官方掩码缺陷区域虽然是白色但有些渠道转发时经过格式转换白色值会被压到200出头甚至更低的灰度。写代码时不要写死等于255统一用大于127作为前景判断才是稳妥的。2.2 读取、裁剪与归一化RSDDs原图分辨率不低如果直接整图缩放成256x256训练缺陷会缩成几个像素甚至消失。我实际验证下来合理的路线是先把所有图像统一缩放到512x512作为网络的输入尺寸虽然依然有信息损失但在单卡显存限制下这是精度和速度的平衡点。图像以灰度模式读入归一化到0到1区间相当于单通道输入。掩码必须以最近邻插值缩放不能用双线性。双线性会把掩码边缘插出介于0和255之间的灰色过渡带而0到1二分类根本不该有中间值。这个坑非常隐蔽稍后我还会专门强调。2.3 数据增强策略与小样本对策RSDDs整个数据集不足200张原始图要训练UNet这样有参数量的模型不做增强基本等于过拟合预定。我这里用的是轻量级在线增强每个epoch对图像随机进行水平翻转、垂直翻转和90度旋转的组合。旋转对轨道裂缝这类方向性强的目标是有争议的有些缺陷形态旋转后物理上不合理但实验显示适度旋转能显著提升模型对缺陷姿态的鲁棒性利大于弊。不推荐的增强手段是随机裁剪。轨道图像里大部分区域是正常轨面随机裁剪出来的patch大概率全是背景反而给模型制造大量无缺陷负样本会让分割结果偏向背景。如果确实要做裁剪用基于缺陷位置的中心裁剪会更合理。3. UNet模型搭建与完整代码实现3.1 UNet网络结构关键点拆解先给UNet的结构定个调。输入是512x512的单通道灰度图经过四个下采样阶段通道数依次为64、128、256、512瓶颈层升到1024然后四个上采样阶段把通道数依次降回来每次上采样之后与对应编码器层拼接最后用1x1卷积输出单通道特征图代表每个像素属于缺陷的logit。每个卷积块由两次3x3卷积批标准化ReLU组成。批标准化对轨道图像这种灰度分布差异大的数据非常重要能抑制光照变化带来的分布漂移。下采样部分直接使用最大池化虽然现代模型普遍用stride2卷积替代池化但在小数据集上最大池化对平移变化更不敏感泛化能力测试下来反而更好。3.2 完整UNet模型定义可直接运行模型部分的完整代码直接从这段开始import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): 两个卷积BNReLU组成的标准卷积块 def __init__(self, in_ch, out_ch): super().__init__() self.block nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.block(x) class DownBlock(nn.Module): 下采样块2x2最大池化 双卷积 def __init__(self, in_ch, out_ch): super().__init__() self.block nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_ch, out_ch) ) def forward(self, x): return self.block(x) class UpBlock(nn.Module): 上采样块双线性插值上采样 跳跃连接拼接 双卷积 def __init__(self, in_ch, out_ch): super().__init__() self.up nn.Upsample(scale_factor2, modebilinear, align_cornersTrue) self.conv DoubleConv(in_ch, out_ch) def forward(self, x_enc, x_dec): x_dec self.up(x_dec) # 处理编码器和解码器特征图尺寸不一致的问题 diff_y x_enc.size()[2] - x_dec.size()[2] diff_x x_enc.size()[3] - x_dec.size()[3] x_dec F.pad( x_dec, [diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2] ) x torch.cat([x_enc, x_dec], dim1) return self.conv(x) class UNet(nn.Module): 标准UNet分割网络 def __init__(self, in_channels1, num_classes1): super().__init__() self.base DoubleConv(in_channels, 64) self.down1 DownBlock(64, 128) self.down2 DownBlock(128, 256) self.down3 DownBlock(256, 512) self.down4 DownBlock(512, 1024) self.up1 UpBlock(1024 512, 512) self.up2 UpBlock(512 256, 256) self.up3 UpBlock(256 128, 128) self.up4 UpBlock(128 64, 64) self.out_conv nn.Conv2d(64, num_classes, 1) def forward(self, x): x1 self.base(x) x2 self.down1(x1) x3 self.down2(x2) x4 self.down3(x3) x5 self.down4(x4) x self.up1(x4, x5) x self.up2(x3, x) x self.up3(x2, x) x self.up4(x1, x) return self.out_conv(x)这段代码没有用任何花哨的trick就是标准的UNet实现也是后面所有实验的基线。需要注意的是上采样部分用了双线性插值而不是转置卷积。转置卷积需要额外的可学习参数在小数据集上训练不当容易产生棋盘格伪影双线性插值则完全避免了这个问题。3.3 数据加载器与数据增强实现模型有了接下来是数据端。这里需要重点处理两件事掩码的二值化逻辑和增强的同步性。import os import random import numpy as np from PIL import Image from torch.utils.data import Dataset class RailDefectDataset(Dataset): 轨道缺陷数据集输入灰度图与二值掩码训练时可选用增强 def __init__(self, image_dir, mask_dir, size(512, 512), augTrue): self.image_paths sorted([ os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.lower().endswith((.png, .jpg, .jpeg, .bmp)) ]) self.mask_paths sorted([ os.path.join(mask_dir, f) for f in os.listdir(mask_dir) if f.lower().endswith((.png, .jpg, .jpeg, .bmp)) ]) assert len(self.image_paths) len(self.mask_paths), 图像和掩码数量不一致 self.size size self.aug aug def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img Image.open(self.image_paths[idx]).convert(L) mask Image.open(self.mask_paths[idx]).convert(L) # 图像用双线性缩放掩码用最近邻缩放 img img.resize(self.size, Image.BILINEAR) mask mask.resize(self.size, Image.NEAREST) img_arr np.array(img, dtypenp.float32) / 255.0 mask_arr (np.array(mask, dtypenp.float32) 127.0).astype(np.float32) if self.aug: if random.random() 0.5: img_arr img_arr[:, ::-1] mask_arr mask_arr[:, ::-1] if random.random() 0.5: img_arr img_arr[::-1, :] mask_arr mask_arr[::-1, :] k random.choice([0, 1, 2, 3]) img_arr np.rot90(img_arr, k) mask_arr np.rot90(mask_arr, k) # 转换为张量图像为1xCxHxW掩码为HxW img_tensor torch.from_numpy(img_arr.copy()).unsqueeze(0) mask_tensor torch.from_numpy(mask_arr.copy()) return img_tensor, mask_tensor这里我特别强调一下resize时双线性和最近邻的问题。图像是连续信号用双线性无损缩放没问题但掩码是离散标签一用双线性就会在边界产生灰色像素。这些灰色像素进损失函数时会被默认当作介于0和1之间的中间值整个模型就会去学习预测一个根本不存在的类别最终导致边界附近predictions糊成一片。3.4 损失函数设计与训练主循环轨道缺陷检测里最典型的困境是一个batch中缺陷像素可能只占全部像素的1%到3%。如果直接用二值交叉熵损失网络很快就会发现全部预测为背景可以让loss非常低这不是我们想要的结果。解决方案是组合损失。import torch import torch.nn as nn import torch.nn.functional as F class DiceLoss(nn.Module): Dice Loss缓解前景背景不平衡 def __init__(self, smooth1e-6): super().__init__() self.smooth smooth def forward(self, logits, targets): preds torch.sigmoid(logits) preds preds.contiguous().view(-1) targets targets.contiguous().view(-1) intersection (preds * targets).sum() dice (2.0 * intersection self.smooth) / ( preds.sum() targets.sum() self.smooth ) return 1.0 - dice def combined_loss(logits, targets): BCE Dice 组合损失 bce F.binary_cross_entropy_with_logits(logits, targets) dice DiceLoss()(logits, targets) return 0.5 * bce 0.5 * diceBCE负责逐像素的分类准确性Dice直接优化的是重叠区域两个指标的侧重点不同组合在一起能让模型既学会判类别又学会画边界。权重比例可以动态调整缺陷特别小的时候把Dice权重往上抬我实测0.5比0.5在RSDDs上表现最均衡。训练主循环这段是标准但需要保持清晰的from torch.utils.data import DataLoader from torch.optim import Adam from torch.optim.lr_scheduler import CosineAnnealingLR def train_model(model, train_loader, val_loader, epochs60, lr1e-4, devicecuda, save_pathunet_rail.pth): optimizer Adam(model.parameters(), lrlr, weight_decay1e-5) scheduler CosineAnnealingLR(optimizer, T_maxepochs) best_dice 0.0 for epoch in range(epochs): # 训练阶段 model.train() train_loss 0.0 for imgs, masks in train_loader: imgs imgs.to(device) masks masks.to(device).unsqueeze(1) logits model(imgs) loss combined_loss(logits, masks) optimizer.zero_grad() loss.backward() optimizer.step() train_loss loss.item() # 验证阶段每5轮跑一次指标 if (epoch 1) % 5 0: val_dice evaluate_model(model, val_loader, device) if val_dice best_dice: best_dice val_dice torch.save(model.state_dict(), save_path) print(fEpoch {epoch1}/{epochs} | Loss: {train_loss/len(train_loader):.4f} | Val Dice: {val_dice:.4f}) scheduler.step()训练过程中损失值稳步下降每5个epoch验证一次Dice系数只保留验证集效果最好的权重。这也是实操里最实用的存档策略避免训练最后阶段过拟合导致保存了更差的模型。3.5 评估指标与可视化函数分割任务只看loss没有说服力还需要一套标准评估指标。这里给出MIoU、Dice、PA像素精度三件套的实现def compute_metrics(logits, masks, threshold0.5): preds (torch.sigmoid(logits) threshold).float() inter (preds * masks).sum() union (preds masks).clamp(0, 1).sum() miou (inter 1e-6) / (union 1e-6) dice (2 * inter 1e-6) / (preds.sum() masks.sum() 1e-6) pa (preds masks).float().mean() return miou.item(), dice.item(), pa.item() def evaluate_model(model, val_loader, device): model.eval() dices [] with torch.no_grad(): for imgs, masks in val_loader: imgs imgs.to(device) masks masks.to(device).unsqueeze(1) logits model(imgs) _, dice, _ compute_metrics(logits, masks) dices.append(dice) return sum(dices) / len(dices)MIoU反映预测区域和真实区域的交并比Dice跟MIoU趋势一致但数值通常更高PA容易受背景像素影响所以只做参考。实际汇报项目进度时工业方更关心Dice和漏检率PA基本没人看。4. 训练配置与调参实战4.1 关键超参数如何确定batch_size直接受显存约束。512x512输入加上UNet本身的中间特征图一张12GB显存的卡最多也就放8张。如果你的卡只有6G老老实实把batch_size降到4或者在Dataset里把输入尺寸改成256x256不要硬撑out-of-memory会打断训练节奏。学习率我建议直接放在1e-4。Adam优化器对学习率不敏感是相对的在分割任务上1e-3起步太激进loss曲线经常抖得像心电图从1e-4起步虽然慢一点但能保证收敛稳定性。epoch数设60轮配合余弦退火调度器逐步降低学习率基本到第30轮以后就可以看到验证集Dice稳定在0.75以上。权重衰减建议开一个非常小的值1e-5就够。它的作用是抑制权重的幅度过大对小数据集能有效防止训练后期在验证集上抖。当然如果你用的是weight_decay默认的0问题也不大主要靠early stopping即保存最优模型来控制。4.2 训练过程记录与结果预期分析训练时记录loss曲线很有必要但不要只看loss。我见过太多人被训练loss骗了——loss在降验证Dice却纹丝不动甚至往下走。这种时候基本可以断定是过拟合或者数据增强不够。RSDDs上我个人经验是BCEDice组合损失在前20轮下降很快到第35轮左右进入平台期这也是模型欠拟合被大部分解决的信号。一个正常的RSDDs训练结果Type-I上的验证Dice应该能到0.78到0.86之间Type-II因为背景更复杂会略低一些0.70到0.80都算正常。如果验证Dice连0.65都到不了先检查掩码读取是不是出了问题再检查是否忘了做归一化。数据端错误的概率比模型问题大得多。4.3 推理预测与检测结果可视化训练结束后的推理代码反而是整个项目里最容易被忽视的部分。这里我给出一个完整的预测函数它把模型输出的概率图转成可视化的缺陷掩码def predict_single(model, image_path, device, size(512, 512), threshold0.5): model.eval() orig_img Image.open(image_path).convert(L) orig_w, orig_h orig_img.size img orig_img.resize(size, Image.BILINEAR) img_arr np.array(img, dtypenp.float32) / 255.0 img_tensor torch.from_numpy(img_arr).unsqueeze(0).unsqueeze(0).to(device) with torch.no_grad(): logits model(img_tensor) prob torch.sigmoid(logits).squeeze().cpu().numpy() mask (prob threshold).astype(np.uint8) * 255 mask_img Image.fromarray(mask).resize((orig_w, orig_h), Image.NEAREST) return mask_img, prob推理时有三个容易踩的点。第一模型输入做了归一化推理时也必须做同样的归一化否则模型看到的数据分布跟训练时不一致效果崩掉。第二resize回原始尺寸必须用最近邻理由跟掩码一致。第三阈值的选择需要看你的业务侧重0.5是默认值想抓全缺陷就把阈值降到0.3牺牲一点精确率追求高精度就提到0.6以上。5. 常见问题与排查技巧实录5.1 掩码全黑或全白问题这个现象太常见了十次里有八次是数据读取的锅。掩码全黑先检查路径是否正确再检查二值化阈值条件是否写倒了。掩码全白大概率是读取到的图像本身就是反色的某些下载源会把黑白反转你需要在读取时做个判断统计一下掩码平均灰度超过128就做一次像素取反。我建议把数据检查做成一个独立的小脚本在训练前打印出一对训练样本的shape、dtype、数值范围、mask的unique值。如果unique值只有0和1数据基本没问题。这一步30秒能做完能省下好几个小时的调试时间。5.2 训练loss不下降或直接NANloss不下降最先检查的不是模型结构而是数据归一化。我见过有人在数据增强里把image乘以255又除以127数值范围乱得一塌糊涂模型当然学不会。确认图像输入在0到1之间掩码只有0和1这两条对了loss不下降的问题就解决了一半。loss直接NAN九成是学习率过大炸了先降学习率再看。另一个诱因是BN层在batch_size1的时候统计量不稳定训练时方差接近0导致梯度爆炸这时候要么加大batch_size要么把BN换成GroupNorm。5.3 预测掩码全是背景怎么办模型预测全背景很多人第一反应是模型坏了实际上通常是推理流程里忘了sigmoid。UNet输出的是logits不经过sigmoid直接阈值化正样本值可能全部小于0.5或者大于0.5但被阈值挡掉结果当然是一片黑。正确的顺序一定是sigmoid拿到概率值再做阈值判断。另外还要检查训练时是否严重类别失衡。如果Dice Loss权重太低而缺陷又极小模型很容易收敛到全是背景的局部最优。把Dice权重从0.5往0.7或者0.8调通常能有效逼迫模型关注正样本。5.4 推理速度优化建议学术demo无所谓速度但落地时一台巡检车上可能挂多个摄像头每个画面都要做实时分割。UNet本身不算重但想要更快有几个简单手段。第一把输入尺寸从512降到384甚至320速度提升非常明显缺陷检测的精度损失可能在可接受范围内。第二推理时用半精度float16一块支持AMP的显卡能带来将近一倍的加速。第三量化到INT8适合部署在边缘设备上但需要校准数据集生产环境里再慢慢调。注意做速度优化时一定要用同一批测试图做精度回归对比不要只盯着帧率。有很多模型压缩到int8后Dice从0.8掉到0.6这种模型到了现场就是批量漏检事故。6. 项目实战心得总结RSDDs这个数据集虽然图少但它是检验一个分割pipeline是否正常的最快方式。在我实际做过这么多分割任务之后回头再看UNet跑轨道缺陷检测这个组合真正考验人的并不是网络结构而是数据管线中掩码插值、归一化、损失函数配比这些细节。只要把这些细节做对了哪怕只用最基础的UNet效果也足够当做一个完整项目来汇报。从我个人的经验来说下一步扩展方向有几个很自然的路径。一是把单通道灰度输入换成Gabor滤波或LBP纹理特征作为额外输入通道对裂纹类缺陷有额外增益二是尝试把UNet的编码器替换为在ImageNet上预训练的ResNet34利用预训练权重加速收敛三是引入注意力模块比如在跳跃连接处加CBAM或者用self-attention去建模缺陷的长距离依赖。每一个方向都够展开做一篇新的实战笔记但前提始终是先把当前这个UNet的baseline稳定跑出来。最后拿我踩过最重的一次坑提醒大家下载数据集之后打开掩码多看几眼确认和原图是对应的。我遇到过一次文件夹内图片排序错位导致图像和掩码完全对不上模型训练了两百个epochloss倒是下降了但实际学到的全是噪声。数据端的严谨程度直接决定你后面花的每一个小时到底有没有意义。