简介本资源面向工业质检方向的深度学习学习者与算法工程师提供一套基于UNet的钢材表面缺陷检测完整项目实战。项目以NEU-DET数据集为基础覆盖划痕、凹坑、夹杂等多类缺陷样本可用于图像分割模型的训练、验证与部署练习适合具备一定Python与深度学习基础、希望积累工业缺陷检测经验的中高级读者。压缩包共约2000个文件整体122.16MB其中1800个xml标注文件与180个jpg缺陷图像构成数据集主体另有8个py脚本、1个ipynb探索性分析笔记、1个onnx导出模型及若干md、txt说明文档便于从数据预处理、模型训练到推理部署全流程复现。目前已有889人学习下载。通过该资源读者可掌握UNet编码器-解码器与跳跃连接在缺陷分割中的落地方式理解数据集划分、损失函数与优化器选择、模型评估等关键环节并借助onnx模型与脚本快速验证推理效果为产线级表面缺陷检测系统开发提供可参考的工程范例。1. 钢材表面缺陷检测为什么值得用 UNet 啃一遍钢厂质检线上一块热轧板以每秒数米的速度掠过相机氧化铁皮、辊印、裂纹、夹杂这些缺陷在图像里往往只占几十个像素对比度还低。传统 OpenCV 那套阈值加形态学的做法遇到光照波动和材质反光就集体翻车漏检率压不下去。UNet 这类编码器-解码器结构之所以在钢材表面缺陷检测里被反复拿出来讲核心就一条它能把浅层的高分辨率纹理和深层的语义信息用跳跃连接拼回去小目标缺陷的边界不至于在降采样里被抹掉。NEU-DET 这个数据集正好提供了热轧板六类缺陷的灰度图规模不大、标注干净是验证 UNet 分割方案能不能跑通、值不值得往产线迁移的最低成本试验田。这篇面向的是想动手复现一套完整缺陷检测流程的工程师从数据加载、模型搭建、训练调参到推理部署每一步都给可抄的命令和参数新手能跟着跑熟手能直接看到边界在哪。2. NEU-DET 数据集的读取与缺陷掩码对齐2.1 先搞清楚 NEU-DET 的目录结构和标注形式NEU-DET 的常见组织方式是按缺陷类别分文件夹每个类别下是若干张 200×200 的灰度图标注以 XML 或掩码图的形式给出。热轧板六类缺陷分别是裂纹crazing、夹杂inclusion、斑块patches、点蚀pitted_surface、氧化铁皮压入rolled-in_scale、划痕scratches。很多公开版本只给了分类标签没有像素级掩码这时候要么找带掩码的版本要么自己用标注工具补。我一般会先写个脚本把目录扫一遍确认图像数量和标注是否一一对应避免训练到一半才发现有图没标签。import os from pathlib import Path def scan_neu_det(root): 扫描 NEU-DET 目录返回每类的图像与标注配对情况 root Path(root) stats {} for cls_dir in sorted(root.iterdir()): if not cls_dir.is_dir(): continue imgs list(cls_dir.glob(*.jpg)) list(cls_dir.glob(*.png)) masks list(cls_dir.glob(*.xml)) list(cls_dir.glob(*_mask.png)) stats[cls_dir.name] {images: len(imgs), masks: len(masks)} return stats if __name__ __main__: print(scan_neu_det(./NEU-DET/train/images))这段脚本只做一件事把每个类别文件夹下的图像数和标注数打出来。参数上注意glob的模式要覆盖你实际拿到的后缀有的版本图像是.jpg掩码是.png有的全是.bmp。如果图像数和标注数对不上先别急着训练去查是不是有图被误删或者标注命名不一致。这一步花十分钟能省掉后面几小时的无效训练。2.2 把 XML 标注转成 UNet 能吃的二值掩码UNet 做分割输出通道数等于类别数加背景。NEU-DET 如果只有分类标签就得先把边界框或标注点转成掩码。常见做法是用 XML 里的bndbox生成矩形掩码虽然粗糙但作为基线够用。更精细的做法是用多边形标注这里以矩形为例。import xml.etree.ElementTree as ET import numpy as np import cv2 def xml_to_mask(xml_path, img_size(200, 200)): 把单个 XML 的 bndbox 转成二值掩码缺陷区域为 1 tree ET.parse(xml_path) root tree.getroot() mask np.zeros(img_size, dtypenp.uint8) for obj in root.findall(object): bbox obj.find(bndbox) xmin int(float(bbox.find(xmin).text)) ymin int(float(bbox.find(ymin).text)) xmax int(float(bbox.find(xmax).text)) ymax int(float(bbox.find(ymax).text)) mask[ymin:ymax, xmin:xmax] 1 return mask # 批量转换并保存 for xml_file in Path(./NEU-DET/train/annotations).glob(*.xml): m xml_to_mask(str(xml_file)) cv2.imwrite(str(xml_file).replace(.xml, _mask.png), m * 255)逻辑上xml_to_mask遍历所有object把每个边界框区域置 1其余为 0。参数img_size必须和实际图像尺寸一致NEU-DET 常见是 200×200如果你做了缩放这里也要同步改。保存时乘 255 是为了让掩码在可视化时能看清训练时再除以 255 归一化。注意矩形掩码会引入大量背景像素如果缺陷本身是细长裂纹矩形框里大部分是正常区域直接训练会让模型学到“框内即缺陷”的错误先验这时候要么用更精细的标注要么在损失函数里对边界做加权。2.3 训练集与验证集的划分要按缺陷类别分层NEU-DET 每类样本数量不均衡随机划分可能导致某类缺陷在验证集里一张都没有。我一般用分层抽样保证每个类别在训练集和验证集里的比例接近。from sklearn.model_selection import train_test_split import numpy as np def split_dataset(image_paths, labels, test_size0.2, seed42): 按类别分层划分返回训练和验证索引 idx np.arange(len(image_paths)) train_idx, val_idx train_test_split( idx, test_sizetest_size, stratifylabels, random_stateseed ) return train_idx, val_idxstratifylabels是关键参数它保证划分后各类别比例一致。random_state固定后结果可复现方便对比不同模型。如果某类样本少于 10 张分层抽样可能报错这时候要么合并稀有类要么改用 K 折交叉验证。划分完记得把图像路径和掩码路径一起存成列表训练时按索引取别在 DataLoader 里再做路径拼接那样容易出错。3. UNet 模型搭建与钢材缺陷分割的训练配置3.1 编码器-解码器结构里跳跃连接为什么不能省UNet 的核心是下采样提语义、上采样恢复分辨率跳跃连接把编码器每层的特征直接拼到解码器对应层。钢材表面缺陷里裂纹可能只有几个像素宽经过四次下采样后特征图只剩原图的十六分之一没有跳跃连接这些细节根本回不来。我见过有人为了省显存把跳跃连接砍掉结果 IoU 直接掉十几个点这就是血泪经验。编码器一般用 VGG 或 ResNet 的前几层解码器用转置卷积或双线性插值加上卷积。输出层用 1×1 卷积把通道数压到类别数二分类就是 1 个通道加 Sigmoid。import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.net 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.net(x) class UNet(nn.Module): def __init__(self, in_ch1, out_ch1, base32): super().__init__() self.enc1 DoubleConv(in_ch, base) self.enc2 DoubleConv(base, base*2) self.enc3 DoubleConv(base*2, base*4) self.enc4 DoubleConv(base*4, base*8) self.pool nn.MaxPool2d(2) self.bottleneck DoubleConv(base*8, base*16) self.up4 nn.ConvTranspose2d(base*16, base*8, 2, stride2) self.dec4 DoubleConv(base*16, base*8) self.up3 nn.ConvTranspose2d(base*8, base*4, 2, stride2) self.dec3 DoubleConv(base*8, base*4) self.up2 nn.ConvTranspose2d(base*4, base*2, 2, stride2) self.dec2 DoubleConv(base*4, base*2) self.up1 nn.ConvTranspose2d(base*2, base, 2, stride2) self.dec1 DoubleConv(base*2, base) self.out nn.Conv2d(base, out_ch, 1) def forward(self, x): e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) e4 self.enc4(self.pool(e3)) b self.bottleneck(self.pool(e4)) d4 self.dec4(torch.cat([self.up4(b), e4], dim1)) d3 self.dec3(torch.cat([self.up3(d4), e3], dim1)) d2 self.dec2(torch.cat([self.up2(d3), e2], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e1], dim1)) return self.out(d1)base32控制第一层通道数显存不够就降到 16但别低于 8否则特征表达能力太弱。in_ch1对应灰度图如果输入是 RGB 就改成 3。out_ch1配 Sigmoid 做二分类多分类就改成类别数加背景。torch.cat的维度是 1也就是通道维拼接前要确保上采样后的尺寸和编码器对应层一致否则会报尺寸不匹配。我一般会在拼接前打印一下 shape确认无误再往下跑。3.2 损失函数选 Dice 还是 BCE钢材缺陷场景怎么定钢材缺陷里正负样本极度不均衡一张 200×200 的图里缺陷可能只占几百个像素BCE 会被背景主导模型倾向于全预测背景。Dice Loss 直接优化重叠度对小目标更友好。常见做法是 BCE 和 Dice 按权重相加比如loss 0.5 * BCE 0.5 * Dice。如果缺陷特别细还可以加 Focal Loss 压制易分样本。我一般先用纯 Dice 跑一版基线看 IoU 能不能过 0.5不行再调权重。class DiceLoss(nn.Module): def __init__(self, smooth1.0): super().__init__() self.smooth smooth def forward(self, pred, target): pred torch.sigmoid(pred) pred pred.view(-1) target target.view(-1) inter (pred * target).sum() dice (2 * inter self.smooth) / (pred.sum() target.sum() self.smooth) return 1 - dice criterion lambda p, t: 0.5 * nn.BCEWithLogitsLoss()(p, t) 0.5 * DiceLoss()(p, t)smooth1.0防止分母为零别设太小否则数值不稳定。BCEWithLogitsLoss内部带 Sigmoid所以模型输出不要再加激活。如果验证集 IoU 震荡厉害先把 Dice 权重降到 0.3 试试有时候纯 Dice 在训练初期梯度太猛反而收敛不稳。3.3 训练循环里学习率和批大小的实操取值NEU-DET 图像小批大小可以设 16 或 32显存不够就降到 8但别低于 4否则 BatchNorm 统计量不准。学习率用 1e-3 配 Adam或者 1e-2 配 SGD 加动量。我一般先用 Adam 1e-3 跑 50 个 epoch看验证集 IoU 曲线如果前期涨得快后期平就加余弦退火。数据增强用随机翻转、旋转 90 度、亮度微调别用太激进的裁剪钢材缺陷位置本身就有随机性裁多了可能把缺陷裁没。from torch.utils.data import DataLoader from torch.optim import Adam from torch.optim.lr_scheduler import CosineAnnealingLR model UNet(in_ch1, out_ch1, base32).cuda() optimizer Adam(model.parameters(), lr1e-3) scheduler CosineAnnealingLR(optimizer, T_max50) loader DataLoader(dataset, batch_size16, shuffleTrue, num_workers4) for epoch in range(50): model.train() for img, mask in loader: img, mask img.cuda(), mask.cuda() optimizer.zero_grad() pred model(img) loss criterion(pred, mask) loss.backward() optimizer.step() scheduler.step() # 每个 epoch 后在验证集上算 IoUnum_workers4根据 CPU 核数调太多反而拖慢。T_max50和总 epoch 一致余弦退火从 1e-3 降到接近 0。验证集 IoU 每 5 个 epoch 记一次别每轮都算浪费时间。如果 loss 不降先检查掩码是不是全黑或者全白这种低级错误我见过不止一次。4. 推理、后处理与产线迁移的避坑清单4.1 推理时阈值怎么选别用默认 0.5训练完模型输出的是 logitsSigmoid 后得到概率图。默认 0.5 做阈值在缺陷检测里往往偏高因为缺陷边缘概率低0.5 会把边缘切掉导致分割区域偏小。我一般会在验证集上扫一遍阈值从 0.3 到 0.7看哪个 IoU 最高。钢材缺陷里裂纹这种细长目标阈值可能要降到 0.3 才能连成线。def evaluate_threshold(model, loader, thresholds): model.eval() results {} with torch.no_grad(): for th in thresholds: ious [] for img, mask in loader: img img.cuda() pred torch.sigmoid(model(img)).cpu().numpy() pred_bin (pred th).astype(np.uint8) # 计算 IoU inter (pred_bin mask.numpy()).sum() union (pred_bin | mask.numpy()).sum() ious.append(inter / (union 1e-6)) results[th] np.mean(ious) return resultsthresholds建议取[0.3, 0.4, 0.5, 0.6, 0.7]跑完看哪个最高。注意验证集要和训练集同分布别拿训练集调阈值那样会过拟合。如果不同类别最优阈值差很多可以考虑按类别分别设阈值但产线上一般统一一个值方便维护。4.2 后处理去噪连通域和形态学怎么配合模型输出难免有零星噪点尤其是背景纹理复杂的地方。常见做法是先做连通域分析去掉面积小于某个阈值的区域再用闭运算把断裂的缺陷连起来。面积阈值根据实际缺陷最小尺寸定NEU-DET 里裂纹可能只有几十个像素阈值设 20 左右。闭运算核大小 3×3 或 5×5太大反而会把相邻缺陷粘在一起。import cv2 def postprocess(pred_mask, min_area20, kernel_size3): pred_mask pred_mask.astype(np.uint8) num, labels, stats, _ cv2.connectedComponentsWithStats(pred_mask, connectivity8) clean np.zeros_like(pred_mask) for i in range(1, num): if stats[i, cv2.CC_STAT_AREA] min_area: clean[labels i] 1 kernel cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) clean cv2.morphologyEx(clean, cv2.MORPH_CLOSE, kernel) return cleanconnectivity8比 4 连通更容易把斜向裂纹连起来。min_area别设太大否则小缺陷直接被滤掉。闭运算放在连通域之后顺序反了会把噪点也连成块。这套后处理在 NEU-DET 上能把误检降三成左右但参数得根据实际图像调没有万能值。4.3 避坑清单钢材缺陷检测里最容易翻车的五个点现象一训练 loss 正常下降但验证集 IoU 始终低于 0.3。原因多半是掩码和图像没对齐比如图像做了归一化而掩码没同步或者路径拼接时把_mask后缀搞错了。解决方法是抽一张图把原图和掩码叠在一起可视化肉眼确认缺陷位置是否重合。现象二模型把所有像素都预测成背景IoU 为 0。这是正负样本极度不均衡的典型表现BCE 被背景主导。解决方法是换 Dice Loss 或者给 BCE 加pos_weight参数把正样本权重调高。pos_weight一般设成负样本数除以正样本数NEU-DET 里这个值可能到几十。现象三推理时显存爆了但训练时没事。推理时忘了加torch.no_grad()或者输入尺寸比训练时大。解决方法是推理代码包在with torch.no_grad():里并且确保输入 resize 到和训练一致的尺寸。如果产线图像分辨率高可以切块推理再拼接别直接整图塞进去。现象四不同批次图像亮度差异大模型表现忽好忽坏。钢材表面反光受光照影响大训练集覆盖不足就会这样。解决方法是在数据增强里加随机亮度、对比度扰动或者用直方图均衡化做预处理。更彻底的做法是加一个自适应归一化层但会增参数量。现象五后处理把细裂纹滤掉了。min_area设太大或者闭运算核太大把裂纹和背景粘在一起。解决方法是先把min_area降到 10 以下闭运算核用 3×3然后逐步往上调每调一次看可视化结果。别一次性设个大值那样调不回来。5. 把 NEU-DET 上的 UNet 迁移到自建数据集的关键动作NEU-DET 跑通只是起点真正值钱的是把这套流程迁到自家产线的图像上。我一般会先做三件事第一用训练好的模型在自建数据上跑推理看哪些缺陷类型漏检最多这能告诉你标注该往哪个方向补第二把自建数据里和 NEU-DET 差异大的样本挑出来比如不同光照、不同材质单独微调最后两层解码器学习率设成 1e-4别动编码器第三如果自建数据标注量少用 NEU-DET 预训练权重做初始化比从头训收敛快一倍不止。迁移时有个参数特别关键输入归一化的均值和方差。NEU-DET 是灰度图均值和方差跟你的相机输出可能差很远。我习惯在自建数据上重新统计一遍把mean和std更新到 DataLoader 的 transform 里。别小看这一步我见过有人直接套 ImageNet 的均值方差结果模型前几个 epoch 完全不收敛查了两天才发现是归一化不匹配。验证迁移效果不能只看 IoU产线上更关心漏检率和过杀率。漏检率是缺陷被预测成背景的比例过杀率是正常区域被预测成缺陷的比例。这两个指标要分开算因为产线对漏检的容忍度远低于过杀。我一般会画一条 PR 曲线看在不同阈值下漏检和过杀的权衡然后选一个业务能接受的阈值。如果漏检率压不下去优先查标注质量其次查模型容量最后才考虑换更复杂的结构。最后说个习惯每次训完模型把权重、配置文件、验证集 IoU 和阈值扫描结果存一个文件夹命名带上日期和关键参数。我吃过亏三个月后想复现某个结果发现权重还在但配置忘了只能重训。现在不管多小的实验我都存一份硬盘不值钱时间值钱。希望帮到你。本文还有配套的精品资源点击获取
简介本资源面向工业质检方向的深度学习学习者与算法工程师提供一套基于UNet的钢材表面缺陷检测完整项目实战。项目以NEU-DET数据集为基础覆盖划痕、凹坑、夹杂等多类缺陷样本可用于图像分割模型的训练、验证与部署练习适合具备一定Python与深度学习基础、希望积累工业缺陷检测经验的中高级读者。压缩包共约2000个文件整体122.16MB其中1800个xml标注文件与180个jpg缺陷图像构成数据集主体另有8个py脚本、1个ipynb探索性分析笔记、1个onnx导出模型及若干md、txt说明文档便于从数据预处理、模型训练到推理部署全流程复现。目前已有889人学习下载。通过该资源读者可掌握UNet编码器-解码器与跳跃连接在缺陷分割中的落地方式理解数据集划分、损失函数与优化器选择、模型评估等关键环节并借助onnx模型与脚本快速验证推理效果为产线级表面缺陷检测系统开发提供可参考的工程范例。1. 钢材表面缺陷检测为什么值得用 UNet 啃一遍钢厂质检线上一块热轧板以每秒数米的速度掠过相机氧化铁皮、辊印、裂纹、夹杂这些缺陷在图像里往往只占几十个像素对比度还低。传统 OpenCV 那套阈值加形态学的做法遇到光照波动和材质反光就集体翻车漏检率压不下去。UNet 这类编码器-解码器结构之所以在钢材表面缺陷检测里被反复拿出来讲核心就一条它能把浅层的高分辨率纹理和深层的语义信息用跳跃连接拼回去小目标缺陷的边界不至于在降采样里被抹掉。NEU-DET 这个数据集正好提供了热轧板六类缺陷的灰度图规模不大、标注干净是验证 UNet 分割方案能不能跑通、值不值得往产线迁移的最低成本试验田。这篇面向的是想动手复现一套完整缺陷检测流程的工程师从数据加载、模型搭建、训练调参到推理部署每一步都给可抄的命令和参数新手能跟着跑熟手能直接看到边界在哪。2. NEU-DET 数据集的读取与缺陷掩码对齐2.1 先搞清楚 NEU-DET 的目录结构和标注形式NEU-DET 的常见组织方式是按缺陷类别分文件夹每个类别下是若干张 200×200 的灰度图标注以 XML 或掩码图的形式给出。热轧板六类缺陷分别是裂纹crazing、夹杂inclusion、斑块patches、点蚀pitted_surface、氧化铁皮压入rolled-in_scale、划痕scratches。很多公开版本只给了分类标签没有像素级掩码这时候要么找带掩码的版本要么自己用标注工具补。我一般会先写个脚本把目录扫一遍确认图像数量和标注是否一一对应避免训练到一半才发现有图没标签。import os from pathlib import Path def scan_neu_det(root): 扫描 NEU-DET 目录返回每类的图像与标注配对情况 root Path(root) stats {} for cls_dir in sorted(root.iterdir()): if not cls_dir.is_dir(): continue imgs list(cls_dir.glob(*.jpg)) list(cls_dir.glob(*.png)) masks list(cls_dir.glob(*.xml)) list(cls_dir.glob(*_mask.png)) stats[cls_dir.name] {images: len(imgs), masks: len(masks)} return stats if __name__ __main__: print(scan_neu_det(./NEU-DET/train/images))这段脚本只做一件事把每个类别文件夹下的图像数和标注数打出来。参数上注意glob的模式要覆盖你实际拿到的后缀有的版本图像是.jpg掩码是.png有的全是.bmp。如果图像数和标注数对不上先别急着训练去查是不是有图被误删或者标注命名不一致。这一步花十分钟能省掉后面几小时的无效训练。2.2 把 XML 标注转成 UNet 能吃的二值掩码UNet 做分割输出通道数等于类别数加背景。NEU-DET 如果只有分类标签就得先把边界框或标注点转成掩码。常见做法是用 XML 里的bndbox生成矩形掩码虽然粗糙但作为基线够用。更精细的做法是用多边形标注这里以矩形为例。import xml.etree.ElementTree as ET import numpy as np import cv2 def xml_to_mask(xml_path, img_size(200, 200)): 把单个 XML 的 bndbox 转成二值掩码缺陷区域为 1 tree ET.parse(xml_path) root tree.getroot() mask np.zeros(img_size, dtypenp.uint8) for obj in root.findall(object): bbox obj.find(bndbox) xmin int(float(bbox.find(xmin).text)) ymin int(float(bbox.find(ymin).text)) xmax int(float(bbox.find(xmax).text)) ymax int(float(bbox.find(ymax).text)) mask[ymin:ymax, xmin:xmax] 1 return mask # 批量转换并保存 for xml_file in Path(./NEU-DET/train/annotations).glob(*.xml): m xml_to_mask(str(xml_file)) cv2.imwrite(str(xml_file).replace(.xml, _mask.png), m * 255)逻辑上xml_to_mask遍历所有object把每个边界框区域置 1其余为 0。参数img_size必须和实际图像尺寸一致NEU-DET 常见是 200×200如果你做了缩放这里也要同步改。保存时乘 255 是为了让掩码在可视化时能看清训练时再除以 255 归一化。注意矩形掩码会引入大量背景像素如果缺陷本身是细长裂纹矩形框里大部分是正常区域直接训练会让模型学到“框内即缺陷”的错误先验这时候要么用更精细的标注要么在损失函数里对边界做加权。2.3 训练集与验证集的划分要按缺陷类别分层NEU-DET 每类样本数量不均衡随机划分可能导致某类缺陷在验证集里一张都没有。我一般用分层抽样保证每个类别在训练集和验证集里的比例接近。from sklearn.model_selection import train_test_split import numpy as np def split_dataset(image_paths, labels, test_size0.2, seed42): 按类别分层划分返回训练和验证索引 idx np.arange(len(image_paths)) train_idx, val_idx train_test_split( idx, test_sizetest_size, stratifylabels, random_stateseed ) return train_idx, val_idxstratifylabels是关键参数它保证划分后各类别比例一致。random_state固定后结果可复现方便对比不同模型。如果某类样本少于 10 张分层抽样可能报错这时候要么合并稀有类要么改用 K 折交叉验证。划分完记得把图像路径和掩码路径一起存成列表训练时按索引取别在 DataLoader 里再做路径拼接那样容易出错。3. UNet 模型搭建与钢材缺陷分割的训练配置3.1 编码器-解码器结构里跳跃连接为什么不能省UNet 的核心是下采样提语义、上采样恢复分辨率跳跃连接把编码器每层的特征直接拼到解码器对应层。钢材表面缺陷里裂纹可能只有几个像素宽经过四次下采样后特征图只剩原图的十六分之一没有跳跃连接这些细节根本回不来。我见过有人为了省显存把跳跃连接砍掉结果 IoU 直接掉十几个点这就是血泪经验。编码器一般用 VGG 或 ResNet 的前几层解码器用转置卷积或双线性插值加上卷积。输出层用 1×1 卷积把通道数压到类别数二分类就是 1 个通道加 Sigmoid。import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.net 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.net(x) class UNet(nn.Module): def __init__(self, in_ch1, out_ch1, base32): super().__init__() self.enc1 DoubleConv(in_ch, base) self.enc2 DoubleConv(base, base*2) self.enc3 DoubleConv(base*2, base*4) self.enc4 DoubleConv(base*4, base*8) self.pool nn.MaxPool2d(2) self.bottleneck DoubleConv(base*8, base*16) self.up4 nn.ConvTranspose2d(base*16, base*8, 2, stride2) self.dec4 DoubleConv(base*16, base*8) self.up3 nn.ConvTranspose2d(base*8, base*4, 2, stride2) self.dec3 DoubleConv(base*8, base*4) self.up2 nn.ConvTranspose2d(base*4, base*2, 2, stride2) self.dec2 DoubleConv(base*4, base*2) self.up1 nn.ConvTranspose2d(base*2, base, 2, stride2) self.dec1 DoubleConv(base*2, base) self.out nn.Conv2d(base, out_ch, 1) def forward(self, x): e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) e4 self.enc4(self.pool(e3)) b self.bottleneck(self.pool(e4)) d4 self.dec4(torch.cat([self.up4(b), e4], dim1)) d3 self.dec3(torch.cat([self.up3(d4), e3], dim1)) d2 self.dec2(torch.cat([self.up2(d3), e2], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e1], dim1)) return self.out(d1)base32控制第一层通道数显存不够就降到 16但别低于 8否则特征表达能力太弱。in_ch1对应灰度图如果输入是 RGB 就改成 3。out_ch1配 Sigmoid 做二分类多分类就改成类别数加背景。torch.cat的维度是 1也就是通道维拼接前要确保上采样后的尺寸和编码器对应层一致否则会报尺寸不匹配。我一般会在拼接前打印一下 shape确认无误再往下跑。3.2 损失函数选 Dice 还是 BCE钢材缺陷场景怎么定钢材缺陷里正负样本极度不均衡一张 200×200 的图里缺陷可能只占几百个像素BCE 会被背景主导模型倾向于全预测背景。Dice Loss 直接优化重叠度对小目标更友好。常见做法是 BCE 和 Dice 按权重相加比如loss 0.5 * BCE 0.5 * Dice。如果缺陷特别细还可以加 Focal Loss 压制易分样本。我一般先用纯 Dice 跑一版基线看 IoU 能不能过 0.5不行再调权重。class DiceLoss(nn.Module): def __init__(self, smooth1.0): super().__init__() self.smooth smooth def forward(self, pred, target): pred torch.sigmoid(pred) pred pred.view(-1) target target.view(-1) inter (pred * target).sum() dice (2 * inter self.smooth) / (pred.sum() target.sum() self.smooth) return 1 - dice criterion lambda p, t: 0.5 * nn.BCEWithLogitsLoss()(p, t) 0.5 * DiceLoss()(p, t)smooth1.0防止分母为零别设太小否则数值不稳定。BCEWithLogitsLoss内部带 Sigmoid所以模型输出不要再加激活。如果验证集 IoU 震荡厉害先把 Dice 权重降到 0.3 试试有时候纯 Dice 在训练初期梯度太猛反而收敛不稳。3.3 训练循环里学习率和批大小的实操取值NEU-DET 图像小批大小可以设 16 或 32显存不够就降到 8但别低于 4否则 BatchNorm 统计量不准。学习率用 1e-3 配 Adam或者 1e-2 配 SGD 加动量。我一般先用 Adam 1e-3 跑 50 个 epoch看验证集 IoU 曲线如果前期涨得快后期平就加余弦退火。数据增强用随机翻转、旋转 90 度、亮度微调别用太激进的裁剪钢材缺陷位置本身就有随机性裁多了可能把缺陷裁没。from torch.utils.data import DataLoader from torch.optim import Adam from torch.optim.lr_scheduler import CosineAnnealingLR model UNet(in_ch1, out_ch1, base32).cuda() optimizer Adam(model.parameters(), lr1e-3) scheduler CosineAnnealingLR(optimizer, T_max50) loader DataLoader(dataset, batch_size16, shuffleTrue, num_workers4) for epoch in range(50): model.train() for img, mask in loader: img, mask img.cuda(), mask.cuda() optimizer.zero_grad() pred model(img) loss criterion(pred, mask) loss.backward() optimizer.step() scheduler.step() # 每个 epoch 后在验证集上算 IoUnum_workers4根据 CPU 核数调太多反而拖慢。T_max50和总 epoch 一致余弦退火从 1e-3 降到接近 0。验证集 IoU 每 5 个 epoch 记一次别每轮都算浪费时间。如果 loss 不降先检查掩码是不是全黑或者全白这种低级错误我见过不止一次。4. 推理、后处理与产线迁移的避坑清单4.1 推理时阈值怎么选别用默认 0.5训练完模型输出的是 logitsSigmoid 后得到概率图。默认 0.5 做阈值在缺陷检测里往往偏高因为缺陷边缘概率低0.5 会把边缘切掉导致分割区域偏小。我一般会在验证集上扫一遍阈值从 0.3 到 0.7看哪个 IoU 最高。钢材缺陷里裂纹这种细长目标阈值可能要降到 0.3 才能连成线。def evaluate_threshold(model, loader, thresholds): model.eval() results {} with torch.no_grad(): for th in thresholds: ious [] for img, mask in loader: img img.cuda() pred torch.sigmoid(model(img)).cpu().numpy() pred_bin (pred th).astype(np.uint8) # 计算 IoU inter (pred_bin mask.numpy()).sum() union (pred_bin | mask.numpy()).sum() ious.append(inter / (union 1e-6)) results[th] np.mean(ious) return resultsthresholds建议取[0.3, 0.4, 0.5, 0.6, 0.7]跑完看哪个最高。注意验证集要和训练集同分布别拿训练集调阈值那样会过拟合。如果不同类别最优阈值差很多可以考虑按类别分别设阈值但产线上一般统一一个值方便维护。4.2 后处理去噪连通域和形态学怎么配合模型输出难免有零星噪点尤其是背景纹理复杂的地方。常见做法是先做连通域分析去掉面积小于某个阈值的区域再用闭运算把断裂的缺陷连起来。面积阈值根据实际缺陷最小尺寸定NEU-DET 里裂纹可能只有几十个像素阈值设 20 左右。闭运算核大小 3×3 或 5×5太大反而会把相邻缺陷粘在一起。import cv2 def postprocess(pred_mask, min_area20, kernel_size3): pred_mask pred_mask.astype(np.uint8) num, labels, stats, _ cv2.connectedComponentsWithStats(pred_mask, connectivity8) clean np.zeros_like(pred_mask) for i in range(1, num): if stats[i, cv2.CC_STAT_AREA] min_area: clean[labels i] 1 kernel cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) clean cv2.morphologyEx(clean, cv2.MORPH_CLOSE, kernel) return cleanconnectivity8比 4 连通更容易把斜向裂纹连起来。min_area别设太大否则小缺陷直接被滤掉。闭运算放在连通域之后顺序反了会把噪点也连成块。这套后处理在 NEU-DET 上能把误检降三成左右但参数得根据实际图像调没有万能值。4.3 避坑清单钢材缺陷检测里最容易翻车的五个点现象一训练 loss 正常下降但验证集 IoU 始终低于 0.3。原因多半是掩码和图像没对齐比如图像做了归一化而掩码没同步或者路径拼接时把_mask后缀搞错了。解决方法是抽一张图把原图和掩码叠在一起可视化肉眼确认缺陷位置是否重合。现象二模型把所有像素都预测成背景IoU 为 0。这是正负样本极度不均衡的典型表现BCE 被背景主导。解决方法是换 Dice Loss 或者给 BCE 加pos_weight参数把正样本权重调高。pos_weight一般设成负样本数除以正样本数NEU-DET 里这个值可能到几十。现象三推理时显存爆了但训练时没事。推理时忘了加torch.no_grad()或者输入尺寸比训练时大。解决方法是推理代码包在with torch.no_grad():里并且确保输入 resize 到和训练一致的尺寸。如果产线图像分辨率高可以切块推理再拼接别直接整图塞进去。现象四不同批次图像亮度差异大模型表现忽好忽坏。钢材表面反光受光照影响大训练集覆盖不足就会这样。解决方法是在数据增强里加随机亮度、对比度扰动或者用直方图均衡化做预处理。更彻底的做法是加一个自适应归一化层但会增参数量。现象五后处理把细裂纹滤掉了。min_area设太大或者闭运算核太大把裂纹和背景粘在一起。解决方法是先把min_area降到 10 以下闭运算核用 3×3然后逐步往上调每调一次看可视化结果。别一次性设个大值那样调不回来。5. 把 NEU-DET 上的 UNet 迁移到自建数据集的关键动作NEU-DET 跑通只是起点真正值钱的是把这套流程迁到自家产线的图像上。我一般会先做三件事第一用训练好的模型在自建数据上跑推理看哪些缺陷类型漏检最多这能告诉你标注该往哪个方向补第二把自建数据里和 NEU-DET 差异大的样本挑出来比如不同光照、不同材质单独微调最后两层解码器学习率设成 1e-4别动编码器第三如果自建数据标注量少用 NEU-DET 预训练权重做初始化比从头训收敛快一倍不止。迁移时有个参数特别关键输入归一化的均值和方差。NEU-DET 是灰度图均值和方差跟你的相机输出可能差很远。我习惯在自建数据上重新统计一遍把mean和std更新到 DataLoader 的 transform 里。别小看这一步我见过有人直接套 ImageNet 的均值方差结果模型前几个 epoch 完全不收敛查了两天才发现是归一化不匹配。验证迁移效果不能只看 IoU产线上更关心漏检率和过杀率。漏检率是缺陷被预测成背景的比例过杀率是正常区域被预测成缺陷的比例。这两个指标要分开算因为产线对漏检的容忍度远低于过杀。我一般会画一条 PR 曲线看在不同阈值下漏检和过杀的权衡然后选一个业务能接受的阈值。如果漏检率压不下去优先查标注质量其次查模型容量最后才考虑换更复杂的结构。最后说个习惯每次训完模型把权重、配置文件、验证集 IoU 和阈值扫描结果存一个文件夹命名带上日期和关键参数。我吃过亏三个月后想复现某个结果发现权重还在但配置忘了只能重训。现在不管多小的实验我都存一份硬盘不值钱时间值钱。希望帮到你。本文还有配套的精品资源点击获取
