简介这份钢材缺陷图像分割数据集面向从事工业质检、缺陷检测与语义分割方向的算法工程师、研究生及竞赛选手解决真实钢材表面缺陷样本稀缺、标注成本高、难以直接投入训练的问题。数据已完成训练集与验证集划分训练集约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, transformNone): 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(imageimg, maskmask) 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(imageimg, maskmask)这一行albumentations 会把同一组随机参数同时作用到图像和掩码上这是保证几何一致性的标准做法。如果你用 torchvision 的 transform就得手动同步随机种子容易出错。参数上img_dir和mask_dir对应前面说的目录结构transform传增强管道。掩码转long是因为分割损失函数如 CrossEntropyLoss要求标签是整型类别 ID转成 float 会直接报错。3.2 钢材缺陷场景的增强选型增强不是越多越好。钢材缺陷的形态有很强的物理约束划痕是细长的、结疤是块状的、裂纹有方向性。如果你上随机旋转 90 度、大角度翻转可能造出物理上不存在的缺陷形态模型学到的就是噪声。我一般会这么配水平翻转、垂直翻转、小角度旋转±15 度以内、随机裁剪、对比度/亮度微调。弹性形变和网格畸变要慎用钢材表面缺陷的几何形态相对刚性过度形变会让边界变得不真实。train_transform A.Compose([ A.HorizontalFlip(p0.5), A.VerticalFlip(p0.5), A.RandomRotate90(p0.0), # 钢材缺陷有方向性关掉 90 度旋转 A.Rotate(limit15, p0.5), # 只做小角度旋转 A.RandomResizedCrop(height512, width512, scale(0.7, 1.0), p0.5), A.RandomBrightnessContrast(brightness_limit0.15, contrast_limit0.15, p0.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_weight0.5, dice_weight0.5, ignore_index255): super().__init__() self.ce_weight ce_weight self.dice_weight dice_weight self.ce nn.CrossEntropyLoss(ignore_indexignore_index) def forward(self, logits, targets): ce_loss self.ce(logits, targets) # 对每个类别单独算 dice 再平均 num_classes logits.shape[1] probs F.softmax(logits, dim1) 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_lossignore_index255是分割里的常见约定把不参与训练的像素比如标注边界标成 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 在中小数据集上收敛更稳。参数推荐值说明输入尺寸512×512兼顾细节与显存batch size48视显存调整初始学习率1e-4预训练/ 1e-3从零AdamW权重衰减1e-4防止过拟合训练轮数50100配合早停损失权重CE 0.4 / Dice 0.6背景占比高时学习率调度用 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_classes5) for i, v in enumerate(ious): print(f类别 {i} IoU: {v:.4f})这个脚本逐类别算交并比union为 0 时返回 nan 表示该类别在验证集里没出现不算入平均。跑完你会看到某些类别 IoU 明显偏低比如裂纹这种细长缺陷因为像素占比小、边界难学IoU 天然吃亏。这时候针对性补该类样本、或者对该类调高损失权重比盲目调全局学习率有效得多。另一个技巧是可视化错误区域把预测和真值的差异图叠到原图上看模型是在边界上错还是在区域内部错。边界错说明感受野或上采样不够区域内部错说明特征判别力不足。这两种问题的解法完全不同不看图根本分不清。从那以后我每次拿到新的分割数据集都强制先跑一遍逐类 IoU 和错误可视化再决定要不要调模型结构。数据本身的信息量往往比换网络更大。希望帮到你。本文还有配套的精品资源点击获取