简介面向钢材表面缺陷检测场景此项目基于UNet架构与NEU-DET数据集构建完整实战方案适合工业视觉质检人员、深度学习初学者及计算机视觉方向学习者参考。包内共2000个文件包含1800个XML标注文件、180张钢材表面缺陷图像、8个Python脚本及配套Pyc另有Jupyter Notebook分析与ONNX模型整体约122.16MB。数据集中涵盖划痕、凹坑、夹杂等多类缺陷项目从数据预处理、模型训练到部署推理均有清晰代码与说明文档可直接复现或迁移。目前已有891人学习浏览项目不仅提供训练推理流程还包含数据划分策略、EDA分析笔记与最终模型文件便于用户理解UNet在缺陷分割任务中的实际应用并快速开展二次开发。1. 钢材缺陷检测实战为什么 UNet NEU-DET 这套组合最稳钢材表面缺陷检测放在产线上本质是“肉眼盯传送带”的升级版人工质检员看久了会疲劳漏检率随班次递增而一张 200×200 的钢材表面图像里裂纹、麻点、划痕这类缺陷常常只占几十个像素肉眼很容易漏掉。UNet 的编码器-解码器结构配合跳跃连接天然适合做像素级分割能直接输出缺陷的精确位置和形状而不是只给一个“有/无”的粗结果。NEU-DET 数据集覆盖了六类常见钢材表面缺陷是训练分割模型最常用的公开数据之一。这个项目从 EDA 分析、数据预处理到 UNet 训练和推理验证全流程都有适合刚接触分割任务、想在一个真实工业数据集上跑通完整链路的人。下面按实际执行顺序把每一步的参数设置和踩坑点摊开讲。2. NEU-DET 数据集与预处理六类缺陷样本怎么变成训练张量2.1 数据集目录结构与标注格式NEU-DET 由东北大学发布原始图像统一为 200×200 像素的灰度图共 1800 张六类缺陷每类 300 张。类别官方命名是 crazing裂纹、inclusion夹杂、patches麻点、pitted_surface氧化铁皮压入、rolled-in_scale轧制氧化皮、scratches划痕。整个压缩包除了原始图像还带了 Exploratory_Data_Analysis.ipynb 和 Data_Separation.jpg前者是 EDA 脚本后者是数据划分的流程示意训练前把这两个文件先过一遍能省掉不少摸索时间。原版 NEU-DET 的标注是 XML 格式的边界框做 UNet 分割必须转成像素级掩码。常见做法是用 OpenCV 解析 XML 里的多边形或 bbox在空白图上填充白色作为缺陷前景。我一般会先写一个批量转换脚本把 XML 转成与图像同名的 PNG 掩码掩码中缺陷像素为 255、背景为 0方便后续 PyTorch Dataset 直接读取也不用每次训练前重复解析。import cv2 import numpy as np import xml.etree.ElementTree as ET def xml_to_mask(xml_path, img_size(200, 200)): # 读原图拿到宽高掩码尺寸与图像严格一致 img cv2.imread(xml_path.replace(.xml, .jpg), cv2.IMREAD_GRAYSCALE) h, w img.shape mask np.zeros((h, w), dtypenp.uint8) tree ET.parse(xml_path) root tree.getroot() for obj in root.iter(object): pts [] for poly in obj.iter(polygon): for point in poly.iter(point): pts.append([int(float(point.attrib[x])), int(float(point.attrib[y]))]) if len(pts) 3: cv2.fillPoly(mask, [np.array(pts)], 255) return mask这段代码的核心是把 XML 里的多边形坐标还原成掩码。实际操作中要注意坐标单位有些版本存的是归一化坐标需要乘以宽高还原如果只有 bbox 没有 polygon就用 cv2.rectangle 填充但矩形掩码会把缺陷周围的正常纹理也标成前景给分割精度埋雷能用 polygon 尽量用 polygon。转换完成后建议随机抽 10 张图检查掩码叠加效果别等训练跑起来才发现标注错位。2.2 归一化与数据增强让模型适应产线光照变化钢材表面灰度图在不同光照、不同轧制批次下差异很大直接拿原始像素训练模型容易过拟合到当前数据集的亮度分布。我习惯先把图像归一化到 [0, 1]再配合数据增强。对小样本缺陷检测来说增强不是可选项而是刚需NEU-DET 每类才 300 张不做增强参数量都喂不饱。from torchvision import transforms train_transform transforms.Compose([ transforms.RandomHorizontalFlip(p0.5), transforms.RandomVerticalFlip(p0.5), transforms.RandomRotation(10), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485], std[0.229]) ])翻转和小角度旋转不改变缺陷语义适合钢材这种纹理方向性强的数据亮度和对比度扰动模拟产线不同工位的光照差异。注意 Normalize 的 mean 和 std 是单通道值NEU-DET 是灰度图如果套三通道 ImageNet 参数会直接报张量维度错误。掩码的增强必须和图像完全同步翻转旋转的随机种子要一致否则图像和标签错位指标看起来正常实际模型已经废了。一个血泪经验RandomRotation 默认会对图像做插值掩码也跟着旋转时要用最近邻插值双线性插值会让掩码边缘出现灰色过渡值二值化时边缘会多出一圈伪缺陷。torchvision 的 transforms 没法直接给 mask 单独指定插值方式常见做法是把图像和掩码拼接成多通道张量一起旋转或者干脆改用 Albumentations它对 image 和 mask 的同步处理更成熟。注意掩码增强必须与图像增强使用同一随机状态建议用 Albumentations 而不是手动组合 torchvision 变换否则极容易出现标签错位这种隐形故障。2.3 数据集划分训练/验证/测试的比例与随机种子数据划分影响模型泛化这个环节最容易翻车。NEU-DET 每类 300 张常见划分是 8:1:1即 1440 张训练、180 张验证、180 张测试。关键是按类别分层采样保证每类样本在三个集合里比例一致否则验证集 mIoU 会因为某类缺失而虚高或骤降。from sklearn.model_selection import train_test_split # filenames 是全部样本路径labels 是对应的类别索引 train_files, temp_files, train_labels, temp_labels train_test_split( filenames, labels, test_size0.2, stratifylabels, random_state42) val_files, test_files, val_labels, test_labels train_test_split( temp_files, temp_labels, test_size0.5, stratifytemp_labels, random_state42)stratify 参数保证划分后每个类别占比与原始数据一致random_state 固定 42 保证实验可复现。这里有个容易忽略的细节划分是在文件路径层面做的增强在训练时实时生成所以训练集的增强数据不会泄漏到验证和测试集。以前有人把增强后的图片写回磁盘再划分导致同一张原始图的不同增强版本同时出现在训练和验证集验证指标虚高一上产线就露馅。3. UNet 模型搭建与训练配置编码器、跳跃连接和损失函数怎么配3.1 UNet 核心结构为什么跳跃连接适合小目标缺陷UNet 名字来自论文里的 U 形结构图左侧编码器逐层下采样提取语义特征右侧解码器逐层上采样恢复空间分辨率中间的跳跃连接把编码器每层的特征图直接拼到解码器对应层。对缺陷检测来说这个结构最大的价值在于保留细节位置信息。钢材表面的裂纹可能只有几个像素宽单纯靠下采样后的高层语义特征根本找不回精确边界跳跃连接把浅层的边缘纹理信息直接送进解码器分割出来的缺陷边界才贴得准。与普通分类网络不同UNet 输出的是和输入同尺寸的分割图每个像素点一个预测值本质上是逐像素分类。NEU-DET 里的缺陷类别是互斥的一个像素要么是背景要么属于某一种缺陷所以输出通道数设为类别数加背景。这里有个常见误区六类缺陷加背景输出通道是 7 而不是 6很多人在这里少算了一路训练时 loss 维度对不上。import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv 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.conv(x)DoubleConv 是 UNet 的基本单元两次 3×3 卷积加 BatchNormpadding 保持特征图尺寸不变。BatchNorm 在 batch size 小于 8 时效果会打折扣如果显存只允许 batch size 设为 4建议换成 InstanceNorm 或者适当增大 batch size否则训练过程会很不稳定。这个模块在编码器和解码器里反复复用是整张网络参数量的主要来源。3.2 损失函数与优化器选择BCE、Dice 与 Adam 怎么组合缺陷分割是典型的类别不平衡问题缺陷像素占比通常在 5% 以下背景占绝对多数。如果直接用 BCE Loss模型会倾向把所有像素预测成背景因为这样 loss 也很低。常见做法是用 BCE Dice Loss 的加权组合Dice 系数对前景和背景一视同仁不会因为背景像素多就把小缺陷淹没。def bce_dice_loss(pred, mask, bce_weight0.5): # pred 是模型原始输出未经过 sigmoid bce nn.functional.binary_cross_entropy_with_logits(pred, mask) pred_prob torch.sigmoid(pred) smooth 1e-6 intersection (pred_prob * mask).sum() dice 1 - (2.0 * intersection smooth) / ( pred_prob.sum() mask.sum() smooth) return bce_weight * bce (1 - bce_weight) * dicebce_weight 一般取 0.5如果想更激进地押注前景召回可以把 Dice 的权重调高到 0.7。smooth 项防止除零数值取 1e-6 足够。注意这里用的是 binary_cross_entropy_with_logits要求模型最后一层不接 sigmoid直接在 loss 内部计算数值稳定性比手动 sigmoid 再算 BCE 好得多。优化器用 Adam初始学习率 1e-4 是比较稳妥的起点配合 ReduceLROnPlateau 在验证 loss 连续多个 epoch 不降时自动衰减学习率。3.3 训练超参数batch size、学习率与 epoch 的取舍NEU-DET 的 200×200 输入比较小显存压力不大。我一般用 batch size 16 起步如果训练过程中 loss 震荡剧烈就减半。epoch 数量参考验证集指标而不是硬性固定常见做法是设 100 个 epoch配合 early stopping 在验证 IoU 连续 15 个 epoch 不提升时终止训练保存最优权重。超参数推荐值说明输入尺寸200×200保持数据集原始分辨率batch size816看显存小于 8 注意 BN 稳定性初始学习率1e-4Adam 配合下足够收敛学习率衰减ReduceLROnPlateaupatience8factor0.5epoch100 early stoppingpatience15lossBCE Dice (0.5/0.5)缓解类别不平衡学习率是最影响结果的超参数。UNet 这种深网络用 1e-3 起步很容易在前期就发散收敛曲线像心跳一样上下跳用 1e-5 又收敛太慢100 个 epoch 跑完还在半山腰。我一般先用 1e-4 跑 20 个 epoch 看 loss 下降趋势不降就调到 1e-3 试一次降得太慢就降到 3e-5这个试探过程几行代码就能自动化别靠肉眼死等。4. 训练实战流程从 EDA 到训练曲线跟踪4.1 EDA 先行先看缺陷分布再定模型策略压缩包里的 Exploratory_Data_Analysis.ipynb 是标准的 EDA 流程。第一步统计每类缺陷的样本数和缺陷像素占比画出柱状图第二步随机挑几类缺陷的样本把原图和掩码叠加显示直观感受缺陷形态差异第三步计算缺陷连通域的大小分布判断哪些是小目标。这一步决定了后处理策略如果大量缺陷连通域只有几十个像素那么预测后做形态学开运算可能会把小缺陷整个滤掉后处理要格外保守。import matplotlib.pyplot as plt import numpy as np def plot_class_distribution(labels): unique, counts np.unique(labels, return_countsTrue) plt.bar(unique, counts) plt.title(NEU-DET Class Distribution) plt.xlabel(Class ID) plt.ylabel(Sample Count) plt.show() def defect_area_ratio(mask): # 统计缺陷像素占整张图的比例判断类别不平衡程度 return (mask 0).sum() / mask.sizeEDA 的价值是避免带着错误的先验去设计模型。比如做完面积统计发现 crazing 类缺陷的平均面积最小、形状细长就可以在训练时给该类的 loss 加权重或者在后处理时改用连通域过滤而不是简单的开闭运算。这类数据洞察写在代码注释里比训练完猜原因靠谱得多。4.2 数据加载与训练脚本核心代码与参数说明PyTorch Dataset 负责把图像路径和掩码路径配对读取后同时应用图像增强和掩码增强。这里的关键是保证 transform 对 image 和 mask 用同一随机种子或者用 Albumentations 的同步机制。import torch from torch.utils.data import Dataset import albumentations as A from albumentations.pytorch import ToTensorV2 class SteelDataset(Dataset): def __init__(self, file_list, transformNone): self.file_list file_list self.transform transform def __len__(self): return len(self.file_list) def __getitem__(self, idx): img_path self.file_list[idx] mask_path img_path.replace(.jpg, _mask.png) image cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) if self.transform: aug self.transform(imageimage, maskmask) image aug[image] mask aug[mask] return image, mask.float() train_aug A.Compose([ A.HorizontalFlip(p0.5), A.RandomBrightnessContrast(p0.3), A.Rotate(limit10, border_modecv2.BORDER_CONSTANT), A.Normalize(mean0.485, std0.229), ToTensorV2() ])Albumentations 的 Compose 会自动对 image 和 mask 应用同一套几何变换不用自己操心随机种子同步。Rotate 的 border_mode 设成 BORDER_CONSTANT旋转后空白区域填 0避免反射填充在图像边缘制造出虚假纹理。Normalize 之后 ToTensorV2 输出的是浮点张量mask 也要转成 float否则和模型输出做 loss 计算时数据类型不匹配。训练循环本身不复杂但有几个细节直接影响结果。模型要设成 train 模式BN 层才会用 batch 统计量每个 epoch 结束在验证集上跑一次记录 mIoU 和 loss保存权重时同时保存 model.state_dict() 和 optimizer.state_dict()这样断点续训时学习率调度状态不会丢。for epoch in range(start_epoch, epochs): model.train() train_loss 0.0 for images, masks in train_loader: images, masks images.to(device), masks.to(device).unsqueeze(1) optimizer.zero_grad() preds model(images) loss bce_dice_loss(preds, masks) loss.backward() optimizer.step() train_loss loss.item() * images.size(0) # 验证阶段切换 eval 模式关闭梯度计算 model.eval() val_iou evaluate(model, val_loader) print(fEpoch {epoch1}/{epochs} loss{train_loss/len(train_loader):.4f} mIoU{val_iou:.4f}) if val_iou best_iou: torch.save(model.state_dict(), best_model.pth) best_iou val_iou验证阶段模型要切 eval 模式并包在 torch.no_grad() 里否则 BN 统计量会被验证数据污染导致训练完的模型和验证时的行为不一致。mask 在送入 loss 前 unsqueeze(1) 补通道维度和模型输出的 [B, 1, H, W] 对齐这是新手最容易漏的一步。4.3 评估指标IoU、Dice 与像素精度怎么读分割任务最常用的指标是 mIoU对所有类别分别算 IoU 再取平均。缺陷检测场景里 IoU 的解读要小心缺陷像素占比本来就小一个 200×200 的图里缺陷可能只有 500 个像素全预测背景的 IoU 是 0预测位置偏移几个像素 IoU 会掉得很快。所以 IoU 对定位精度非常敏感这是它的优点。def compute_iou(pred_mask, true_mask, threshold0.5): pred_bin (torch.sigmoid(pred_mask) threshold).float() intersection (pred_bin * true_mask).sum() union ((pred_bin true_mask) 0).sum() return (intersection 1e-6) / (union 1e-6)threshold 取 0.5预测概率大于 0.5 的像素判为缺陷。评估时不要只看 mIoU 一个数把六个类别的 IoU 分开打印通常能发现某一类特别差。NEU-DET 里 crazing 类的 IoU 普遍偏低因为它细长、面积小边界稍微偏移几个像素 IoU 就降到 0.3 以下这是数据本身的难度不代表模型坏了。提示评估时把六类 IoU 分开打印比只看 mIoU 更容易定位问题类别尤其是 crazing 这类细长缺陷。5. 钢材缺陷训练避坑指南五个实际翻车现场与排查方法5.1 现象与原因速查表入坑 UNet 做 NEU-DET 的常见问题先给一张速查表后面逐个展开排查过程。现象可能原因快速排查loss 不降或 NaN学习率过大 / 输入含 NaN打印 loss 前 10 步降到 1e-5 验证验证 mIoU 高但实测差数据泄漏 / 增强不一致检查划分是否在增强之后全预测背景类别不平衡 / 阈值过高打印预测均值调低 threshold 或加 Dice显存不足batch size 过大减半或开启梯度累积掩码边缘锯齿插值方式错误改成最近邻插值5.2 五个高频踩坑记录踩坑一loss 一开始就 NaN权重全变 NaN现象训练第一个 batch 的 loss 就是 nan后续所有参数变成 nan模型直接报废。 原因最常见的是输入图像存在全黑或全白区域Normalize 后出现极端值或者学习率 1e-3 起步对深网络太激进。还有一个隐蔽原因是 mask 里出现了 0 和 1 之外的值比如标注转换时没二值化255 被当成类别值。 解决先在 DataLoader 里加断言检查 mask 的取值集合必须是 {0, 1}学习率固定 1e-4 起步如果输入归一化用的是 ImageNet 统计量确认灰度图的通道数写法正确。踩坑二验证集 mIoU 高达 0.85实测新图却一塌糊涂现象验证指标很好看把模型拿到产线新拍的图上测试缺陷几乎全漏。 原因这是典型的数据泄漏。之前有人把增强后的图片直接保存到磁盘再统一划分数据集同一张原始图的不同增强版本同时进了训练集和验证集模型实际上见过验证集的“亲戚”。 解决先划分文件路径再做增强增强只发生在训练 DataLoader 内部。验证集和测试集只做 Normalize不做任何随机增强。这是吃过最大的亏从那以后每次划分数据都强制检查训练集和验证集的文件路径是否完全无交集。踩坑三模型把所有像素预测为背景mIoU 停在 0 附近现象训练 loss 在下降但预测结果全黑没有任何前景像素。 原因BCE Loss 在缺陷占比极低时全预测背景的 loss 也很低模型找到了这个“捷径”。默认 threshold 0.5 也可能偏高sigmoid 输出整体被压到 0.5 以下。 解决换成 BCE Dice 组合 loss提高 Dice 权重到 0.7把 threshold 调低到 0.3 观察如果还是没有前景检查 mask 读取路径是否正确很可能读到的全是 0。踩坑四200×200 图像加大输入尺寸后显存爆炸现象尝试把输入 resize 到 512×512 提高精度显存直接 OOM。 原因UNet 的显存占用和输入尺寸近似平方关系200×200 升到 512×512特征图内存翻了 6 倍以上。 解决保持原始 200×200 训练确实需要大图就在编码器第一层加 stride2 的降采样或改用 DeepLab 类轻量结构也可以梯度累积batch size 减半、累积 2 步效果接近但显存省一半。踩坑五训练曲线正常但推理速度不够产线用现象验证 mIoU 0.7 以上但 GPU 上单张推理 30ms产线节拍要求 10ms。 原因UNet 是重型结构参数量大推理时 BN 和 sigmoid 都算在耗时里。 解决先用 torch.jit.script 或 ONNX 导出省掉 Python 调度开销推理时把输入换成半精度 float16如果还超时把编码器换轻量 backbone 或裁剪通道数。精度和速度的取舍是部署阶段的常态没有免费午餐。6. 模型推理与部署验证从单张测试到产线适配6.1 单张图像推理加载权重输出分割掩码训练完成后最直接的价值就是把模型用起来。推理脚本要做的第一件事是确认输入预处理和训练时完全一致包括灰度读取、归一化参数任何一边不一致都会导致输入分布偏移模型表现断崖式下跌。import torch import cv2 import numpy as np model.load_state_dict(torch.load(best_model.pth, map_locationcpu)) model.eval() img cv2.imread(test.jpg, cv2.IMREAD_GRAYSCALE) img cv2.resize(img, (200, 200)) img_t torch.from_numpy(img).float().div(255) img_t (img_t - 0.485) / 0.229 img_t img_t.unsqueeze(0).unsqueeze(0).to(device) with torch.no_grad(): pred model(img_t) pred_mask torch.sigmoid(pred)[0, 0].cpu().numpy() pred_bin (pred_mask 0.5).astype(np.uint8) * 255这段代码的关键是 unsqueeze 两次一次补 batch 维度一次补通道维度因为模型输入是 [B, 1, H, W]。推理全程包在 no_grad 里省显存也省时间。输出 pred_mask 是原始概率图建议保存一份概率图和一份二值图概率图方便后面调 threshold二值图直接给产线判断用。6.2 输出验证与产线适配技巧模型跑通之后我习惯做三个验证第一把预测掩码按 50% 透明度叠加在原图上肉眼确认缺陷位置是否和实际吻合第二统计预测缺陷的连通域面积分布和训练时 EDA 的面积分布对比如果出现大面积异常连通域多半是背景误检第三专门收集一批产线新图覆盖不同光照、不同钢材表面状态构成一个独立的小测试集这个集合永远不进训练流程。产线部署时还有一个实用技巧给预测结果加形态学后处理时先对概率图做高斯模糊再二值化能滤掉单像素噪声比直接二值化后做开运算保留的小缺陷更多。阈值不要拍脑袋定 0.5拿一批验证图扫一遍 0.3 到 0.7画出准确率和召回率曲线挑的是曲线拐点而不是默认值这是我在调试阶段最常用的习惯。从那以后我每次做完缺陷检测项目都会强制走一遍“先分层划分、再查 mask 取值、训练中盯验证集、推理前核对预处理”这四步。这套流程在钢材缺陷项目里救了我好几次也适用于其他工业视觉分割任务。希望帮到你。本文还有配套的精品资源点击获取
