ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

钢材缺陷图像分割数据集实战:4100张标注图直接训练与调优指南

钢材缺陷图像分割数据集实战:4100张标注图直接训练与调优指南 简介这份钢材缺陷图像分割数据集面向从事工业质检、表面缺陷检测的算法工程师与深度学习学习者提供可直接投入训练的多类别分割样本。数据已按训练集与验证集划分完毕训练集约2900张图像及对应mask验证集约1200张合计约4100张标签采用png像素值区分背景与4种缺陷对应像素1至4具体类别可查阅classes文件。压缩包共2000个文件以1273个png与725个jpg为主另含1个txt类别说明和1个py可视化脚本整体约102.78MB7z格式便于解压。脚本可随机抽取一张图片展示原图、GT图像及GT在原图上的蒙板效果并保存至当前目录方便快速核验标注质量。目前已有67人学习适合直接用于UNet、SwinUNet、TransUNet等分割网络的训练与改进实验省去数据清洗与格式转换环节。1. 钢材缺陷图像分割数据集4100 张标注图怎么直接喂进训练脚本拿到一个「已处理完可以直接训练」的钢材缺陷图像分割数据集第一反应往往不是兴奋而是怀疑4100 张到底分了几类、掩码是单通道索引图还是多通道二值图、目录结构能不能被torchvision.datasets或segmentation_models_pytorch直接吃进去。钢材表面缺陷检测在工业质检里是刚需热轧板、冷轧板上的划痕、夹杂、氧化铁皮、裂纹靠人工肉眼在高速产线上盯漏检率压不下来。图像分割比目标检测多给一层像素级边界对后续的缺陷面积统计、等级判定、切割路径规划都有直接价值。这个数据集适合两类人一类是想跑通语义分割全流程但缺工业数据的算法工程师另一类是产线质检团队想验证「分割模型能不能替代人工标注」。下面按「先看清数据长什么样再决定怎么训」的顺序拆开讲。2. 先搞清 4100 张钢材缺陷图的标注格式与类别分布2.1 语义分割与实例分割的选型差别钢材缺陷分割最常见的两种标注形态语义分割semantic segmentation给每个像素一个类别标签同类缺陷连成一片实例分割instance segmentation在语义基础上再区分「这是第几个缺陷」。4100 张这个量级如果缺陷之间粘连严重、需要数个数实例分割更合适代表方案是 Mask2Former 或 YOLOv8-seg如果只关心「这块区域是不是缺陷、是什么缺陷」语义分割足够DeepLabV3、U-Net、SegFormer 都能跑。判断方法很直接打开一张标注图看同类缺陷挨在一起时掩码是连成一块还是分成两块。连成一块就是语义标注分成两块且颜色不同就是实例标注。钢材缺陷里划痕经常细长且多条并行如果标注时没做实例区分后期想数条数就得返工。提示拿到数据集先别急着写 DataLoader用下面这段脚本统计类别像素占比能直接暴露类别不平衡问题。import os import numpy as np from PIL import Image from collections import Counter # 假设掩码是单通道索引图像素值即类别 id mask_dir dataset/masks counter Counter() total_pixels 0 for name in os.listdir(mask_dir): mask np.array(Image.open(os.path.join(mask_dir, name))) # 若掩码是 RGB 三通道需先映射回类别 id这里按单通道处理 values, counts np.unique(mask, return_countsTrue) for v, c in zip(values, counts): counter[int(v)] int(c) total_pixels mask.size for cls_id, cnt in sorted(counter.items()): print(fclass {cls_id}: {cnt} px, ratio {cnt/total_pixels:.4f})这段脚本的逻辑是逐张读取掩码、统计每个像素值出现次数最后输出各类别像素占比。参数上要注意mask_dir指向掩码目录如果掩码是 RGB 三通道比如用颜色区分类别需要先建立颜色到类别 id 的映射表再统计否则会把颜色值当成类别 id。跑完你会看到类似「背景占 92%、划痕占 3%、夹杂占 2%」的结果这就是后面选损失函数和采样策略的依据。2.2 目录结构决定你能不能「直接训练」「已处理完可以直接训练」通常意味着目录已经整理成images/和masks/两个平行文件夹文件名一一对应。但不同来源的数据集差异很大常见三种结构结构类型目录形态适配框架注意点平行文件夹images/train、masks/train自定义 Dataset文件名必须严格对应索引图images 单张 label 索引图需转换索引图要拆成单类掩码COCO/VOCannotations.json 或 XML需转换分割任务要转成掩码如果掩码是单张索引图所有类别画在一张图上像素值 0/1/2/3语义分割可以直接用但多类别训练时要把索引图转成 one-hot 或直接用 CrossEntropyLoss 接受索引图。如果掩码是每类一张二值图那就要在 Dataset 里堆叠成多通道。我一般会先写一个 20 行的检查脚本确认三件事图片和掩码数量是否一致、文件名是否一一对应、掩码像素值范围是否在预期类别数内。这三件事任何一件出问题训练时都会以「loss 不下降」或「全预测背景」的形式翻车而且很难从日志里看出来。2.3 类别不平衡在钢材缺陷里的具体表现钢材缺陷数据集的类别分布几乎必然长尾背景占绝对多数划痕、裂纹这类缺陷可能只占几个百分点。更麻烦的是有些缺陷类别在 4100 张里只出现几十张模型很容易学会「全预测背景」也能拿到很高的像素准确率。这时候像素准确率pixel accuracy就是个骗人的指标必须看 mIoU 和每类的 IoU。处理长尾的常见做法有三种一是损失函数用 Dice Loss 或 Focal Loss 替代纯 CrossEntropy二是采样时对含稀有类别的图做过采样三是训练时对稀有类别像素加权。我一般先跑一版不加权的 baseline看每类 IoU再决定要不要加权。如果稀有类 IoU 低于 0.1加权基本是必须的。3. 用 segmentation_models_pytorch 跑通第一版训练3.1 环境与依赖的最小安装钢材缺陷分割训练环境不复杂核心就几个包。CUDA 版本按你显卡驱动选别盲目装最新。# 创建独立环境避免和已有项目冲突 conda create -n steel_seg python3.10 -y conda activate steel_seg # 安装 PyTorch按官方对应 CUDA 版本选择这里以 cu118 为例 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 # 分割训练常用库 pip install segmentation-models-pytorch albumentations opencv-python tqdm tensorboardsegmentation-models-pytorch的好处是编码器encoder可以直接换从 ResNet 到 EfficientNet 到 MiTSegFormer 的 backbone都有预训练权重。钢材缺陷这种纹理型任务ImageNet 预训练权重能明显加快收敛别从随机初始化开始训。3.2 自定义 Dataset 与增强策略下面是一个能直接套用的 Dataset 写法假设目录是images/和masks/平行结构掩码是单通道索引图。import os import numpy as np import cv2 import torch from torch.utils.data import Dataset 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.names [f for f in os.listdir(img_dir) if os.path.exists(os.path.join(mask_dir, f))] def __len__(self): return len(self.names) def __getitem__(self, idx): name self.names[idx] img cv2.imread(os.path.join(self.img_dir, name)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask cv2.imread(os.path.join(self.mask_dir, name), cv2.IMREAD_GRAYSCALE) if self.transform: aug self.transform(imageimg, maskmask) img, mask aug[image], aug[mask] # 图像归一化到 [0,1] 并转 CHW掩码保持 long 类型供 CrossEntropyLoss 使用 img torch.from_numpy(img).permute(2, 0, 1).float() / 255.0 mask torch.from_numpy(mask).long() return img, mask # 训练增强几何变换对钢材缺陷有效颜色抖动要克制 train_transform A.Compose([ A.Resize(512, 512), A.HorizontalFlip(p0.5), A.VerticalFlip(p0.5), A.RandomRotate90(p0.5), A.RandomBrightnessContrast(p0.3), ])逻辑说明__init__里做了一次文件名交集过滤这是血泪经验——很多数据集图片和掩码数量对不上不提前过滤会在训练到一半时崩掉。__getitem__里图像转 RGB、掩码保持灰度是因为 OpenCV 默认读 BGR直接送进模型会导致颜色通道错位模型也能训但精度会掉。增强部分几何变换给得比较足钢材缺陷方向不固定翻转旋转都合理颜色抖动只给 0.3因为钢材表面颜色本身就是判别特征之一抖太狠会破坏纹理信息。参数上Resize(512,512)是权衡显存和精度的常见选择4100 张图如果原图很大可以先缩到 512 训练验证阶段再上原图。掩码用最近邻插值albumentations 对 mask 默认就是最近邻不能用双线性否则会插出 0.5 这种不存在的类别 id。3.3 模型、损失与训练循环的关键参数import torch import torch.nn as nn import segmentation_models_pytorch as smp # 多类别语义分割classes 按你实际类别数改含背景 model smp.DeepLabV3Plus( encoder_nameresnet34, encoder_weightsimagenet, in_channels3, classes5, # 背景 4 类缺陷按实际改 ) # 类别不平衡时给稀有类更高权重权重按像素占比倒数估算 class_weights torch.tensor([0.1, 1.0, 2.0, 3.0, 3.0]).cuda() criterion nn.CrossEntropyLoss(weightclass_weights) optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max50)encoder_name选 resnet34 是起步配置显存够可以换 resnet50 或 mit_b2。classes一定要含背景这是新手最常犯的错——把背景漏掉模型输出通道数对不上训练直接报维度错误。class_weights的数值不是拍脑袋是按 2.1 节统计出的像素占比倒数归一化后估的背景权重压到 0.1稀有类提到 3.0。学习率 1e-4 配 AdamW 是分割任务的稳妥起点太大容易在早期把预训练权重冲掉。训练循环里我习惯每轮验证后打印每类 IoU而不是只看总 loss。总 loss 下降但稀有类 IoU 不动说明模型在偷懒预测背景这时候要回去调权重或加过采样。4. 训练过程中的避坑与排查清单4.1 现象loss 降到某个值就不动了mIoU 卡在 0.3 以下原因通常有三个学习率太大导致在局部震荡、类别权重设置不合理、或者掩码类别 id 和模型输出通道没对齐。先检查掩码最大值是否等于classes-1如果掩码里有 255 这种「忽略区域」标记要在 Dataset 里把它映射成 255 并在损失函数里设ignore_index255否则会被当成一个真实类别去学。解决顺序先用小学习率1e-5跑 5 个 epoch 看 loss 是否继续下降再把类别权重全设成 1 跑一版对比最后检查掩码 id 分布。三步下来基本能定位。4.2 现象验证集 mIoU 很高但可视化预测结果全是背景这是典型的「指标欺骗」。如果验证集里背景占比 95%模型全预测背景也能拿到很高的像素准确率mIoU 如果只算前景类平均反而会暴露问题。但如果你的 mIoU 实现把背景也算进去且没做类别平均就会虚高。检查你的 mIoU 计算是不是按类求 IoU 再平均而不是把所有类像素混在一起算。另一个可能是验证集和训练集分布不一致比如训练集里缺陷图多、验证集里正常图多。用 2.1 的统计脚本分别跑训练集和验证集对比类别占比。4.3 现象训练中途报 CUDA out of memory4100 张图如果原图分辨率很高比如 4000×3000即使 Resize 到 512如果 DataLoader 的num_workers开太大、batch_size设太高也会爆显存。排查顺序先把batch_size降到 4 或 2再把num_workers降到 2最后考虑用梯度累积模拟大 batch。# 梯度累积小显存模拟大 batch 的常用写法 accum_steps 4 for i, (img, mask) in enumerate(loader): pred model(img.cuda()) loss criterion(pred, mask.cuda()) / accum_steps loss.backward() if (i 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()注意 loss 要除以累积步数否则梯度会放大accum_steps倍等效学习率变大容易训崩。4.4 现象预测掩码边缘锯齿严重缺陷边界对不齐原因多半是上采样方式太粗暴。DeepLabV3 默认用双线性插值上采样边缘会糊。可以在解码器最后加一层nn.ConvTranspose2d或者用 CRF条件随机场做后处理。更简单的做法是训练时把输入分辨率提高比如从 512 提到 768边缘细节会明显改善代价是显存和训练时间增加。另一个容易忽略的点如果标注本身边缘就不精细人工标注常见模型学出来的边界也不会准。这种情况先抽查几十张标注图确认标注质量再谈模型优化。4.5 现象换用 SegFormer 或 Mask2Former 后训练不收敛这两个模型对学习率和 warmup 很敏感。SegFormer 官方配置用 AdamW 多项式衰减 warmup直接套 1e-4 固定学习率往往不收敛。Mask2Former 更复杂它依赖匈牙利匹配做实例分配如果类别数少、每张图实例少匹配会不稳定。钢材缺陷如果只是语义分割没必要上 Mask2FormerDeepLabV3 或 SegFormer 足够别为了追新而追新。5. 把 4100 张用到极致小样本下的进阶技巧4100 张在工业分割里不算多想再往上提点精度有几个我实际用过有效的方向。第一个是测试时增强TTA推理时对同一张图做水平翻转、垂直翻转、多尺度缩放把多次预测结果平均。这个技巧不增加训练成本推理慢几倍但 mIoU 通常能涨 1 到 2 个点适合最终交付前跑一版。def tta_predict(model, img): # img: 1x3xHxW tensor preds [] for flip in [False, True]: for scale in [1.0, 1.25]: x torch.nn.functional.interpolate( img, scale_factorscale, modebilinear, align_cornersFalse) if flip: x torch.flip(x, dims[3]) with torch.no_grad(): p model(x) p torch.nn.functional.interpolate( p, sizeimg.shape[2:], modebilinear, align_cornersFalse) if flip: p torch.flip(p, dims[3]) preds.append(torch.softmax(p, dim1)) return torch.stack(preds).mean(0)第二个方向是伪标签。如果手头还有未标注的钢材表面图用训好的模型跑一遍预测把置信度高的像素当伪标签加入训练集迭代一到两轮。这个方法在工业数据上很实用因为产线每天产生的图远多于标注量。注意伪标签要设置信度阈值低于阈值的像素设成ignore_index别硬塞进去。第三个是类别特定的后处理。钢材缺陷里划痕是细长结构用形态学闭运算能把断开的划痕连起来夹杂是块状用面积阈值过滤小噪点。这些后处理规则不通用要按你数据集的缺陷形态调但调好了比换模型管用。技巧预期收益成本适用阶段TTAmIoU 1~2推理慢 4 倍交付前伪标签mIoU 2~5需未标注数据有额外数据时形态学后处理特定类 3~8调参时间类别形态明确时高分辨率训练边缘 2~4显存翻倍显存充足时最后说个我自己的习惯每次训完模型一定抽 20 张验证集图做可视化把原图、真值掩码、预测掩码并排存成一张图。看指标是一回事看图是另一回事很多问题比如某类缺陷系统性漏检、边界偏移只有看图才能发现。4100 张的数据集认真跑一轮加调优语义分割 mIoU 做到 0.6 以上是合理预期稀有类能到 0.4 以上就算可用。别指望一次训到完美工业缺陷分割本来就是迭代出来的。希望帮到你。本文还有配套的精品资源点击获取
返回列表