小样本铁路轨道故障图像识别：从CNN到YOLOv5的分类实践

简介面向铁路轨道视觉检测场景的已标注图像分类数据集围绕轨道表面损坏与未损坏的二分类任务构建适合正在开展缺陷识别、智慧运维或图像分类相关课题的算法工程师、科研人员与学生使用。数据包含约380张轨道图像标签分为损坏、未损坏两类并已抽分为训练集、验证集和测试集按目录分别存放同时提供json标签文件与Python可视化脚本便于直接查看样本、理解类别映射并快速接入CNN分类网络或YOLOv5分类项目进行训练与评估。整个压缩包共387个文件以jpg、jpeg轨道图像为主体辅以py脚本、json类别说明和png示例图包体约159.73MB目录结构简洁标签与数据一一对应支持脚本化批量读取。目前已有191人学习参考适合作为轨道伤损检测研究的基准数据集用于数据增强、模型调优、迁移学习及不同分类算法的对比验证。1. 铁路轨道故障图像识别数据集380张已标注图像能撑起一次完整的二分类实践第一次拿到这个铁路轨道故障图像识别数据集的时候我第一反应是怀疑自己看错了——只有380张上下后来我把整个流程跑通才意识到这恰恰是轨道运维场景里最真实的样本状态现场能拍到的故障图永远不可能像通用视觉数据集那样动辄几万张能用的标注图就这么多。这个数据集只做一件事区分轨道的损坏和未损坏两类每一张图都对应json标注文件作者已经把训练集、验证集、测试集按目录分好还带一个show脚本方便可视化。对刚接触图像识别分类落地的人来说它就是那种「麻雀虽小、五脏俱全」的练习场对有经验的人来说它更像个测试台用来验证小样本分类流程的稳定性。2. 数据集结构拆解损坏/未损坏二分类与JSON标注规则2.1 训练集、验证集、测试集的目录组织方式拿到的压缩包解开之后不是一堆散图而是按「训练集、验证集、测试集」三个目录分好的。每个目录下又按类别放图也就是说你可以不写任何解析代码直接用torchvision的ImageFolder或者YOLOv5的dir格式去读。这个设计很务实——做分类任务最忌讳的就是数据集到手先花半天时间去理文件结构。rail_track_dataset/ ├── train/ │ ├── damaged/ # 损坏 │ │ ├── download (1).jpeg │ │ ├── images (4).jpeg │ │ └── ... │ └── undamaged/ # 未损坏 │ ├── download (2).jpeg │ └── ... ├── valid/ │ ├── damaged/ │ └── undamaged/ ├── test/ │ ├── damaged/ │ └── undamaged/ └── label_config.json目录命名和文件命名有个细节值得注意源文件的图片名是download (N).jpeg和images (N).jpeg混着的说明这批图是巡检人员从不同渠道收集来的不是用固定相机机位连续拍的。好处是样本的拍摄角度、光线条件天然有差异模型学到的特征不容易过拟合到单一机位坏处是不同渠道的图分辨率可能不一致后面训练图像识别模型时需要对输入尺寸做统一resize。我一般会先跑一遍show脚本确认两张图的尺寸分布再做后续处理。2.2 读懂json标注文件类别映射与图像文件名的关联逻辑这个数据集的类别信息没有写在图片文件名里而是集中在json文件里。打开label_config.json里面维护的就是类别名和编号的映射。常见做法是写成下面这种结构这个数据集也遵循类似的约定{ categories: [ {id: 0, name: damaged, description: 轨道表面存在损伤、裂纹或明显缺陷}, {id: 1, name: undamaged, description: 轨道表面正常无明显缺陷} ], train: { download (1).jpeg: 0, images (4).jpeg: 1 }, valid: {}, test: {} }这里有两个信息要读透第一个是「分类个数2」这个描述里强调的东西——这是个二分类问题损失函数用CrossEntropyLoss(2)即可第二个是类别0和1的语义damaged对应的照片是模型真正要抓的目标undamaged是负样本。实际训练时你会发现这个数据集的负样本未损坏数量往往略多这符合巡检场景里坏轨是少数派的真实分布。如果你拿到的json不是这种结构而是按文件名数组组织的也别慌只要保持「文件名→类别id」的映射关系不变加载逻辑都是通用的。2.3 show脚本一条命令把标注结果可视化出来作者给了一个show脚本我建议拿到数据的第一件事不是训练而是先跑它。可视化的目的不是看个热闹而是确认三件事图片能不能被正常解码、标注的类别和图片内容是否对得上、是否存在名实不符的脏数据。python show.py --data_dir ./rail_track_dataset --save_dir ./preview脚本逻辑上一般是这样遍历train/valid/test目录下的所有类别文件夹随机抽若干张图把图片路径和类别名拼到画面上然后保存到save_dir。如果某个目录里混进了损坏的图片文件比如JPEG解码失败这一步就会直接报错。跑完show脚本你手上就有了一份可视化的「数据体检报告」。我习惯把每类图各看20张尤其是damaged这一类确认里边没有把「未损坏」标成「损坏」的低级错误。数据标注的质量直接决定模型上限这个环节偷懒后面所有训练都是白费。3. 用CNN分类网络吃下这批数据训练流程与参数设置3.1 自定义Dataset类从目录结构到tensor流的转换数据量小torchvision的ImageFolder可以直接用但我更推荐自己写一个Dataset类原因有两个一是后面加数据增强、加样本权重方便二是对数据集的掌控感更强出问题的时候能快速定位是数据加载的锅还是模型的锅。import os import json from PIL import Image from torch.utils.data import Dataset import torchvision.transforms as T class RailTrackDataset(Dataset): def __init__(self, root_dir, split, transformNone): self.root_dir root_dir self.split split self.transform transform # 读取json拿到当前split下的文件名到类别id的映射 with open(os.path.join(os.path.dirname(root_dir), label_config.json)) as f: config json.load(f) self.name2id config[split] self.image_paths [] self.labels [] for fname, label in self.name2id.items(): for sub in [damaged, undamaged]: cand os.path.join(root_dir, split, sub, fname) if os.path.exists(cand): self.image_paths.append(cand) self.labels.append(label) break def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img Image.open(self.image_paths[idx]).convert(RGB) if self.transform: img self.transform(img) return img, self.labels[idx]这段代码的核心逻辑是从json里读出当前splittrain/valid/test的文件名和类别映射再去对应类别的子目录里找真实路径。有个细节是convert(RGB)如果源图是灰度图或者带alpha通道的PNG这一步能保证所有图统一成三通道避免后面模型输入维度不一致。3.2 轻量CNN选型与训练超参小数据集防过拟合是头等大事380张左右的样本量ResNet50这个规模直接上必然过拟合。我一般会在这种量级下选ResNet18或者退一步用一个4层卷积的小网络关键是配合数据增强把「模型容量 vs 样本数量」压到平衡点。import torch import torch.nn as nn import torchvision.models as models import torchvision.transforms as T # 训练阶段的数据增强resize 随机翻转 颜色抖动 train_transform T.Compose([ T.Resize((224, 224)), T.RandomHorizontalFlip(p0.5), T.ColorJitter(brightness0.2, contrast0.2), T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) model models.resnet18(pretrainedTrue) model.fc nn.Linear(model.fc.in_features, 2) criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-4) # 训练回合控制在30-50之间配合早停参数选择上有几个讲究pretrainedTrue代表用ImageNet上预训练的权重初始化。只有380张图从零开始训一个ResNet18几乎不可能收敛到可用状态但迁移学习可以让模型在几十个epoch内快速达到80%以上的准确率。学习率这里用了1e-4而不是默认的1e-3也是因为样本少、微调阶段的学习率必须更低否则预训练权重会被快速破坏。python train.py --epochs 50 --batch_size 16 --lr 1e-4batch_size我用16。样本太少时batch越大每个batch对梯度的估计越偏小batch反而能带来一定的正则化效果。3.3 评估策略二分类问题别看Acc看混淆矩阵训练结束之后评估阶段有个典型的误区只看accuracy。在损坏/未损坏这种二分类场景里如果未损坏样本占多数模型只要全部猜undamaged就能拿到80%的准确率但这个模型没有任何生产价值。我习惯把测试集的预测结果整理成混淆矩阵重点看Recall和Precision两个指标。from sklearn.metrics import confusion_matrix, classification_report y_true [] y_pred [] model.eval() with torch.no_grad(): for images, labels in test_loader: outputs model(images) _, predicted torch.max(outputs, 1) y_true.extend(labels.tolist()) y_pred.extend(predicted.tolist()) print(classification_report(y_true, y_pred, target_names[damaged, undamaged])) print(confusion_matrix(y_true, y_pred))classification_report输出的Precision和Recall是针对每个类别的尤其看damaged那一行如果damaged的Recall偏低说明模型把很多实际损坏的轨道漏掉了——这在运维场景里是不可接受的。宁可误报Precision低一点也不能漏报这个认知要提前跟评价指标对齐。4. YOLOv5分类模式把380张图喂给yolov5-cls4.1 目录格式转换从train/valid/test拆分的目录到YOLO目录约定YOLOv5分类模式要求的数据组织方式和上一章的数据集原始结构非常接近它的约定是根目录下建train和val两个目录每个目录下按类别名建子目录子目录里放对应类别的图片。这个数据集天然满足这个条件只是验证集目录名从valid变成了val做一次符号链接或目录重命名就行。# 建YOLOv5分类训练所需目录结构 cd rail_track_dataset cp -r valid val # 得到 yolo_cls_data/train 和 yolo_cls_data/val如果你用的是基于YOLOv5分类项目的代码仓库那么数据目录结构就是上面这个样子。注意test集在这里用不上YOLOv5分类模式用val目录做验证。4.2 训练命令与超参调整图像尺寸、epochs与batch的取舍基于yolov5的分类训练入口是classify/train.py命令行参数和检测模式有差异主要的超参集中在img、batch、epochs三个地方。以这个数据集的体量我给的参考配置是python classify/train.py \ --model yolov5s-cls.pt \ --data ./yolo_cls_data \ --img 224 \ --batch 16 \ --epochs 80 \ --workers 4 \ --project ./runs_cls \ --name rail_track_clsimg设224而不是640是因为分类任务对这种局部纹理裂纹、破损并不需要检测模式那么高的分辨率224够用且显存占用小。epochs设80配合YOLOv5内置的cosine学习率调度小数据集能充分收敛。训练结束后在runs_cls/rail_track_cls/weights/下会产出best.pt和last.pt直接取best.pt做推理。推理命令也简单python classify/predict.py \ --weights runs_cls/rail_track_cls/weights/best.pt \ --source ./yolo_cls_data/test \ --img 224命令跑完后会在runs_cls/rail_track_cls/predict/下生成带类别标签和置信度的结果图。这一步有一个容易被忽略的细节source路径指向test目录时YOLOv5会按目录读取并推理但默认不计算准确率。想看准确率需要把图片放到val目录结构里再跑predict它才会同时输出top1 accuracy和confusion matrix。4.3 类别标签的解析逻辑为什么说YOLOv5分类是「目录即标签」YOLOv5分类模式里类别id不是写在json里的而是由目录名排序生成的。比如damaged和undamaged两个子目录排序后damaged在前对应class 0undamaged对应class 1。这和数据集json文件里的映射刚好一致damaged0undamaged1所以不需要额外改代码但如果你自己换了目录名比如把damaged改成broken类别顺序就会变训练脚本不会报错但语义就错位了。5. 避坑指南小样本二分类最常见的五个翻车现场5.1 现象训练loss持续下降验证集准确率纹丝不动原因这基本就是过拟合的典型信号。380张图对ResNet这类模型来说容量太大了模型开始硬记训练集里的噪声而不是学习轨道的通用特征。解决换resnet18甚至更小的网络开启更强的数据增强比如加上RandomRotation和RandomAffine同时把训练轮次降到30以内。还有一个检查点看看验证集里undamaged占比是否过高如果模型全猜成undamaged指标也会看起来「纹丝不动」。5.2 现象换一台机器训练训练结果完全对不上原因torchvision的ImageFolder和YOLOv5分类模式都是按目录名的字母顺序生成类别编号如果一台机器上目录名是damaged、undamaged另一台机器上被人改成了1_damaged、0_undamaged编号顺序就反了。模型本身感知不到这个变化但推理结果标签全反。解决每次训练前固定一份类别映射表写死在代码里或者和模型权重一起保存别依赖目录排序。5.3 现象show脚本跑出来图片全是黑的原因这批图的来源渠道多样有一部分图片是CMYK色彩空间的JPEGPIL读取后如果没做convert显示的时候通道映射就错乱。解决show脚本里或者Dataset的__getitem__里统一用Image.open(path).convert(RGB)确保所有图都是RGB三通道。这个问题在训练阶段不会报错但可视化时非常迷惑人。5.4 现象验证集准确率很高一到真实巡检照片就翻车原因数据集的test集和train集来自同源分布模型在「家庭作业」上考得好不代表能处理现场的新场景。380张图的覆盖度有限不同光照、不同轨枕、不同锈蚀程度的图都可能影响判断。解决用数据增强模拟光照变化做训练更实际的做法是把预训练模型当成特征提取器只微调最后一层泛化能力通常比全量微调更稳。5.5 现象YOLOv5分类训练报错ImageFolder not found原因classify/train.py对数据目录的要求是train和val两个子目录必须存在于--data路径下只给了test没给train就会直接报错。解决确认data目录下train和val都存在val目录里的类和train完全一致。还有个隐藏坑val目录下如果包含某些系统自动生成的.DS_Store或Thumbs.db文件torchvision在读取时可能把它们也当图片路径处理建议清一下。6. 从380张到能用的模型数据增强与可解释性验证6.1 数据增强是第一道防线但不是越多越好小样本场景下数据增强不只是为了让模型多看点数据更是为了告诉模型「这些因素变化了类别不变」。对轨道故障这种任务轨道的纹理方向、光照角度、拍摄距离是变量但裂纹和破损的本质特征不变。我把增强策略分为两级第一级是轻度增强包括随机翻转、小角度旋转、轻微颜色抖动用来日常训练第二级是重度增强加上随机擦除、高斯模糊、灰度化用来做对抗性的模型鲁棒性测试。train_transform_heavy T.Compose([ T.Resize((224, 224)), T.RandomHorizontalFlip(p0.5), T.RandomRotation(degrees15), T.ColorJitter(brightness0.3, contrast0.3, saturation0.2), T.RandomErasing(p0.3, scale(0.02, 0.2)), T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])这里的RandomErasing可以把随机区域抹掉强迫模型关注轨道的整体纹理而不是某个局部亮点。用它训练后再用原始test集验证你会发现准确率略有下降但对真实拍摄抖动、遮挡场景的适应力明显上升。6.2 用Grad-CAM验证模型在看什么准确率再高也得确认模型关注的区域是不是合理。我一般在训练完最后一个模型后固定跑一个Grad-CAM做可解释性验证把damaged类别的激活热力图叠加到原图上。如果热力点集中在轨道表面的裂纹区域说明模型学到的是物理特征如果热力点散落在背景草丛或者轨枕边缘那这个模型就是在走捷径要立刻回炉。from pytorch_grad_cam import GradCAM from pytorch_grad_cam.utils.image import show_cam_on_image cam GradCAM(modelmodel, target_layers[model.layer4[-1]]) # input_tensor 为单张预处理后的图 grayscale_cam cam(input_tensorinput_tensor, targetsNone) visualization show_cam_on_image(rgb_img / 255.0, grayscale_cam[0])target_layers指向的是模型最后一组卷积层resnet18里就是model.layer4[-1]。越深的层语义越强空间细节越少如果想要更细粒度的定位可以换到layer3但噪音也会变多具体用哪层需要根据热力图效果来回调。6.3 一个能落地的习惯把训练配置、评估指标和模型文件一起固化被小样本数据集坑过几次之后我养成了一个习惯每次训练结束除了best.pt和最后的混淆矩阵截图一定把训练命令、数据增强参数、类别映射表、评估指标打包存成一个config.json放在模型目录里命名带日期和epoch数。下次想复现或者换一个人接手照着一个json就能完整还原训练环境。这个习惯救过我一次——项目隔了三周要重新改版原训练脚本被同事动过数据增强里多了一行RandomResizedCrop如果当时没留配置记录我根本没法定位准确率下降的根源。从那以后我每训一次模型都强制走一遍「数据体检→训练→混淆矩阵→Grad-CAM→配置固化」这个流程缺一步都觉得不踏实。380张图的数据集看着小但把这一套流程完整走一遍你收获的不只是一个模型而是一条以后面对任何小样本图像识别任务都不会慌的处理流水线。希望帮到你。本文还有配套的精品资源点击获取