钢材缺陷分割实战：4100张像素级标注数据训练与避坑指南

简介：这份资源面向从事工业缺陷检测与语义分割的算法工程师、研究生及竞赛选手，提供一套可直接投入训练的钢材缺陷图像分割数据集，解决多类别缺陷标注数据获取难、预处理繁琐的问题。压缩包共约2000个文件，以png掩膜与jpg原图为主，另含1个txt类别说明和1个py可视化脚本，整体约102.62MB。数据已按训练集约2900张、验证集约1200张划分完毕，标签共5类：0为背景，4种缺陷对应像素值1、2、3、4，类别定义见classes文件。配套脚本可随机抽取一张图片，将原图、GT图像及GT在原图上的蒙板叠加结果一并展示并保存到当前目录，便于快速核验标注质量。资源还关联UNet、SwinUnet、TransUnet等改进网络专栏，方便读者直接套用模型开展实验。目前已有63人学习，适合需要快速搭建钢材缺陷分割基线、验证改进网络效果的用户。

1. 钢材缺陷分割这套 4100 张带标签数据，到底能省下多少标注成本

做过工业质检的都知道，钢材表面缺陷检测最磨人的不是模型结构，而是数据。划痕、夹杂、结疤、氧化铁皮、裂纹这几类缺陷在整卷带钢上分布极不均匀，有的缺陷只占几十个像素，有的又连成大片，用目标检测框去标，框里一半是正常纹理，模型学得稀里糊涂。语义分割把每个像素归到缺陷类别或背景，边界贴合得多，但代价是标注量成倍上涨。这套资源给的就是 4100 张已经打好标签的钢材缺陷分割数据，多类别、像素级掩码，直接对应工业缺陷检测里最典型的语义分割任务。它适合两类人：一类是刚接触深度学习缺陷检测、想找一个真实工业场景练手的学生和转行者；另一类是手里有产线图像但缺标注、想先拿现成数据把训练和推理链路跑通的工程师。省下的不是一点标注工时，而是从零搭一套像素级标注流程的那几周。

2. 语义分割数据怎么读：标签格式、类别映射与训练前自检

拿到一份分割数据集，第一件事不是急着喂进网络，而是把标签的存储格式、类别编码和图像尺寸对齐关系搞清楚。钢材缺陷分割这类工业数据，标签常见两种存法：一种是单通道灰度掩码图，像素值 0 表示背景，1、2、3…对应不同缺陷类别；另一种是彩色掩码，用不同 RGB 值区分类别。两种都能用，但读取和转换方式完全不同，搞错了模型会学出一堆玄学结果。

2.1 灰度掩码与彩色掩码的读取差异

灰度掩码最省事，直接按单通道读进来就是类别索引。彩色掩码则必须先做颜色到类别的映射，否则同一类缺陷因为颜色抖动被拆成好几类。我一般会先写个脚本统计所有标签图里出现过的唯一像素值，确认类别数对不对得上。

import numpy as np from PIL import Image import os label_dir = "labels" unique_vals = set() for name in os.listdir(label_dir): if not name.lower().endswith((".png", ".bmp", ".tif")): continue img = np.array(Image.open(os.path.join(label_dir, name))) # 灰度掩码直接取唯一值；彩色掩码需先转成类别索引 if img.ndim == 2: unique_vals.update(np.unique(img).tolist()) else: # 彩色掩码：把 RGB 三元组映射成整数再统计 flat = img.reshape(-1, img.shape[-1]) codes = flat[:, 0].astype(np.int32) * 65536 + flat[:, 1].astype(np.int32) * 256 + flat[:, 2] unique_vals.update(np.unique(codes).tolist()) print("唯一像素/颜色编码数量:", len(unique_vals)) print(sorted(unique_vals)[:20])

这段脚本的作用是训练前的第一道体检。灰度掩码下，如果打印出来的唯一值只有 0 到 4，说明是 5 类（含背景）；如果出现 255 这种值，多半是标注工具导出的调色板图，需要重新映射。彩色掩码下，唯一编码数量应该等于类别数，如果远大于类别数，说明标注时颜色不统一，得先做颜色聚类再统一映射。参数上，label_dir指向标签目录，img.ndim == 2判断是否为单通道，这是区分两种格式最直接的依据。

2.2 图像与标签的一一对应检查

工业数据经常出现图像和标签文件名对不上、数量不一致的情况，尤其是分批标注、多人协作的项目。训练前必须做一次配对检查，否则 DataLoader 会在某个 batch 突然报文件找不到，训练中断，白跑几个小时。

import os img_dir = "images" label_dir = "labels" img_names = {os.path.splitext(f)[0] for f in os.listdir(img_dir)} label_names = {os.path.splitext(f)[0] for f in os.listdir(label_dir)} only_img = img_names - label_names only_label = label_names - img_names print("有图无标签:", len(only_img), list(only_img)[:5]) print("有标签无图:", len(only_label), list(only_label)[:5]) print("配对成功:", len(img_names & label_names))

逻辑很直白：用集合差集找出单边存在的文件。常见做法是先把这些异常文件挪到_orphan目录，不要直接删，因为有时候是命名后缀大小写不一致导致的假异常。参数上，os.path.splitext去掉扩展名，避免.jpg和.png混用造成的误判。这一步跑完，配对数量应该和数据集宣称的 4100 张对得上，对不上就得回头查标注批次。

2.3 类别不平衡的统计与应对

钢材缺陷里，裂纹、结疤这类严重缺陷样本往往很少，氧化铁皮、划痕又多得离谱。直接训练，模型会偏向多数类，少数类召回率低得没法看。我一般会先统计每个类别的像素占比，再决定要不要加权损失或过采样。

import numpy as np from PIL import Image import os label_dir = "labels" num_classes = 5 pixel_count = np.zeros(num_classes, dtype=np.int64) for name in os.listdir(label_dir): if not name.lower().endswith((".png", ".bmp", ".tif")): continue mask = np.array(Image.open(os.path.join(label_dir, name))) if mask.ndim == 3: mask = mask[:, :, 0] # 灰度化，前提是已确认单通道编码 for c in range(num_classes): pixel_count[c] += np.sum(mask == c) total = pixel_count.sum() for c in range(num_classes): print(f"类别 {c}: 像素占比 {pixel_count[c] / total:.4f}")

统计结果里，如果背景占比超过 95%，说明缺陷本身就很稀疏，这是工业分割的常态，不用慌，但损失函数要选对。常见做法是用交叉熵加类别权重，权重取该类像素占比的倒数再归一化；或者用 Dice Loss、Focal Loss 这类对不平衡更鲁棒的损失。参数上，num_classes要和你前面统计出的类别数一致，别硬编码成 5 结果数据里只有 4 类。

3. 从掩码到训练：语义分割模型选型与最小可跑通流程

数据检查完，接下来是选模型和搭训练流程。钢材缺陷分割属于典型的工业语义分割，输入分辨率通常不小（产线相机拍出来动辄 2048 宽），缺陷尺度跨度大，小目标多。选型上不用一上来就追最新结构，先把一个稳定 baseline 跑通，再谈涨点。

3.1 U-Net、DeepLab 与 SegFormer 的取舍

U-Net 是最稳的起点，编码器-解码器加跳跃连接，对小缺陷的边界恢复友好，数据量不大时也不容易过拟合。DeepLab 系列靠空洞卷积扩大感受野，适合缺陷尺度变化大的场景，但训练时对学习率更敏感。SegFormer 是 Transformer 路线，全局建模能力强，但显存吃得凶，4100 张数据如果分辨率高，单卡可能 batch size 只能开到 2 或 4。我的建议是：先用 U-Net 把整条链路跑通，确认数据读取、损失、评估都没问题，再换 DeepLab 或 SegFormer 做对比实验。工业项目里，能稳定复现的 baseline 比刷高一个点的花哨结构值钱得多。

3.2 一个最小可跑通的训练脚本骨架

下面这个骨架用 PyTorch 写，覆盖数据加载、模型前向、损失计算和一轮验证。不追求完整工程化，目的是让你拿到数据后能最快看到 loss 在降。

import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader import numpy as np from PIL import Image import os class SteelDefectDataset(Dataset): def __init__(self, img_dir, label_dir, img_size=512): self.img_dir = img_dir self.label_dir = label_dir self.names = [os.path.splitext(f)[0] for f in os.listdir(img_dir) if os.path.exists(os.path.join(label_dir, os.path.splitext(f)[0] + ".png"))] self.img_size = img_size def __len__(self): return len(self.names) def __getitem__(self, idx): name = self.names[idx] img = Image.open(os.path.join(self.img_dir, name + ".jpg")).convert("RGB") mask = Image.open(os.path.join(self.label_dir, name + ".png")) img = img.resize((self.img_size, self.img_size)) mask = mask.resize((self.img_size, self.img_size), Image.NEAREST) img = np.array(img, dtype=np.float32) / 255.0 mask = np.array(mask, dtype=np.int64) if mask.ndim == 3: mask = mask[:, :, 0] img = torch.from_numpy(img).permute(2, 0, 1) mask = torch.from_numpy(mask) return img, mask # 这里用一个简化的 U-Net 占位，实际可替换为 smp.Unet 等实现 class SimpleUNet(nn.Module): def __init__(self, in_ch=3, num_classes=5): super().__init__() self.enc = nn.Sequential( nn.Conv2d(in_ch, 32, 3, padding=1), nn.ReLU(), nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), ) self.dec = nn.Conv2d(64, num_classes, 1) def forward(self, x): return self.dec(self.enc(x)) dataset = SteelDefectDataset("images", "labels") loader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=2) model = SimpleUNet(num_classes=5).cuda() criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) for epoch in range(5): model.train() total_loss = 0.0 for img, mask in loader: img, mask = img.cuda(), mask.cuda() pred = model(img) loss = criterion(pred, mask) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() print(f"epoch {epoch}, loss {total_loss / len(loader):.4f}")

逻辑说明：SteelDefectDataset负责配对读取图像和掩码，Image.NEAREST保证掩码缩放时不引入新类别值，这点很关键，用双线性插值会把类别索引插成小数，训练直接崩。SimpleUNet只是占位，真实项目里换成segmentation_models_pytorch的 U-Net 或自己搭的编码器都行。参数上，img_size根据显存调，512 是工业分割的常见起点；batch_size=4是保守值，显存够可以往上加；lr=1e-3配合 Adam 是 baseline 常用组合，换 DeepLab 时建议降到 1e-4 量级。

3.3 评估指标别只看像素准确率

工业缺陷分割里，背景占绝大多数，像素准确率轻松上 98%，但模型可能一个缺陷都没分出来。必须看 IoU 和 Dice，尤其是每个类别的 IoU。常见做法是算 mIoU 的同时，单独打印少数类的 IoU，如果裂纹类 IoU 低于 0.1，说明模型根本没学到，得回头查标签里裂纹样本是不是太少，或者损失权重没设对。评估代码可以在验证集上跑混淆矩阵，再按类别算交并比，这里不展开贴，但这一步不能省。

4. 避坑与排查：钢材缺陷分割训练里最容易翻车的五件事

这一章是我自己踩过和帮别人排查过的真实问题，按「现象 → 原因 → 解决」写，你训练时如果遇到类似情况，可以直接对号入座。

4.1 现象：loss 一直不降，输出全是背景

原因通常是标签读取错了。彩色掩码被当成灰度读，或者灰度掩码里背景是 255 而不是 0，模型学到的类别索引和损失函数期望的对不上。解决：回到 2.1 的统计脚本，确认唯一像素值和类别数一致，背景必须是 0，类别从 1 开始连续编号。如果标注工具导出的是 255 背景，写一行映射mask[mask == 255] = 0再训练。

4.2 现象：训练正常，验证集 IoU 极低

多半是图像和标签的预处理不一致。训练时图像做了归一化，验证时忘了；或者训练 resize 用了双线性，验证用了最近邻，导致输入分布不同。解决：把预处理封装成一个函数，训练和验证共用，别在两处各写一遍。参数上，归一化的均值和标准差要固定，不要用每个 batch 自己算。

4.3 现象：显存爆了，batch size 只能开到 1

工业图像分辨率高，U-Net 这种全分辨率解码的结构显存占用随分辨率平方增长。解决：先降img_size到 512 或 384 跑通，再考虑用混合精度训练或梯度累积。常见做法是torch.cuda.amp加GradScaler，显存能省三成左右，速度也快。别一上来就上 1024，除非你确认卡够大。

4.4 现象：少数类缺陷完全检测不到

类别不平衡没处理。解决：在损失里加类别权重，权重取该类像素占比倒数的归一化值；或者对含少数类的图像做过采样。注意权重别设得太极端，否则模型会把正常纹理也判成缺陷，误检率飙升。我一般会把权重上限卡在 10 倍以内，超过就说明数据本身需要补标。

4.5 现象：推理时边界毛糙，小缺陷断成几截

这是分割任务的老问题，跟模型感受野和损失都有关。解决：训练时加 Dice Loss 或 Boundary Loss 辅助，推理后处理用连通域分析把小碎块合并。常见做法是cv2.connectedComponents过滤掉面积小于阈值的区域，阈值按实际缺陷最小像素面积设，别拍脑袋定。

5. 把 4100 张数据用透：进阶技巧与验证习惯

数据量固定的时候，涨点靠的是把每一张都用到位。我自己的习惯是先把数据按缺陷类型分层，保证训练集和验证集里每类缺陷都有代表，而不是随机切分导致验证集里某类缺陷一张都没有。分层切分用sklearn.model_selection.StratifiedKFold或者自己按类别统计后手动分配，虽然麻烦，但验证结果可信度高得多。

进阶用法上，工业分割很吃数据增强，但增强要符合物理意义。水平翻转、垂直翻转、小角度旋转一般没问题，钢材纹理本身方向性不强；但大角度旋转和随机裁剪要小心，裁剪可能把缺陷裁没，旋转可能引入不存在的边界。我一般会用 Albumentations 配一套保守增强，参数如下表。

验证习惯上，我每次改完模型或损失，都会固定跑一遍同一批验证图，把预测掩码叠加到原图上肉眼过一遍。指标涨了但肉眼看着不对的情况太常见了，尤其是边界处把正常纹理判成缺陷。这一步花不了几分钟，但能拦住很多「指标好看、上线翻车」的模型。从那以后我每次训练完都强制走一遍可视化验证，不看图不认指标。希望这套数据和流程能帮你把钢材缺陷分割这条链路真正跑起来。

本文还有配套的精品资源，点击获取