UNet实战钢材表面缺陷检测：从NEU-DET数据集到语义分割落地

简介面向工业质检场景的UNet钢材表面缺陷检测实战项目基于NEU-DET数据集覆盖从数据预处理到模型部署的完整流程。资源包含2000个文件以1800个XML标注文件与180张JPG缺陷图像为主另有Python脚本、Pyc、Jupyter Notebook、ONNX模型及说明文档压缩包约122MB。目前已有888人学习下载适合计算机视觉学习者和工业质检开发者参考也适用于相关课题设计与毕设扩展。项目中包含探索性数据分析、数据划分、训练工作流等模块可帮助读者理解划痕、凹坑、夹杂等典型缺陷的样本分布掌握UNet的编码器-解码器结构及跳跃连接在图像分割中的应用并学会通过数据增强提升模型鲁棒性。提供可运行的Python脚本和ONNX导出模型便于在真实场景中快速部署和验证是一份兼顾算法原理与工程实现的优质实战资料。1. 用 UNet 啃下钢材表面缺陷检测NEU-DET 这个数据集到底难在哪一条热轧带钢产线钢坯从加热炉出来经过粗轧、精轧表面温度大几百摄氏度氧化铁皮、冷却水、辊道振动全混在一起。质检工位上的老师傅要在几秒内判断这一卷钢有没有裂纹、划伤、麻点眼睛盯着屏幕一盯八小时漏检一次后面整卷报废。这个场景一直想自动化传统机器视觉用阈值分割、边缘检测做表面缺陷识别前提是缺陷和背景的光照分布要干净产线上根本做不到。所以这几年工业缺陷检测基本都转向了深度学习而 UNet 就是切入这个方向性价比最高的模型之一结构简单、标签需求不算苛刻、单卡就能训练。NEU-DET 是东北大学公开的热轧带钢表面缺陷数据集1800 张灰度图覆盖 6 类典型缺陷正好拿来验证 UNet 在这个任务上的可行性。这篇从数据准备、模型搭建、训练评估到避坑完整过一遍适合想入行工业视觉但手里只有一张消费级显卡的工程师也适合课程设计、毕业设计要快速出效果的读者。2. 从标注框到分割 maskNEU-DET 数据准备与数据增强的落地做法2.1 NEU-DET 数据集结构六类缺陷和 1800 张图的信息量NEU-DET 全称是 Northeastern University Surface Defect Database图像是热轧带钢表面的灰度图原始分辨率 200x200每张图只包含一类缺陷的单个或多个实例。整个数据集 1800 张六类缺陷每类 300 张。这个规模和 ImageNet 比是很小但在工业缺陷场景里已经算能用很多工厂内部能拿到的标注数据往往只有几百张。六类缺陷如下类别英文名典型外观检测难度裂纹Crazing网状细线对比度低难夹杂Inclusion颗粒状暗斑形状不规则中斑块Patches大面积灰暗区域中麻点Pitted Surface细小点状凹坑难氧化铁皮压入Rolled-in Scale片状深色压痕中划伤Scratches线性亮条或暗条易下载解压后的常见组织方式是一个 images 目录放图一个 annotations 目录放标注文件。标注格式在不同渠道的打包里不完全一样有的给 VOC 风格的 XML有的给 YOLO 风格的 txt每行是类别和归一化中心点坐标、宽高。这给第一次用的人造成不少困扰因为网上很多教程默认你拿到的是某种固定格式结果代码一跑就报错。我先说结论不管拿到的是哪种格式都要先写一个小脚本统计一遍标注框数量、类别 ID 分布不要直接开训。这一步几十秒能挡住后面大半的翻车。另外要注意NEU-DET 的标注是目标检测框不是像素级分割掩码。UNet 做的是语义分割理论上需要每个像素都标注但项目里常见做法是先用检测框生成粗 mask 跑通全流程再针对重点缺陷类型人工补细标注。这种方式训练出来的模型边界会偏大但作为基线验证完全够用。后面第 5 章会专门讲这个做法的坑。2.2 把 bbox 标注批量转成 UNet 需要的像素级 mask先定义一个函数把 XML 标注文件解析成 numpy 数组形式的 mask。mask 的形状和原图一样是 (H, W)像素值用类别 ID 填充背景为 0六类缺陷依次映射为 1 到 6。解析 XML 用 Python 自带的 xml.etree.ElementTree 就够了不需要额外装库。import os import numpy as np import xml.etree.ElementTree as ET # 类别名到ID的映射顺序按NEU-DET常见命名梳理 CLASS_MAPPING { crazing: 1, inclusion: 2, patches: 3, pitted_surface: 4, rolled-in_scale: 5, scratches: 6, } def xml_to_mask(xml_path, img_h200, img_w200): 把VOC风格的XML标注转成分割mask tree ET.parse(xml_path) root tree.getroot() mask np.zeros((img_h, img_w), dtypenp.uint8) for obj in root.iter(object): name obj.find(name).text.lower() if name not in CLASS_MAPPING: print(f未知类别: {name}, 文件: {xml_path}) continue class_id CLASS_MAPPING[name] bndbox obj.find(bndbox) xmin int(float(bndbox.find(xmin).text)) ymin int(float(bndbox.find(ymin).text)) xmax int(float(bndbox.find(xmax).text)) ymax int(float(bndbox.find(ymax).text)) # 截断到图像范围内防止标注越界 xmin max(0, min(xmin, img_w - 1)) xmax max(0, min(xmax, img_w - 1)) ymin max(0, min(ymin, img_h - 1)) ymax max(0, min(ymax, img_h - 1)) mask[ymin:ymax 1, xmin:xmax 1] class_id return mask这段代码核心就两件事解析出每个 object 的类别名和 bbox 坐标然后把 bbox 矩形区域内的像素值填成对应的类别 ID。这里有几个参数要说明。CLASS_MAPPING 的类别名大小写不统一是常态所以先做个 lower() 归一实际项目里我遇到过类别名写成 Crazing 带空格的解析后 strip() 更稳这里没加是因为怕把代码写啰嗦你用的时候建议补上。bbox 坐标如果在图像边缘外面不做截断的话填充循环会报错所以 min/max 那两行很关键。如果你拿到的是 YOLO txt 格式每个框的坐标是归一化到 0-1 的相对值转 mask 时要乘回图像宽高并处理成整数坐标def txt_to_mask(txt_path, img_h200, img_w200): 把YOLO txt标注转成分割mask每行: class_id cx cy w h mask np.zeros((img_h, img_w), dtypenp.uint8) with open(txt_path, r) as f: for line in f: parts line.strip().split() if len(parts) 5: continue cls_id int(parts[0]) 1 # YOLO从0开始mask从1开始 cx, cy, w, h map(float, parts[1:5]) xmin int((cx - w / 2) * img_w) xmax int((cx w / 2) * img_w) ymin int((cy - h / 2) * img_h) ymax int((cy h / 2) * img_h) xmin max(0, min(xmin, img_w - 1)) xmax max(0, min(xmax, img_w - 1)) ymin max(0, min(ymin, img_h - 1)) ymax max(0, min(ymax, img_h - 1)) mask[ymin:ymax 1, xmin:xmax 1] cls_id return maskYOLO 格式常见的问题是把 w 和 h 写成像素值而不是归一化值新旧版本标注工具混用容易出现这个差异。转换前先打印几行标注内容确认一下数值范围全是 0-1 之间的小数就是归一化的出现 200 左右的数值就要按像素值处理。转换成 mask 后建议顺手保存成 PNG 或者直接以 numpy 数组形式缓存下来避免每轮训练都重新解析一次标注文件1800 张图解析虽然只要几秒钟但数据集扩大到上万张时这个习惯能省不少训练准备时间。2.3 灰度图增强策略别把缺陷增强没了工业图像增强和自然图像增强有个关键差别不能随便做大幅几何变换。钢材表面缺陷里划伤和裂纹有方向性旋转 90 度后视觉特征还在但旋转 20 度、30 度后裂纹走向变斜模型学到的是各种角度的裂纹这本身不算坏事问题是现场相机安装角度固定缺陷方向分布基本一致过度增强反而让模型把精力花在处理不可能出现的角度上。我一般对 NEU-DET 这种小数据集用下面这套增强import albumentations as A from albumentations.pytorch import ToTensorV2 def build_train_transforms(): 训练集增强管道适合200x200小图 return A.Compose([ A.HorizontalFlip(p0.5), A.VerticalFlip(p0.5), A.RandomBrightnessContrast( brightness_limit0.15, contrast_limit0.15, p0.3 ), A.Rotate(limit15, border_mode0, p0.3), A.GaussNoise(var_limit(5.0, 15.0), p0.2), A.Normalize(mean(0.5,), std(0.5,)), ToTensorV2(), ]) def build_valid_transforms(): 验证集只做归一化不做增强 return A.Compose([ A.Normalize(mean(0.5,), std(0.5,)), ToTensorV2(), ])上面这套用的是 albumentations 库它是工业分割任务里最常用的图像增强库速度和 GPU 利用率都比徒手写 numpy 版本好。参数说明RandomBrightnessContrast 的亮度、对比度扰动范围控制在 15% 以内因为钢材表面反光在产线上变化较大这个扰动模拟光照变化是合理的但超过 30% 会把暗色缺陷增强成和背景一样的灰度值属于自己给自己加噪声。Rotate 的 limit15 是个折中允许小角度旋转增强多样性又不至于让划伤方向完全失真。GaussNoise 加的方差也别太大NEU-DET 本身是实验室条件下的成像质量真实产线图像噪声更大可以在训练后期把噪声增强的权重调高一点。还要强调一个容易忽略的点mask 的增强必须和图像同步。albumentations 的 Compose 会自动处理这个问题因为你在 transform 里同时传 image 和返回的 mask它内部用的是同一组随机参数。如果是自己手写增强函数最容易犯的错误是图像做了随机翻转而 mask 没翻训练时损失一直在降但验证 mIoU 上不去八成就是这个原因。另外 Normalize 的 mean 和 std 用 (0.5, 0.5)因为 NEU-DET 是单通道灰度图PyTorch 要求输入是 (B, 1, H, W) 或者 (B, 3, H, W)单通道就写一个值写成三通道的 mean(0.5, 0.5, 0.5) 会把灰度图复制成三通道浪费显存还容易在推理时搞混通道顺序。3. 搭建 UNet编码器-解码器结构与损失函数为什么要用组合3.1 UNet 主体结构双卷积、跳连接和上采样UNet 的核心思想是编码器逐层下采样提取高维特征解码器逐层上采样恢复空间分辨率中间靠跳连接把同尺度的编码器特征拼到解码器上。跳连接为什么对缺陷检测这么重要因为钢材表面的裂纹、麻点这类缺陷尺寸只有几十个像素经过四次下采样后小缺陷在深层特征图里几乎只剩一个点信息大量丢失。跳连接把浅层的边缘纹理信息直接送到解码器对应层等于让网络在恢复分辨率时一直有原始细节可以参考。我在自己的项目里做过消融去掉跳连接的 UNet 在 NEU-DET 上 mIoU 掉了 12 个百分点以上这个数字远超我的预期。下面是一个标准 UNet 的 PyTorch 实现基础通道数 base_channels64import torch import torch.nn as nn class DoubleConv(nn.Module): 两个卷积 BatchNorm ReLUUNet的基本块 def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, kernel_size3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, kernel_size3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.conv(x) class Down(nn.Module): 下采样最大池化 双卷积 def __init__(self, in_ch, out_ch): super().__init__() self.mpconv nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_ch, out_ch) ) def forward(self, x): return self.mpconv(x) class Up(nn.Module): 上采样转置卷积 跳连接拼接 双卷积 def __init__(self, in_ch, out_ch): super().__init__() # 转置卷积把特征图尺寸翻倍 self.up nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size2, stride2) self.conv DoubleConv(in_ch, out_ch) def forward(self, x1, x2): x1 self.up(x1) # 输入尺寸可能因奇数分辨率差1个像素做padding对齐 diffY x2.size()[2] - x1.size()[2] diffX x2.size()[3] - x1.size()[3] x1 nn.functional.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2]) x torch.cat([x2, x1], dim1) # 跳连接拼在通道维度 return self.conv(x) class UNet(nn.Module): def __init__(self, n_channels1, n_classes7, base_channels64): super().__init__() self.inc DoubleConv(n_channels, base_channels) self.down1 Down(base_channels, base_channels * 2) self.down2 Down(base_channels * 2, base_channels * 4) self.down3 Down(base_channels * 4, base_channels * 8) self.down4 Down(base_channels * 8, base_channels * 8) self.up1 Up(base_channels * 8, base_channels * 4) self.up2 Up(base_channels * 4, base_channels * 2) self.up3 Up(base_channels * 2, base_channels) self.up4 Up(base_channels, base_channels) self.outc nn.Conv2d(base_channels, n_classes, kernel_size1) def forward(self, x): x1 self.inc(x) x2 self.down1(x1) x3 self.down2(x2) x4 self.down3(x3) x5 self.down4(x4) x self.up1(x5, x4) x self.up2(x, x3) x self.up3(x, x2) x self.up4(x, x1) return self.outc(x)几个实现细节值得展开。DoubleConv 里每个卷积后面都跟 BatchNorm这个对工业数据的灰度图尤其重要因为不同批次的灰度分布可能差很多BN 能起到一定的归一化作用。kernel_size 固定 3padding 1 保证特征图尺寸不缩水避免后面拼跳连接时尺寸对不上。Up 模块里有一个容易被忽略的 padding 处理当输入图尺寸是奇数时下采样会损失一个像素直接 cat 会报错所以要先算出差值手动补零。200x200 的输入经过 4 次下采样变成 12x12上采样回去正好对齐但如果你后续替换成 224x224 或者其他不整除的尺寸这段逻辑就是必写项。OutConv 用 1x1 卷积把通道数压到类别数 7等于每个像素输出 7 个 logits代表 6 类缺陷加背景的概率。3.2 损失函数BCE 和 Dice 的组合比单用哪个都强钢材缺陷分割的类别不平衡问题非常典型背景像素占了图像面积的绝大多数麻点和裂纹这种细长缺陷的像素占比往往不到 1%。直接交叉熵训练模型很快就学会把所有像素预测成背景因为这么做损失就很小了。纯 Dice Loss 能缓解不平衡但梯度在小目标上容易振荡训练初期收敛慢。所以工业项目里最常见的做法是交叉熵和 Dice 加权组合让两个损失互相牵制。import torch.nn.functional as F class SoftDiceLoss(nn.Module): 多分类Dice Loss支持类别权重 def __init__(self, n_classes, class_weightsNone): super().__init__() self.n_classes n_classes self.class_weights class_weights # shape: (C,) 或 None def forward(self, logits, targets): # logits: (B, C, H, W) targets: (B, H, W) 值范围0~C-1 probs F.softmax(logits, dim1) # 转成概率 targets_onehot F.one_hot(targets, num_classesself.n_classes) # 把通道维移到后面形状变成 (B, H, W, C) targets_onehot targets_onehot.permute(0, 3, 1, 2).float() smooth 1.0 dims (0, 2, 3) # 在batch和空间维度上求和 intersection torch.sum(probs * targets_onehot, dimdims) cardinality torch.sum(probs targets_onehot, dimdims) dice (2.0 * intersection smooth) / (cardinality smooth) if self.class_weights is not None: dice dice * self.class_weights.to(dice.device) return 1.0 - dice.mean() class CombinedLoss(nn.Module): CE Dice 组合损失 def __init__(self, n_classes, ce_weight0.6, dice_weight0.4, class_weightsNone): super().__init__() self.n_classes n_classes self.ce_weight ce_weight self.dice_weight dice_weight self.dice_loss SoftDiceLoss(n_classes, class_weights) # 给少数类更高的权重背景权重压低 ce_w torch.ones(n_classes) if class_weights is not None: ce_w torch.tensor(class_weights, dtypetorch.float32) self.ce_loss nn.CrossEntropyLoss(weightce_w) def forward(self, logits, targets): ce self.ce_loss(logits, targets) dice self.dice_loss(logits, targets) return self.ce_weight * ce self.dice_weight * dice这里有几个参数要说明白。class_weights 不是必须的但加上能在类别不平衡场景里明显提高少数类的召回。我在 NEU-DET 上用的权重是类别权重背景0.5crazing1.8inclusion1.5patches1.0pitted_surface1.5rolled-in_scale1.0scratches1.0背景权重 0.5 而不是更低的原因背景占比太高压太低会导致背景区域大量假阳性推理时整个图像都被预测成缺陷这比漏检还难受。crazing 和 pitted_surface 权重高是因为这两类像素面积小、灰度对比度低不加权重的话基本学不出来。ce_weight 和 dice_weight 的搭配我试过 1:1、0.7:0.3、0.5:0.50.6:0.4 在新手期最省心Dice 占比太高训练指标波动大CE 占比太高小缺陷找回能力弱。你如果觉得自己的场景缺陷更小可以慢慢把 dice_weight 往上调但每次只调 0.05别一次跳到 0.7。3.3 训练脚本和超参数配置数据集加载部分我直接用 Dataset 类包一层重点说超参数的选择理由from torch.utils.data import Dataset, DataLoader import cv2 class SteelDataset(Dataset): def __init__(self, image_paths, mask_paths, transformNone): self.image_paths image_paths self.mask_paths mask_paths self.transform transform 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) # 确保mask类别ID连续 img img.astype(np.float32) / 255.0 mask mask.astype(np.int64) if self.transform: augmented self.transform(imageimg, maskmask) # albumentations返回的image是(H,W)要补通道维 img augmented[image].unsqueeze(0) mask augmented[mask] return img, mask训练主体按 PyTorch 常规写法这里给出核心配置BATCH_SIZE 16 EPOCHS 100 LR 3e-4 DEVICE cuda if torch.cuda.is_available() else cpu model UNet(n_channels1, n_classes7).to(DEVICE) optimizer torch.optim.AdamW(model.parameters(), lrLR, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxEPOCHS) loss_fn CombinedLoss(n_classes7, class_weights[0.5, 1.8, 1.5, 1.0, 1.5, 1.0, 1.0])超参数的选择依据输入 200x200 的小图显存占用很小6GB 显存跑 batch_size 16 绰绰有余如果是 2GB 的显卡把 batch_size 降到 8学习率相应降到 1.5e-4。学习率 3e-4 是 AdamW 在小数据集上的保守选择太大容易前期振荡后期跑飞太小则 100 个 epoch 收益不够。CosineAnnealingLR 配合 100 个 epoch 是经典组合前 40 个 epoch 学习率在缓慢下降后面 60 个 epoch 逐步逼近最低点这比固定学习率训练 100 个 epoch 的收敛质量高。weight_decay 1e-4 主要防止模型在小数据集上过拟合可以对比不加的情况验证集 mIoU 通常能高 2 到 3 个点。4. 训练与验证mIoU、可视化与效果评估全流程4.1 训练循环与验证指标计算训练循环本身没什么特别的重点是验证阶段要把 mIoU 和像素准确率算准。mIoU 是按类别算 IoU 再取平均不是把类别混淆矩阵拉平。同一个模型PA 可能达到 95% 以上但 mIoU 只有 60%因为背景像素占比高把背景猜对了 PA 就很高所以工业缺陷检测里更看重 mIoU。def compute_ious(preds, targets, n_classes7): 计算每个类别的IoU和mIoUpreds和targets都是像素标签 ious [] for cls in range(n_classes): pred_mask (preds cls) target_mask (targets cls) intersection (pred_mask target_mask).sum() union (pred_mask | target_mask).sum() if union 0: # 如果真实和预测都没有这个类不计入平均 ious.append(float(nan)) else: ious.append(intersection.item() / union.item()) valid_ious [v for v in ious if not np.isnan(v)] miou np.mean(valid_ious) if valid_ious else 0.0 return ious, miou def validate(model, dataloader, device): model.eval() all_preds [] all_targets [] with torch.no_grad(): for images, masks in dataloader: images images.to(device) masks masks.to(device) logits model(images) preds logits.argmax(dim1).cpu().numpy() all_preds.append(preds) all_targets.append(masks.cpu().numpy()) all_preds np.concatenate(all_preds) all_targets np.concatenate(all_targets) ious, miou compute_ious(all_preds, all_targets) return ious, miou这段代码要说明的细节predict 阶段先把 logits 沿类别维度取 argmax得到的是每个像素的预测类别形状是 (B, H, W)。然后累计所有验证集样本最后统一算 IoU不要在 batch 级别算完再平均那样小 batch 里缺类的情况会导致 mIoU 被低估。union 为 0 的类别直接跳过意思是这个类在整批预测和真实里都没出现不应该参与平均。如果某个类别的 IoU 是 0说明这类缺陷完全没被检测出来单独打印出来看这是最有诊断价值的输出。4.2 结果可视化把预测 mask 叠回原图看指标是抽象的缺陷分割模型的成败最终要回到图像上看。我见过 mIoU 看起来很高但某个类别的边缘全是毛刺的情况也见过指标一般但关键裂纹全检出来的模型所以每轮验证结束都生成几张可视化图存档。叠加方式如下import matplotlib.pyplot as plt # 6类缺陷对应的BGR颜色视觉区分度高 COLOR_MAP { 1: (0, 0, 255), # crazing 红色 2: (0, 255, 0), # inclusion 绿色 3: (255, 0, 0), # patches 蓝色 4: (0, 255, 255), # pitted_surface 黄色 5: (255, 0, 255), # rolled-in_scale 品红 6: (255, 255, 0), # scratches 青色 } def visualize_prediction(image, mask, pred, save_path): 把分割结果叠回灰度图上左边原图右边叠加便于对比 img_bgr cv2.cvtColor(image, cv2.COLOR_GRAY2BGR) overlay img_bgr.copy() for cls_id, color in COLOR_MAP.items(): # 预测区域高亮 overlay[pred cls_id] color # 真实区域用白色描边画轮廓方便看漏检和误检 contours, _ cv2.findContours( (mask cls_id).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE ) cv2.drawContours(overlay, contours, -1, (255, 255, 255), 1) combined np.hstack([img_bgr, overlay]) cv2.imwrite(save_path, combined)可视化图里两个信息很关键彩色区域是模型预测的缺陷位置白色轮廓是真实标注位置。如果彩色区域明显比白色轮廓大一圈说明 bbox 转 mask 的粗标注让模型学出了冗余边界如果白色轮廓附近完全没有彩色说明这类缺陷漏检了。看可视化图的顺序也有讲究先看最难类的查全率再看整体误检率最后才看数值指标。一张画满绿色预测的图mIoU 有可能是 0.6因为 patches 类面积大拉高了平均分其他类的表现被掩盖了。5. 钢材缺陷检测避坑记录五个高频翻车现场5.1 缺陷其实很小放大输入分辨率反而掉点现象很多新手拿到 200x200 的原始图总觉得分辨率太低先把图 resize 到 512 再训练。训练出来的模型 mIoU 反而比 200x200 直接训练低 3-5 个百分点推理速度还慢了不少。原因NEU-DET 的缺陷尺寸是相对固定的裂纹和麻点在 200x200 下占 20-50 个像素放大到 512 后缺陷变成线性插值出来的模糊斑块反而丢失了原始纹理信息。而且 UNet 输入分辨率变大后下采样层数不变感受野相对变小模型更难建立缺陷的全局上下文。解决NEU-DET 就用原生 200x200 训练不要放大。如果担心信息不够可以试试随机裁剪 160x160 加多尺度训练让模型见到不同的缺陷尺寸。迁移到自己的产线数据时先统计缺陷的像素分布再定输入尺寸而不是拍脑袋定一个标准值。5.2 bbox 直接转 mask框内背景被当成了缺陷现象第一版模型训练完可视化发现每个缺陷周围都有一圈预测出来的假阳性区域mIoU 不错但实际效果没法用。原因NEU-DET 原始标注是检测框2.2 节转 mask 时把整个矩形区域内所有像素都标成了缺陷类别。矩形框里除了缺陷本体还有大量正常背景模型被迫学习把背景预测成缺陷。这在划伤和裂纹上特别明显因为框的面积比实际缺陷大很多。解决有两个方向。一是训练完做后处理对预测 mask 做形态学腐蚀把边界收缩一圈代码就几行但效果直接。二是重新标注像素级 mask只对模型表现差的类别做常见做法是先用粗标注模型预测一遍再人工在标注工具里修正边界每张图几十秒。如果你是做个基线验证第一种就够了如果项目要交付必须走第二种。5.3 crazing 和 inclusion 分不清类别混淆要单独处理现象mIoU 整体 65%但查看每类 IoU 发现 crazing 只有 30%inclusion 也不高。可视化图上裂纹附近经常出现夹杂的预测色。原因这两类缺陷在灰度图像上长得很像。crazing 是细网状裂纹inclusion 是颗粒状夹杂但两者都是低对比度、形状不规则的暗色区域加上 bbox 转 mask 的粗标签进一步模糊了边界网络很难区分。解决先确认是不是标签本身造成的问题单独打印几组真实 mask 看看两类缺陷的像素分布有没有重叠区域。如果标注确实模糊给这两类在损失函数里加更高权重比如 crazing 从 1.8 提到 2.5inclusion 从 1.5 提到 2.0。还有一个思路是改成二分类模型先把所有缺陷从背景里分出来再用另一个小模型区分具体类别工程上叫级联方案在小数据集上往往比端到端多分类更可靠。5.4 验证 mIoU 不错但可视化 mask 噪声大现象模型训练完mIoU 到了 60% 以上但可视化图里到处是零散的小色块有的只有几个像素看起来像盐噪声。数量一多下游统计缺陷面积就没法用了。原因分割网络的逐像素预测天然会输出孤立小区域这是 softmax 独立预测每个像素的结果缺少空间一致性约束。产线上的图像有光照不均和粉尘干扰小噪声更容易被激发。解决后处理阶段用两个经典操作。中值滤波去掉椒盐式噪声再用连通域面积阈值滤掉太小的预测块import cv2 import numpy as np from scipy import ndimage def clean_prediction(pred, min_area20, kernel_size3): 后处理中值滤波 去除小连通域 # 对每个类别分别处理避免类别间互相污染 cleaned np.zeros_like(pred) for cls_id in np.unique(pred): if cls_id 0: continue cls_mask (pred cls_id).astype(np.uint8) # 中值滤波让边缘更规整 cls_mask cv2.medianBlur(cls_mask, kernel_size) # 标记连通域面积小于阈值的直接清零 labeled, num ndimage.label(cls_mask) sizes ndimage.sum(cls_mask, labeled, range(1, num 1)) for i, size in enumerate(sizes, start1): if size min_area: cleaned[labeled i] cls_id return cleanedmin_area 的取值不是拍脑袋定的先统计一下高铁缺陷和误检小块的真实像素面积分布取一个能保住最小真实缺陷的值。在 NEU-DET 上裂纹的连通域往往不到 30 像素min_area 设 15-20 比较安全如果设 50 会把真实小缺陷一起滤掉。kernel_size 用 3 就好大了会让细长裂纹断掉。5.5 训练不收敛的排查顺序先查数据再调参数现象训练了 30 个 epochtrain loss 只从 1.8 降到 1.6 就再也不动验证 mIoU 一直在 0.1 左右徘徊。看起来像是模型容量不够或者学习率问题于是换更大的模型、调学习率折腾半天没效果。原因这类问题在工业数据上最常见的原因反而是数据管线出问题。有一次我的训练集和验证集没有按类别做分层划分随机切分后 crazing 全部进了验证集训练集里完全没有这个类。损失降不下去因为模型在努力拟合训练集里七个类但验证集里出现一个训练时根本没见过的类mIoU 一直拉不起来。还有一次是 mask 读进来的时候opencv 默认按 BGR 三通道读灰度 mask 被读成三通道后类别 ID 全变了。解决遇到训练不收敛按这个顺序排查别先动模型结构。第一步打印一个 batch 的 image 和 mask 的 shape、数值范围、唯一值列表确认 mask 类别 ID 是 0-6 且分布正常。第二步检查 train 和 val 的类别分布是否接近用每个类别样本数除以总样本数差太多就是切分有问题用 StratifiedShuffleSplit 重切。第三步检查增强管道里有没有把 mask 转成 float 导致类别 ID 出现小数。这三步走完如果还没解决再看学习率和模型结构。损失函数在 1.6 不动但验证 mIoU 在缓慢上升的有可能是学习率太小把 3e-4 提到 1e-3 试二十个 epoch 对比曲线。我已经养成了训练时把 loss 曲线和 mIoU 曲线画在一张图上的习惯哪个先出问题一眼就知道是数据问题还是模型问题。6. 从能跑到好用轻量化改进、验证方法和部署落地技巧6.1 改进前先立基线谁改了什么必须用同一套指标说话拿到一个效果还行的 UNet 模型下一步自然是改进。我最反感的是没有基线就乱改网络结构改完说 mIoU 从 50 涨到 55结果基线是谁、用了什么数据增强、跑了多少 epoch 全是黑匣子。正确的做法是固定数据划分、输入尺寸、损失函数、训练轮数四件事只改你想验证的那个变量。我自己的固定配置是8:1:1 划分训练验证测试集200x200 输入0.6 CE 0.4 Dice100 epochAdamW 3e-4。在这套配置下跑出来的 UNet 基线 mIoU 在 65-68 左右作为后续所有改进的参照物。6.2 轻量化和注意力改进怎么选有人一上来就想把编码器换成 MobileNetV3、EfficientNet 或者加各种注意力模块。改进方向没有对错但要看你的交付场景。如果钢板检测工位上已经有工控机推理要在 30ms 内完成那轻量化编码器是合理的。如果只是实验室验证显卡跑无压力先别折腾结构把数据增强和类别权重调好收益来得更快。下面是我验证过的几个方向改进手段原理适用场景预期代价编码器换 MobileNetV2深度可分离卷积大幅减少参数部署到嵌入式或旧工控机理论上 mIoU 略降 1-3 点但推理速度快 3 倍在跳连接处加 SE 模块对通道做注意力重标定增强有用通道抑制噪声背景复杂、类别多参数量增加很小训练时间略涨Deep Supervision每个解码器层都计算监督损失深层网络训练慢、收敛差代码改动集中在训练循环涨点不明显测试时多尺度推理原图 200、放大 1.25 倍、缩小 0.75 倍分别推理再融合不追求实时性、只追求精度推理时间翻 3 倍以上我自己的经验是NEU-DET 这种 200x200 的小图瓶颈不在网络容量而在小缺陷的表达能力。与其换网络不如在数据侧花时间把裂纹类做针对性裁剪补丁让模型多看到 40x40 的局部细节配合后处理联通域过滤实际收益比换 MobileNet 骨干大得多。改进方向的选择标准很简单——你的瓶颈是速度还是精度然后对症下药。6.3 分割结果接入产线从 mask 到缺陷统计分割模型输出的是像素 mask产线要的往往是缺陷记录位置、面积、类别、数量。把 mask 转成结构化数据需要一次连通域分析每个连通域算面积和中心坐标按面积阈值判断是否真的需要报警。这一步可以把小噪声滤掉同时生成报表。另外部署时如果只是跑 PyTorch 推理工控机上要装全套 Python 环境现场实施成本高。如果检测节拍允许、工控机性能一般我一般会把模型导出成 ONNX 再转 TensorRT输入输出尺寸固定成 1x1x200x200这样推理单张稳定在 5-10ms且不依赖 Python 环境。导出时的细节是 opset 版本要设在 12 以上BatchNorm 和 ReLU 的融合交给 TensorRT 自己处理。产线项目里「能跑」和「好用」的差距就在这些地方模型训练出来只算完成了三分之一后面还有数据回流、误报率监控、定期增量训练这些持续工作要做。我自己带项目时最大的教训是别急着上最好的模型先把推理链路打通用最朴素的 UNet 做一轮端到端验证让现场人员看到真实效果再逐步迭代。这样即使模型改进走弯路产线那边也知道你在解决什么问题。希望这些经验和踩过的坑能帮你在钢材表面缺陷检测这条路上少走几步弯路。本文还有配套的精品资源点击获取