工业语义分割实战：4100张钢材缺陷标签从数据到部署全流程

简介本资源为钢材缺陷图像分割数据集面向从事工业缺陷检测、语义分割方向的深度学习开发者与研究人员尤其适合需要多类别分割实战数据的学生和工程师。数据集约4100张图像及对应标签已按训练集2900张、验证集1200张划分完毕可直接投入训练。标签共5类0为背景4种缺陷分别对应png像素值1、2、3、4类别定义见classes文件。压缩包为7z格式共2000个文件其中1273个png标签图、725个jpg原图另含1个txt类别说明与1个py可视化脚本整体约102.62MB。该脚本可随机抽取一张图片展示原图、GT图像及GT在原图上的蒙板效果并保存至当前目录便于快速核验标注质量。目前已有63人学习配套UNet、SwinUnet、TransUnet等改进网络专栏方便读者对照实验、替换主干并复现分割流程。1. 钢材缺陷分割4100 张标签背后工业语义分割到底在做什么产线上刚轧出来的钢板表面那点氧化铁皮、划痕、夹杂人眼在强反光下盯久了基本靠猜。钢材缺陷分割要解决的就是这件事给一张钢板表面图像逐像素判断它属于背景、划痕、夹杂、氧化铁皮还是孔洞。注意这不是目标检测画个框就完事——缺陷的边界往往模糊、细长、互相粘连框住它没用下游要的是精确到像素的掩膜才能算面积、判等级、决定这块板是降级还是回炉。4100 张带标签数据在工业语义分割里算是一个能起步、但远谈不上宽裕的量级。它够训一个中等规模的分割网络但不够你随便堆大模型。所以这篇笔记的主线很明确怎么用这 4100 张多类别标签把钢材缺陷分割从数据检查一路做到能推理、能评估、能上线。适合已经懂点深度学习、手里有标注数据、但还没跑通工业分割全流程的工程师也适合做深度学习毕设、想找一个真实工业场景练手的人。语义分割算法在这里不是炫技是要在缺陷边界这种玄学区域上稳住。2. 从 4100 张标签到可训练掩膜数据这一关怎么过工业分割项目翻车十次有八次不是模型问题是数据问题。4100 张标签听起来不少但如果类别分布歪、边界标注不一致、掩膜格式混乱训出来的模型就是学了个寂寞。这一章把数据从原始标签变成网络能吃的掩膜每一步都落到能跑的代码上。2.1 先搞清楚你的标签是什么格式钢材缺陷分割的数据标签常见有三种来源标注工具导出的多边形 JSON、直接画好的 PNG 掩膜、或者 COCO 格式的 polygon 标注。你得先确认手里是哪种因为后面转换逻辑完全不同。多类别语义分割要求每张图对应一张单通道掩膜像素值就是类别索引0 背景、1 划痕、2 夹杂……而不是每种缺陷一张二值图。我一般先写个脚本统计类别分布和图像尺寸别急着训。这一步能提前暴露长尾问题——比如孔洞只有几十张划痕占了一大半。import os import numpy as np from PIL import Image from collections import Counter # 遍历掩膜目录统计每个类别像素占比和图像尺寸分布 mask_dir data/masks class_pixel Counter() size_dist Counter() for name in os.listdir(mask_dir): if not name.endswith(.png): continue m np.array(Image.open(os.path.join(mask_dir, name))) size_dist[m.shape] 1 # 掩膜像素值即类别索引直接统计 for cls, cnt in zip(*np.unique(m, return_countsTrue)): class_pixel[int(cls)] int(cnt) total sum(class_pixel.values()) for cls in sorted(class_pixel): print(fclass {cls}: {class_pixel[cls]} px, {class_pixel[cls]/total*100:.2f}%) print(size distribution:, size_dist.most_common(5))这段代码的关键在np.unique(m, return_countsTrue)它把掩膜里出现的所有像素值和数量一次统计出来。参数上没什么可调的但你要盯两个输出一是类别像素占比如果某个缺陷类低于 1%后面损失函数就得考虑加权二是尺寸分布如果图像尺寸五花八门训练前必须统一 resize 或裁剪否则 dataloader 会直接报错。2.2 多类别掩膜的生成与对齐如果标签是多边形 JSON你需要把它栅格化成掩膜。这里最容易踩的坑是坐标顺序和图像尺寸对不上——标注时用的显示尺寸和原图尺寸不一致栅格化出来整体偏移。常见做法是用标注里的 imageWidth/imageHeight 做归一化再乘回原图尺寸。import json import numpy as np from PIL import Image, ImageDraw # 把多边形标注栅格化成单通道类别掩膜 def poly_to_mask(ann_file, img_size, class_map): with open(ann_file, r, encodingutf-8) as f: data json.load(f) w, h img_size mask Image.new(L, (w, h), 0) # 0 为背景 draw ImageDraw.Draw(mask) for shape in data[shapes]: label shape[label] if label not in class_map: continue cls_idx class_map[label] # 标注点可能是归一化坐标先还原到像素 pts [(float(x), float(y)) for x, y in shape[points]] draw.polygon(pts, fillcls_idx) return np.array(mask) class_map {scratch: 1, inclusion: 2, oxide: 3, hole: 4} mask poly_to_mask(data/ann/001.json, (1280, 1024), class_map) Image.fromarray(mask).save(data/masks/001.png)逻辑说明Image.new(L, ...)建的是 8 位单通道图最多支持 255 类工业缺陷通常不超过 10 类够用。draw.polygon按类别索引填充后画的会覆盖先画的所以如果缺陷有重叠你得决定优先级——一般让面积小的缺陷后画避免被大缺陷吃掉。参数上class_map必须固定训练、验证、推理三处用同一份否则类别索引错位模型学出来的东西全是乱的。提示栅格化完一定要抽几张图叠在原图上肉眼核对尤其是细长划痕多边形点太少会把它画断。2.3 划分数据集时别按图像随机分工业图像常常是同一块板连续拍多张如果按图像随机划分训练验证集相邻帧会同时出现在两边验证指标虚高上线就露馅。我一般按板号或拍摄批次分组划分保证同一批次的图只进一个集合。import random from collections import defaultdict # 按批次分组划分避免同批次图像泄漏到验证集 groups defaultdict(list) for name in os.listdir(data/masks): batch_id name.split(_)[0] # 假设文件名前缀是批次号 groups[batch_id].append(name) batch_ids list(groups.keys()) random.seed(42) random.shuffle(batch_ids) split int(len(batch_ids) * 0.8) train_batches set(batch_ids[:split]) val_batches set(batch_ids[split:]) train [n for b in train_batches for n in groups[b]] val [n for b in val_batches for n in groups[b]] print(ftrain {len(train)}, val {len(val)})这里random.seed(42)是为了可复现工业项目里实验对比必须固定种子。split0.8是常见比例但 4100 张如果批次少验证集可能太小可以调到 0.85。分组划分会让验证集更接近真实泛化表现代价是指标可能比随机划分低几个点——这是好事说明你没自欺欺人。3. 语义分割模型选型钢材缺陷这种边界模糊场景该用谁数据理顺了接下来是选模型。钢材缺陷分割的难点在于缺陷细长、边界对比度低、类别间形态相似所以模型得同时有足够大的感受野和足够细的空间分辨率。这一章讲选型逻辑和最小可跑通训练。3.1 U-Net、DeepLabV3、SegFormer 在钢材场景的取舍先说结论4100 张这个量级我一般从 U-Net 起步想再压边界就上 DeepLabV3 的 ASPP追求速度和精度的平衡可以试 SegFormer 的轻量版。别一上来就上大 backbone数据不够过拟合比欠拟合更难救。模型优势钢材场景注意点数据量要求U-Net结构简单、边界恢复好编码器别用太深否则小缺陷丢中低DeepLabV3多尺度 ASPP适合不同宽度划痕空洞卷积输出步长要调不然细划痕断中SegFormerTransformer 全局建模速度快需要足够数据小数据要强增强中高U-Net 的跳跃连接对细长划痕特别友好因为浅层高分辨率特征能直接补回上采样丢掉的边界。DeepLabV3 的 ASPP 用不同膨胀率抓多尺度适合宽度差异大的缺陷但输出步长设成 16 时很细的划痕会断成几截我一般设成 8。SegFormer 在工业缺陷检测里近两年用得多它的分层结构对多类别分割友好但 4100 张要配强增强才稳。3.2 用 PyTorch 搭一个最小可训练分割网络下面是一个能直接跑的 U-Net 精简版编码器用预训练 ResNet 太占篇幅这里用纯卷积版本方便你先把流程跑通再换 backbone。import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.net 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.net(x) class UNet(nn.Module): def __init__(self, in_ch3, num_classes5): super().__init__() self.d1 DoubleConv(in_ch, 32) self.d2 DoubleConv(32, 64) self.d3 DoubleConv(64, 128) self.pool nn.MaxPool2d(2) self.up2 nn.ConvTranspose2d(128, 64, 2, stride2) self.u2 DoubleConv(128, 64) self.up1 nn.ConvTranspose2d(64, 32, 2, stride2) self.u1 DoubleConv(64, 32) self.out nn.Conv2d(32, num_classes, 1) def forward(self, x): c1 self.d1(x) c2 self.d2(self.pool(c1)) c3 self.d3(self.pool(c2)) x self.u2(torch.cat([self.up2(c3), c2], dim1)) x self.u1(torch.cat([self.up1(x), c1], dim1)) return self.out(x) model UNet(in_ch3, num_classes5) print(sum(p.numel() for p in model.parameters()) / 1e6, M params)逻辑说明DoubleConv是两次卷积加 BN 加 ReLU这是 U-Net 的基本砖块。下采样用MaxPool2d(2)上采样用ConvTranspose2dtorch.cat把上采样结果和对应编码器特征拼接这就是跳跃连接。num_classes5对应背景加四类缺陷必须和你掩膜的类别索引一致。参数量打印出来大概几 M4100 张够训不会一上来就过拟合。参数上in_ch3是 RGB 输入如果你用灰度图就改 1。num_classes千万别写错写少了会丢类别写多了浪费。BN 层在小 batch 下不稳定如果显存只够 batch2建议换成 GroupNorm。3.3 损失函数和类别不平衡的处理钢材缺陷里背景像素通常占 90% 以上直接用交叉熵模型会学会全预测背景也能拿高准确率。常见做法是交叉熵加 Dice 的组合损失Dice 对前景占比敏感能把小缺陷拉回来。import torch.nn.functional as F def dice_loss(logits, target, num_classes, eps1e-6): probs F.softmax(logits, dim1) target_onehot F.one_hot(target, num_classes).permute(0, 3, 1, 2).float() dice 0.0 for c in range(1, num_classes): # 跳过背景 p probs[:, c] t target_onehot[:, c] inter (p * t).sum() dice 1 - (2 * inter eps) / (p.sum() t.sum() eps) return dice / (num_classes - 1) def combined_loss(logits, target, num_classes): ce F.cross_entropy(logits, target) dc dice_loss(logits, target, num_classes) return ce dc # 权重可按需调dice_loss里跳过背景类只算前景这样小缺陷的梯度不会被背景淹没。eps防止除零。combined_loss里两项等权如果某类缺陷特别少可以给 Dice 更大权重比如ce 2*dc。注意target是长整型类别索引不是 one-hotF.cross_entropy要求这样。注意Dice 对空掩膜整张图没缺陷会不稳定如果数据里有纯背景图要么过滤要么在 dice 里加判断跳过。4. 训练、评估与推理把 4100 张数据真正跑起来模型和数据都齐了这一章讲训练循环、指标评估和推理落地。工业项目里训练只是中间一步评估和推理才是决定能不能用的关键。4.1 训练循环里必须盯住的三个量训练循环本身不复杂但工业分割有三个量必须每轮记录训练损失、验证集 mIoU、以及每个类别的 IoU。只看总 loss 会被背景主导看不出小缺陷学没学会。import torch from torch.utils.data import DataLoader def train_one_epoch(model, loader, optimizer, device, num_classes): model.train() total_loss 0.0 for img, mask in loader: img, mask img.to(device), mask.to(device) optimizer.zero_grad() logits model(img) loss combined_loss(logits, mask, num_classes) loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(loader) torch.no_grad() def evaluate(model, loader, device, num_classes): model.eval() inter torch.zeros(num_classes) union torch.zeros(num_classes) for img, mask in loader: img, mask img.to(device), mask.to(device) pred model(img).argmax(dim1) for c in range(num_classes): inter[c] ((pred c) (mask c)).sum().item() union[c] ((pred c) | (mask c)).sum().item() iou inter / (union 1e-6) return iou, iou.mean().item()evaluate里用混淆矩阵的思路累计每个类的交集和并集最后算 IoU。iou.mean()就是 mIoU但你要把iou数组打出来看每一类——如果背景 IoU 0.99、划痕 0.3说明模型基本没学会划痕。argmax(dim1)把 logits 变成类别索引和掩膜对齐。训练时optimizer我一般用 AdamW学习率 1e-3 起步配合余弦退火。batch size 看显存能到 8 就别用 2BN 更稳。4.2 评估指标mIoU 之外还要看边界mIoU 是语义分割的标准指标但钢材缺陷的边界质量它反映不出来。两条划痕一条边界糊但 IoU 还行一条边界准但细mIoU 可能差不多实际可用性差很多。我一般额外看边界 F1 或者直接抽图肉眼比对。import numpy as np from scipy.ndimage import binary_erosion def boundary_f1(pred, gt, num_classes, width2): # 用腐蚀差提取边界再算 F1 f1_list [] for c in range(1, num_classes): p (pred c) g (gt c) p_edge p ^ binary_erosion(p, iterationswidth) g_edge g ^ binary_erosion(g, iterationswidth) tp (p_edge g_edge).sum() fp (p_edge ~g_edge).sum() fn (~p_edge g_edge).sum() f1 2 * tp / (2 * tp fp fn 1e-6) f1_list.append(f1) return np.mean(f1_list)binary_erosion腐蚀后和原图异或得到边界带。width控制边界宽度一般 2 到 3 像素。这个指标对细缺陷敏感能补 mIoU 的盲区。注意它计算量比 IoU 大评估时抽样算就行不用每张都跑。4.3 推理部署从 PyTorch 到产线可用的输出推理阶段要把模型输出变成产线能用的东西缺陷掩膜、面积、类别。常见做法是导出 ONNX 再用推理引擎跑但先确保 PyTorch 推理结果对。torch.no_grad() def infer(model, img_tensor, device, num_classes, min_area50): model.eval() img_tensor img_tensor.to(device) logits model(img_tensor) pred logits.argmax(dim1).squeeze(0).cpu().numpy() results [] for c in range(1, num_classes): area int((pred c).sum()) if area min_area: # 过滤噪点 results.append({class: c, area: area}) return pred, resultsmin_area是面积阈值小于它的连通区域当噪点丢掉工业上很实用能压掉一堆假阳性。argmax后squeeze(0)去掉 batch 维。返回的pred是整张掩膜results是每个缺陷类的面积统计下游按面积判级。提示导出 ONNX 前先确认输入尺寸固定动态尺寸在部分推理引擎上会掉精度。5. 钢材缺陷分割避坑那些让我返工的坑这一章全是血泪经验每条都按现象、原因、解决写。工业分割项目里这些坑踩一次就够你加班一周。5.1 验证集 mIoU 很高上线一塌糊涂现象本地验证 mIoU 0.85产线跑起来缺陷漏检严重。原因验证集和训练集同批次划分图像高度相似模型记住了这批板的纹理换一批就崩。解决按批次或板号分组划分验证集必须来自不同批次指标会降但真实。5.2 细长划痕总是断成几截现象推理掩膜里一条连续划痕被切成好几段。原因下采样倍数太高输出步长 16 或 32细划痕在低分辨率特征里丢了。解决把输出步长降到 8或者用 U-Net 这种高分辨率跳跃连接的结构再不行就在损失里加边界权重。5.3 某个缺陷类 IoU 一直是 0现象训练日志里孔洞类 IoU 始终接近 0。原因孔洞样本太少或者掩膜里孔洞像素值写错了和背景混了。解决先查掩膜像素值分布确认类别索引没错再对稀有类加权重或过采样Dice 损失里给它更大系数。5.4 训练 loss 震荡不收敛现象loss 上下跳降不下去。原因学习率太大或者 batch 太小导致 BN 统计不稳。解决学习率降到 1e-4 试batch 太小就把 BN 换 GroupNorm或者用梯度累积凑等效 batch。5.5 推理速度跟不上产线节拍现象单张推理几百毫秒产线要求几十毫秒。原因模型太大或者没做推理优化。解决换轻量 backbone导出 ONNX 用推理引擎输入尺寸别盲目用原图按缺陷最小可辨尺寸缩。6. 把 4100 张用到极致小数据下的增强与半监督技巧4100 张在工业分割里不算多想再往上提点靠的不是换更大的模型而是把数据榨干。这一章讲几个我实际用过、性价比高的技巧。先说增强。钢材缺陷的增强不能乱来翻转、旋转一般安全但颜色抖动要小心——氧化铁皮的颜色本身就是判别特征抖太狠会把类别搞混。我一般用水平垂直翻转、90 度旋转、轻微缩放颜色只做很小幅度的亮度对比度扰动。import albumentations as A train_tf A.Compose([ A.HorizontalFlip(p0.5), A.VerticalFlip(p0.5), A.RandomRotate90(p0.5), A.RandomBrightnessContrast(brightness_limit0.1, contrast_limit0.1, p0.3), A.Resize(512, 512), ]) val_tf A.Compose([A.Resize(512, 512)])brightness_limit和contrast_limit都压到 0.1就是怕破坏颜色特征。Resize统一到 512训练验证必须一致。albumentations 对分割任务会自动同步图像和掩膜省得你自己写。再说半监督。4100 张里如果有大量无标签图可以用伪标签先用有标签数据训一个模型对无标签图推理取置信度高的像素当伪标签混进训练集再训一轮。关键是置信度阈值要卡高我一般用 0.9低了会引入噪声越训越歪。torch.no_grad() def pseudo_label(model, unlabeled_loader, device, num_classes, conf_th0.9): model.eval() pseudo [] for img, name in unlabeled_loader: img img.to(device) probs torch.softmax(model(img), dim1) conf, pred probs.max(dim1) # 只保留高置信度像素其余标为忽略 mask torch.where(conf conf_th, pred, torch.full_like(pred, 255)) pseudo.append((name, mask.cpu())) return pseudoconf_th0.9是经验值卡太高伪标签太少卡太低噪声大。255作为忽略索引训练时在损失里跳过这些像素。半监督不是万能有标签数据质量差的时候伪标签只会放大错误先把有标签的洗干净再说。最后一个技巧是测试时增强TTA推理时对图像做翻转、多尺度把多次预测平均。能涨一两个点代价是推理时间翻几倍产线节拍紧就别用。我自己做这类项目最大的教训是别在模型上反复折腾先把数据检查、分组划分、类别平衡这三件事做扎实往往比换个 backbone 涨得多。4100 张够你跑通全流程也够你踩完大部分坑剩下的就是耐心调。希望帮到你。本文还有配套的精品资源点击获取