钢材表面缺陷检测：基于PyTorch的语义分割实战

简介面向计算机相关专业毕业设计、期末大作业及项目实战学习者这一基于Python的钢材表面缺陷检测与分割竞赛解决方案覆盖了从数据增强、网络构建、损失函数设计到训练评估的完整流程。压缩包共7个文件包含5个Python脚本——分别实现在线数据增强、模型定义、Lovász Softmax损失、训练及带增强评估等核心模块——外加1个可直接加载的预训练权重文件和1份说明文档整体仅3.5MB结构紧凑、易于复现。项目源自导师指导并认可的高分设计评审98分源码均经过本地编译与严格调试可稳定运行既能帮助初学者快速理解语义分割的工程代码组织方式也为进阶者提供了良好的二次开发基础。目前已有83人学习下载适用于课程设计、竞赛备赛练习也可作为钢铁表面质量检测方向的基础参考工程。1. 为什么钢材表面缺陷检测要先做分割而不是只做分类钢材表面缺陷检测与自然图像分类的差异不在网络结构而在输出维度。裂纹、麻点、划痕这类缺陷往往细长、低对比度灰度分布和氧化铁皮背景高度重叠全图分类只能告诉你“有没有缺陷”却给不出缺陷的位置、面积和形态而这些恰恰是竞赛评分和产线复检最看重的信息。竞赛评测通常提交像素级 mask、用 mIoU 计分所以完整解决方案必然落到语义分割或实例分割。这套基于 Python 的源码包把整条链路串好了train.py 管理训练流程model.py 定义分割网络lovasz_softmax.py 提供面向 IoU 的损失OnlineAugment.py 做在线增强eval_with_aug.py 负责带增强的评估。正在做毕业设计、期末大作业或第一次接触分割任务的 Python 开发者可以直接对照源码跑通再替换成自己的数据。2. 数据解构标注格式、类别分布与 Dataset 输入侧设计2.1 钢材表面缺陷的类型与分布拿到源码包先打开 README.md确认数据目录和标注格式。公开钢材表面缺陷数据的标注粒度有两种框和像素 mask。竞赛里为了算 mIoU几乎都是像素级标注具体编码又分两种——单通道 PNG 和 RLE 字符串。图像本身多数是灰度图缺陷类型基本落在以下六类缺陷类型典型外观标注难点裂纹细线状、低对比易与划痕混淆夹杂颗粒状、亮度突变边界判定困难斑块大面积灰度异常边缘过渡带长麻点密集小圆坑单个体积小、数量大氧化铁皮压入块状、纹理不均匀和背景融为一体划痕长条状、方向性强可能横跨整幅图像注意一个容易被忽略的分布事实无缺陷的负样本占比经常超过三分之一。如果直接拿全图分类的思路去做分割模型很容易学会把所有像素预测成背景loss 看起来在下降mIoU 却一直是 0。所以第一个环节要做的不是选模型而是先统计 mask 的面积分布和类别占比再决定损失函数和类别权重怎么配。2.2 带预处理的 Dataset 类train.py 里的数据入口通常是一个自定义 Dataset读取图像和 mask缩放、转 tensor、归一化。最常见的实现是 OpenCV 读图训练时缩放到统一尺寸参考写法如下import cv2 import numpy as np import torch from torch.utils.data import Dataset class SteelDefectDataset(Dataset): def __init__(self, image_paths, mask_paths, size(256, 512), augmentNone): self.image_paths image_paths self.mask_paths mask_paths self.size size self.augment augment def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img cv2.imread(self.image_paths[idx], cv2.IMREAD_GRAYSCALE) mask cv2.imread(self.mask_paths[idx], cv2.IMREAD_GRAYSCALE) img cv2.resize(img, self.size, interpolationcv2.INTER_LINEAR) mask cv2.resize(mask, self.size, interpolationcv2.INTER_NEAREST) # mask 必须用最近邻重采样线性插值会在缺陷边缘生成 0.x 的中间值 # 这些中间值会被当成额外类别污染交叉熵损失 if self.augment is not None: augmented self.augment(imageimg, maskmask) img augmented[image] mask augmented[mask] img_t torch.from_numpy(img).float().unsqueeze(0) / 255.0 mask_t torch.from_numpy(mask).long() return img_t, mask_t这段代码里有两个点值得展开。第一unsqueeze(0)把灰度图从(H, W)变成(1, H, W)因为 PyTorch 的 Conv2d 要求输入是四维(B, C, H, W)这里 batch 维度还没加先补上 channel 维度。第二归一化直接除以 255 而不是逐通道做 mean/std速度快而且对钢材这种灰度分布相对固定的场景和完整归一化没有明显差别。真正影响精度的不是归一化方式而是 mask 的 resample 方式所以 mask 统一用cv2.INTER_NEAREST。提示如果 mask 的值域出现 0 和 255说明原数据集把标签存成了可视化 PNG训练前必须把 255 改成 1否则交叉熵会把它当成独立的第二个类别。2.3 RLE 解码与缺陷感知划分如果竞赛给的是 RLE 字符串要先解码成二维 mask。解码逻辑本身不复杂但有边界条件要处理空字符串表示无缺陷应当返回全零矩阵不能直接 split 报错。def rle_decode(rle_string, shape): if rle_string is None or rle_string : return np.zeros(shape[0] * shape[1], dtypenp.uint8) s rle_string.split() starts np.asarray(s[0::2], dtypeint) - 1 lengths np.asarray(s[1::2], dtypeint) flat np.zeros(shape[0] * shape[1], dtypenp.uint8) for start, length in zip(starts, lengths): flat[start:start length] 1 return flat.reshape(shape)RLE 的坐标区间经常是 1-based所以starts要减 1lengths描述的是开区间长度直接用切片flat[start:start length]赋值即可。钢材缺陷的 RLE 通常是按行扫描的所以 decode 完 reshape 成(H, W)时顺序要和标注约定一致否则 mask 会整体错位。数据划分也不能直接 random split要按缺陷类别做分层。同一张图里可能同时出现多个缺陷类别一个简单的做法是先统计每张图包含的类别 id再按类别频率做 StratifiedGroupKFold保证 train/val 中的类别比例接近。训练尺寸和最终推理尺寸不一致时记得在验证阶段把预测 mask 还原到原图尺寸再做 mIoU 统计否则分数会虚高或虚低 1-2 个点。3. model.py 与 lovasz_softmax.py分割网络结构和损失函数选型3.1 为什么选 encoder-decoder 而不是直接接全连接钢材缺陷的尺度跨度很大划痕可能覆盖整幅图像麻点却只有几个像素。如果主干一路下采样到 1/32一个 8×8 像素的麻点在特征图上只剩 1×1几乎不可能被恢复。所以 model.py 里常见的分割网络是 DeepLabV3 或 UNet 这类 encoder-decoder 结构encoder 提供语义decoder 恢复空间分辨率。工业上最稳的做法是拿一个预训练好的 ResNet 做 backbone然后在最后一层接 ASPP空洞空间金字塔池化用不同膨胀率的空洞卷积并行捕捉多尺度上下文最后上采样回原分辨率。参考实现import torch import torch.nn as nn import torchvision class ASPP(nn.Module): def __init__(self, in_channels, out_channels256, rates(6, 12, 18)): super().__init__() self.branches nn.ModuleList() for rate in rates: self.branches.append( nn.Conv2d(in_channels, out_channels, kernel_size3, paddingrate, dilationrate, biasFalse) ) self.project nn.Sequential( nn.Conv2d(len(rates) * out_channels, out_channels, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue), nn.Conv2d(out_channels, 1, kernel_size1) ) def forward(self, x): outs [branch(x) for branch in self.branches] return self.project(torch.cat(outs, dim1))这里 rates 是膨胀率。rate6 表示卷积核在 3×3 的采样点上隔 6 个像素取一次感受野瞬间变大但参数不增加。把三个不同膨胀率的特征拼到一起模型就能同时看到细划痕和大的氧化铁皮压入区域。output 通道最后接 1是因为很多竞赛只有一个“缺陷”类别输出一张(B, 1, H, W)的 logits 图就够了多类别场景把out_channels改成类别数即可。3.2 Lovász softmax 为什么比 cross entropy 适合 mIoU 竞赛竞赛评分是 mIoU交叉熵却是逐像素独立优化两者在语义上是错位的。IoU 是集合层面的重叠度量缺陷像素少的时候交叉熵会把大量梯度分配给背景前景区域学得很慢。Lovász softmax 的思路是把离散的 Jaccard 指数扩展成连续可导的损失让梯度直接优化 IoU 的代理值。那个 lovasz_softmax.py 文件的核心是排序梯度的过程简化为 lovasz_graddef lovasz_grad(gt_sorted): p len(gt_sorted) gts gt_sorted.sum() intersection gts - gt_sorted.float().cumsum(0) union gts (1 - gt_sorted).float().cumsum(0) jaccard 1.0 - intersection / union if p 1: jaccard[1:p] jaccard[1:p] - jaccard[0:-1] return jaccardgt_sorted是按预测概率排序后的标签0 表示背景1 表示缺陷。cumsum 是前缀和intersection通过“总正样本数减去已扫描的正样本累计数”不断更新union则是把背景的累计数加到总正样本上。最终jaccard计算每个排序位置的 IoU 增量再取差分得到梯度权重。简单理解就是预测越离谱的 hard negativeLovász 给它的梯度权重越大。3.3 损失搭配和类别权重实际训练很少单独用 Lovász因为它对冷启动阶段不友好。我一般会把它和交叉熵按loss 0.7 * lovasz 0.3 * ce混合。前几个 epoch交叉熵提供稳定的逐像素梯度等 mask 预测有一点形状之后Lovász 再主导优化方向。这个比例不固定数据类别特别不平衡时可以把 Lovász 权重提到 0.8。损失函数优化目标适用场景注意点CrossEntropy逐像素分类通用基线对类别不平衡敏感Dice Loss前景区域重叠二分类 mask训练初期容易震荡Lovász SoftmaxJaccard 指数mIoU 竞赛 / 掩膜评价需要对 logits 先过 softmax还有一个细节由于 mask 的前背景比例极端可以在交叉熵里把背景的类别权重调低例如weighttorch.tensor([0.5, 1.0])。这不会直接提升 mIoU但能防止 CNN 前几轮就把所有像素判定为背景导致 loss 下降到一定数值后 mIoU 卡死在 0。4. OnlineAugment.py 与 train.py在线增强、训练循环与 checkpoint 管理4.1 在线增强为什么要放在数据读取端钢材图像往往来自同一卷钢带的连续拍摄直接连续采样会碰到大量高度相似的帧模型很快就会记住“背景统计量”泛化性极差。OnlineAugment 的好处是每个 epoch 都用不同的随机变换组合相当于每个 epoch 得到一份新数据而离线增强会把磁盘占用放大 N 倍。放在 Dataset 里还避免了在 GPU 上做增强的同步开销。4.2 增强组合的常见做法这是一个稳健、不会破坏缺陷物理形态的增强管线使用 albumentationsimport albumentations as A train_aug A.Compose([ A.HorizontalFlip(p0.5), A.VerticalFlip(p0.5), A.RandomRotate90(p0.5), A.OneOf([ A.GaussianBlur(blur_limit(3, 5), p0.5), A.MotionBlur(blur_limit3, p0.5), ], p0.3), A.RandomBrightnessContrast(brightness_limit0.15, contrast_limit0.15, p0.5), ])注意这里几乎没有使用 Scale/Translate/Rotate 这类任意角度形变。钢材表面的氧化铁皮压入和麻点是刚性的不存在物体被扭曲的物理可能过度形变反而会让模型学到错误的形状先验。RandomRotate90 本质是重排像素不会产生新的插值伪影在钢材灰度图上非常安全。如果担心图像上下方向在生产线上有语义可以去掉 VerticalFlip。4.3 train.py 里的训练循环训练循环看起来和普通分割任务差别不大但有几个容易被忽略的步骤。一个标准 epoch 的骨架如下for epoch in range(start_epoch, total_epochs): model.train() for images, masks in train_loader: images images.cuda(non_blockingTrue) masks masks.cuda(non_blockingTrue) logits model(images) probas torch.softmax(logits, dim1) lovasz_loss lovasz_softmax(probas, masks) ce_loss nn.CrossEntropyLoss()(logits, masks) loss 0.7 * lovasz_loss 0.3 * ce_loss optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() val_miou validate(model, val_loader)clip_grad_norm_是为了处理偶发的大梯度。钢材图像里的极端像素值比如强反光会产生很大的 loss 尖刺把梯度范数限制在 1.0 能避免一个 step 就把权重冲坏。validate函数每次返回验证集 mIoU用它来判断是否保存新的 best checkpoint而不是看训练 loss。4.4 checkpoint 与断点续训竞赛工程里 checkpoint 不能只存模型权重至少要包含四样东西torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), best_miou: best_miou, }, checkpoint.pth)恢复训练时把 epoch 和 optimizer 状态一并载入才能继续原本的学习率调度。常见的错误是只 loadmodel_state_dict导致恢复后 learning rate scheduler 从头开始学习率瞬间跳高模型几步之内就发散。超参数上以这块竞赛数据为参考常用范围如下参数参考值说明输入尺寸256×256 / 256×512显存不足时优先缩短短边batch size8-32与学习率同步调整优化器AdamWlr 从 1e-4 开始调度器CosineAnnealingLRT_max 设为总 epoch 数训练轮数50-100验证 mIoU 连续 3 轮不升则早停梯度裁剪max_norm1.0防止异常像素值导致发散5. eval_with_aug.pyTTA 推理、mIoU 统计与断点续训的细节5.1 从 checkpoint 恢复模型状态eval_with_aug.py 先要解决的是“怎么从训练状态切到推理状态”。关键是把权重载进来再切换成评估模式ckpt torch.load(checkpoint.pth, map_locationcuda:0) model.load_state_dict(ckpt[model_state_dict]) start_epoch ckpt[epoch] 1 best_miou ckpt[best_miou] model.eval()如果你的 checkpoint 里有优化器状态并且想继续训练就把optimizer.load_state_dict(ckpt[optimizer_state_dict])也加上再进入 train 循环如果只是做验证和推理则不要加载 optimizer。model.eval()必须放在循环外它只会改 BatchNorm 和 Dropout不会影响 Conv 权重但漏掉它会让 mIoU 因为 BatchNorm 统计量不一致而掉 1-3 个点。5.2 多尺度 TTA 的叠加TTA 的思路是让模型在推理时看到更多“视角”然后把概率平均。一个常用组合是三种尺度加水平翻转import torch.nn.functional as F def predict_tta(model, img, scales(0.8, 1.0, 1.2)): H, W img.shape[-2:] probs [] for scale in scales: scaled F.interpolate(img, scale_factorscale, modebilinear, align_cornersFalse) logits model(scaled) logits F.interpolate(logits, size(H, W), modebilinear, align_cornersFalse) p torch.softmax(logits, dim1) logits_f model(torch.flip(scaled, dims[3])) logits_f F.interpolate(logits_f, size(H, W), modebilinear, align_cornersFalse) p p torch.flip(torch.softmax(logits_f, dim1), dims[3]) probs.append(p) return torch.stack(probs).mean(dim0)每个尺度内部先各自处理翻转再把所有尺度求平均避免不同尺度权重失衡。多尺度 TTA 通常能带来 1-2 个百分点的 mIoU 提升代价是推理时间乘以 6所以要权衡是在线评分还是离线提交。5.3 评估口径要统一最后的 mIoU 统计要保证训练、验证、推理三个阶段对 mask 的插值方式一致。训练时代码用INTER_NEAREST缩 mask验证时就别改用INTER_LINEAR否则边缘像素的对齐方式不同mIoU 的波动会掩盖真实的增益。统计指标时按类别分别计算 IoU再求平均得到 mIoU不要把背景类算进去。最后提一个实际踩过的坑把model.pth导出给后端推理时模型输出 logits 的形状是(B, C, H, W)但如果脚本里有一次.permute(0, 2, 3, 1)最后 mask 会以(B, H, W, C)输出。在 torchvision 或 Flask 接口里这两者混用时 argmax 的轴对不上得到的 mask 就是错的。建议在 eval 脚本末尾加一行断言assert pred_mask.shape (batch_size, height, width), mask 形状必须是 B,H,W这样就算之后有人改动了前向分支提交前也能被挡在门口。本文还有配套的精品资源点击获取