简介：这份钢材缺陷图像分割数据集面向从事工业质检、缺陷检测与语义分割方向的算法工程师、研究生及竞赛选手，解决真实钢材表面缺陷样本稀缺、标注成本高、难以直接投入训练的问题。数据已完成训练集与验证集划分：训练集约2900张图像及对应掩膜，验证集约1200张，合计约4100张，标签采用png像素值区分背景与4类缺陷，对应像素值1至4，具体类别可查阅classes文件。压缩包为7z格式，共2000个文件，以1273个png掩膜与725张jpg原图为主，另含1个txt类别说明和1个py可视化脚本，整体约102.78MB。脚本可随机抽取一张图片，展示原图、GT图像及GT在原图上的蒙板效果并保存至当前目录，便于快速核验标注质量。目前已有67人学习，适合直接用于UNet、SwinUNet、TransUNet等分割网络的训练与改进实验。

1. 钢材缺陷图像分割数据集：4100 张带标签的产线级素材到底怎么用

产线上做表面质检的兄弟大概率都经历过这个阶段：算法框架跑通了，模型结构也调了，卡在数据上——要么缺陷样本太少，要么标注格式对不上，要么类别定义跟实际产线对不齐。这份钢材缺陷图像分割数据集就是冲着这个痛点来的：约 4100 张图像，配套像素级分割标签，已经处理成可以直接喂给训练脚本的状态，覆盖多类常见钢材表面缺陷。它解决的不是"有没有数据"的问题，而是"数据能不能直接进训练循环"的问题。适合谁？做工业质检方向、手上有分割模型但缺标注数据、或者想快速验证某个分割网络在钢材场景下表现的从业者。如果你还在用分类标签凑合做缺陷定位，这份像素级标签的价值会体现得很直接。

2. 先搞清楚分割标签的形态：掩码、类别与目录约定

拿到一份分割数据集，第一件事不是急着写训练脚本，而是把标签的物理形态摸清楚。钢材缺陷分割和普通自然图像分割有个明显差别：缺陷区域往往细长、边缘模糊、对比度低，标签的精度直接决定模型能不能学到边界。所以这一章先把"标签长什么样、类别怎么定义、目录怎么组织"讲透，再谈训练。

2.1 像素级掩码与类别映射

图像分割数据集的核心是掩码（mask），每个像素都要有类别归属。钢材缺陷常见的类别包括划痕、夹杂、结疤、裂纹、氧化铁皮压入等，不同数据集类别命名和数量不一样，这份数据是多类别分割，意味着单张图里可能同时出现多种缺陷，掩码里每个像素值对应一个类别 ID。

常见做法是两种标签组织方式：一种是每张图对应一张单通道 PNG 掩码，像素值 0 表示背景，1、2、3…表示不同缺陷类别；另一种是每类一张二值掩码。前者更省空间、加载更快，后者在多类别重叠场景下更清晰。钢材缺陷一般不会大面积重叠，所以单通道索引掩码是主流选择。

我一般拿到数据先跑一段统计脚本，把类别分布和掩码取值摸清楚，避免后面训练时出现"标签里有 5 类但配置文件只写了 3 类"这种低级翻车。

import os import numpy as np from PIL import Image from collections import Counter mask_dir = "dataset/masks" counter = Counter() sizes = [] for name in os.listdir(mask_dir): if not name.lower().endswith((".png", ".bmp")): continue m = np.array(Image.open(os.path.join(mask_dir, name))) # 统计该掩码里出现的所有像素值，即类别 ID counter.update(np.unique(m).tolist()) sizes.append(m.shape) print("类别像素值分布:", dict(sorted(counter.items()))) print("掩码尺寸样例:", sizes[:5])

这段脚本做两件事：用np.unique提取每张掩码里实际出现的像素值，用Counter汇总全数据集的类别分布；同时记录掩码尺寸，方便判断是否需要统一 resize。逻辑上先确认类别 ID 集合，再确认尺寸一致性。参数上mask_dir指向掩码目录，如果你的标签是每类一张二值图，这里要改成遍历子目录。跑完如果发现某个类别像素占比极低（比如低于 0.1%），训练时就要考虑类别加权或过采样，否则模型会直接忽略它。

2.2 目录结构与命名对齐

分割数据集最容易出问题的地方不是模型，是图像和掩码的对应关系。4100 张这个量级，靠人工核对不现实，必须靠命名约定。常见组织方式是图像和掩码同名不同目录：

dataset/ ├── images/ │ ├── 0001.jpg │ ├── 0002.jpg │ └── ... └── masks/ ├── 0001.png ├── 0002.png └── ...

同名不同后缀是最稳的做法，加载时用os.path.splitext替换后缀即可。如果数据集已经划分好 train/val/test，那就在 images 和 masks 下再各分一层。这里有个血泪经验：有些数据集图像是.jpg、掩码是.png，但文件名里带了额外前缀或后缀，直接按 stem 匹配会漏掉一批。所以加载前先做一次配对校验。

import os img_dir = "dataset/images" mask_dir = "dataset/masks" img_stems = {os.path.splitext(f)[0] for f in os.listdir(img_dir)} mask_stems = {os.path.splitext(f)[0] for f in os.listdir(mask_dir)} only_img = img_stems - mask_stems only_mask = mask_stems - img_stems print("有图无标签:", len(only_img), list(only_img)[:5]) print("有标签无图:", len(only_mask), list(only_mask)[:5]) print("成功配对:", len(img_stems & mask_stems))

用集合差集找出单边缺失的样本，这是训练前必做的一步。参数很直白，两个目录路径。如果only_img或only_mask不为空，要么是命名不一致，要么是数据本身有缺失，必须先处理掉，否则训练时 DataLoader 会在某个 batch 直接抛异常，而且报错位置往往离真正的问题很远，排查起来很痛苦。

提示：配对校验建议写进数据加载类的__init__里，每次实例化时自动跑一遍，比事后 debug 省事得多。

3. 把数据接进训练流程：从 Dataset 到增强策略

标签形态清楚了，接下来是让它真正流动起来。分割任务和分类任务在数据管道上的最大区别是：图像和掩码必须做完全一致的几何变换，否则标签就错位了。这一章把 Dataset 封装、增强同步、以及钢材缺陷场景下的增强选型讲清楚。

3.1 自定义 Dataset 与同步增强

不管你用 PyTorch 还是别的框架，核心都是把"读图 + 读掩码 + 同步变换"封装成一个可迭代对象。钢材缺陷图像普遍是灰度或低饱和度，颜色增强意义不大，几何变换和对比度调整才是重点。

import os import numpy as np import torch from torch.utils.data import Dataset from PIL import Image import albumentations as A class SteelDefectDataset(Dataset): def __init__(self, img_dir, mask_dir, transform=None): self.img_dir = img_dir self.mask_dir = mask_dir self.transform = transform # 只保留图像和掩码都存在的样本 self.stems = sorted( set(os.path.splitext(f)[0] for f in os.listdir(img_dir)) & set(os.path.splitext(f)[0] for f in os.listdir(mask_dir)) ) def __len__(self): return len(self.stems) def __getitem__(self, idx): stem = self.stems[idx] img = np.array(Image.open(os.path.join(self.img_dir, stem + ".jpg")).convert("RGB")) mask = np.array(Image.open(os.path.join(self.mask_dir, stem + ".png"))) if self.transform: # 图像和掩码走同一套几何变换，保证像素对齐 augmented = self.transform(image=img, mask=mask) img, mask = augmented["image"], augmented["mask"] img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0 mask = torch.from_numpy(mask).long() return img, mask

关键点在self.transform(image=img, mask=mask)这一行：albumentations 会把同一组随机参数同时作用到图像和掩码上，这是保证几何一致性的标准做法。如果你用 torchvision 的 transform，就得手动同步随机种子，容易出错。参数上img_dir和mask_dir对应前面说的目录结构，transform传增强管道。掩码转long是因为分割损失函数（如 CrossEntropyLoss）要求标签是整型类别 ID，转成 float 会直接报错。

3.2 钢材缺陷场景的增强选型

增强不是越多越好。钢材缺陷的形态有很强的物理约束：划痕是细长的、结疤是块状的、裂纹有方向性。如果你上随机旋转 90 度、大角度翻转，可能造出物理上不存在的缺陷形态，模型学到的就是噪声。

我一般会这么配：水平翻转、垂直翻转、小角度旋转（±15 度以内）、随机裁剪、对比度/亮度微调。弹性形变和网格畸变要慎用，钢材表面缺陷的几何形态相对刚性，过度形变会让边界变得不真实。

train_transform = A.Compose([ A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.5), A.RandomRotate90(p=0.0), # 钢材缺陷有方向性，关掉 90 度旋转 A.Rotate(limit=15, p=0.5), # 只做小角度旋转 A.RandomResizedCrop(height=512, width=512, scale=(0.7, 1.0), p=0.5), A.RandomBrightnessContrast(brightness_limit=0.15, contrast_limit=0.15, p=0.5), A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ])

RandomRotate90设成 0 是刻意的，钢材轧制方向决定了缺陷有明确的方向分布，90 度旋转会破坏这个先验。Rotate限制在 15 度，RandomResizedCrop的 scale 下限设 0.7，避免裁得太狠把缺陷裁没。Normalize用的是 ImageNet 统计量，如果你从零训练可以用数据集自身均值方差，但用预训练权重的话就保持 ImageNet 的。

注意：验证集和测试集不要加任何随机增强，只做 resize 和 normalize。否则每次验证结果都在抖，你根本分不清是模型在学还是增强在捣乱。

4. 训练配置与损失函数：多类别分割怎么设才不崩

数据管道通了，接下来是训练本身。多类别分割和单类别在配置上有几个关键差异：输出通道数、损失函数选择、类别不平衡处理。这一章把这些参数落到具体数值和代码上。

4.1 输出通道与损失函数匹配

分割网络的输出通道数必须等于类别数（含背景）。假设这份数据有 4 类缺陷加 1 个背景，那输出就是 5 通道。常见翻车是配置文件里写 1 通道（二分类思维），结果训练不报错但 mask 全是 0 或 1，模型完全学不到多类别。

损失函数上，多类别分割首选 CrossEntropyLoss，如果类别不平衡严重就上 Dice Loss 或两者加权组合。钢材缺陷里背景像素通常占 90% 以上，纯 CE 会让模型倾向于全预测背景，所以组合损失更稳。

import torch import torch.nn as nn import torch.nn.functional as F class ComboLoss(nn.Module): def __init__(self, ce_weight=0.5, dice_weight=0.5, ignore_index=255): super().__init__() self.ce_weight = ce_weight self.dice_weight = dice_weight self.ce = nn.CrossEntropyLoss(ignore_index=ignore_index) def forward(self, logits, targets): ce_loss = self.ce(logits, targets) # 对每个类别单独算 dice 再平均 num_classes = logits.shape[1] probs = F.softmax(logits, dim=1) dice_loss = 0.0 for c in range(num_classes): pred_c = probs[:, c] target_c = (targets == c).float() intersection = (pred_c * target_c).sum() dice_loss += 1 - (2 * intersection + 1e-6) / (pred_c.sum() + target_c.sum() + 1e-6) dice_loss /= num_classes return self.ce_weight * ce_loss + self.dice_weight * dice_loss

ignore_index=255是分割里的常见约定，把不参与训练的像素（比如标注边界）标成 255 跳过。Dice 部分逐类别计算，加1e-6防止除零。ce_weight和dice_weight是你要调的参数，背景占比高时可以把 dice 权重提到 0.6 甚至 0.7。逻辑上 CE 管像素级分类准确，Dice 管类别区域重叠度，两者互补。

4.2 学习率、batch size 与显存权衡

4100 张图在分割任务里属于中小规模。batch size 受显存限制，512×512 输入下，8GB 显存大概能跑 batch 4 到 8。学习率用预训练权重时从 1e-4 起步，从零训练可以到 1e-3。优化器 AdamW 比 SGD 在中小数据集上收敛更稳。

学习率调度用 CosineAnnealing 或 ReduceLROnPlateau 都行，前者更平滑。如果验证集 mIoU 连续 10 轮不涨就早停，4100 张图过拟合来得比你想的快。

5. 避坑与排查：分割训练里最容易翻车的五件事

这一章全是踩过的坑，按"现象 → 原因 → 解决"写，能帮你省掉大量 debug 时间。

现象一：训练 loss 正常下降，但预测结果全是背景。原因基本是类别极度不平衡，背景像素占比过高，CE 损失被背景主导。解决：把 Dice 权重提上去，或者用带类别权重的 CE，给稀有类别更高权重。也可以先统计各类别像素占比，按反比设权重。

现象二：图像和掩码错位，预测边界整体偏移。原因是增强时图像和掩码没走同一套变换，或者 resize 时用了不同的插值方式。解决：统一用 albumentations 的 Compose 同时处理，掩码 resize 必须用最近邻插值（INTER_NEAREST），用双线性会把类别 ID 插成小数，直接毁掉标签。

现象三：训练时突然报 "Target x is out of bounds"。原因是掩码里的像素值超出了你设定的类别数。比如配置写了 4 类，但掩码里出现了 5。解决：回到 2.1 的统计脚本，把实际类别 ID 集合跑出来，配置里的类别数必须覆盖最大值加一。

现象四：验证集指标远低于训练集，且差距越来越大。典型过拟合。4100 张图对分割网络来说不算多，尤其是深层网络。解决：加数据增强、加权重衰减、用预训练编码器、或者换更轻量的解码器。别一上来就上超大模型。

现象五：显存溢出（OOM），但 batch size 已经调到 1。原因可能是输入尺寸太大，或者模型输出层没做下采样。解决：先把输入降到 256×256 验证流程能跑通，再逐步往上加。另外检查是否有中间特征图没释放，用torch.cuda.empty_cache()辅助排查。

提示：这五条里前三条出现频率最高，建议在正式训练前用一个小子集（比如 50 张）跑通全流程，确认无报错再上全量。

6. 进阶技巧：用 mIoU 逐类分析定位模型的真实短板

训练跑通只是开始，真正决定这份数据能不能产出可用模型的，是你能不能从指标里读出问题。整体 mIoU 是个平均数，会掩盖掉某些类别的糟糕表现。我一般会写一个逐类 mIoU 的评估脚本，把每个缺陷类别的 IoU 单独打出来，再结合混淆矩阵看误分类方向。

import numpy as np import torch def per_class_iou(preds, targets, num_classes): # preds/targets: (N, H, W) 整型 ious = [] for c in range(num_classes): pred_c = (preds == c) target_c = (targets == c) intersection = (pred_c & target_c).sum() union = (pred_c | target_c).sum() iou = intersection / union if union > 0 else float("nan") ious.append(iou) return ious # 假设 all_preds / all_targets 是累积的 numpy 数组 ious = per_class_iou(all_preds, all_targets, num_classes=5) for i, v in enumerate(ious): print(f"类别 {i} IoU: {v:.4f}")

这个脚本逐类别算交并比，union为 0 时返回 nan 表示该类别在验证集里没出现，不算入平均。跑完你会看到某些类别 IoU 明显偏低，比如裂纹这种细长缺陷，因为像素占比小、边界难学，IoU 天然吃亏。这时候针对性补该类样本、或者对该类调高损失权重，比盲目调全局学习率有效得多。

另一个技巧是可视化错误区域：把预测和真值的差异图叠到原图上，看模型是在边界上错还是在区域内部错。边界错说明感受野或上采样不够，区域内部错说明特征判别力不足。这两种问题的解法完全不同，不看图根本分不清。

从那以后我每次拿到新的分割数据集，都强制先跑一遍逐类 IoU 和错误可视化，再决定要不要调模型结构。数据本身的信息量，往往比换网络更大。希望帮到你。

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