简介本资源为钢材缺陷图像分割数据集面向从事工业质检、表面缺陷检测的深度学习开发者与图像分割学习者可直接用于训练与验证多类别分割模型。数据集已划分完毕训练集约2900张图片及对应mask验证集约1200张图片及对应mask标签共5类0为背景4种缺陷对应png像素值1、2、3、4类别定义可查看classes文件。压缩包共2000个文件以1273个png与725个jpg为主另含1个txt标签说明和1个py可视化脚本整体约102.78MB采用7z格式打包。配套脚本可随机抽取一张图片展示原图、GT图像及GT在原图上的蒙板效果并保存至当前目录便于快速核验标注质量。目前已有67人学习适合直接上手训练UNet、SwinUnet、TransUnet等分割网络省去数据清洗与划分环节。1. 钢材缺陷图像分割数据集从4100张标注图到可训练模型的落地路径拿到一个约4100张、已经标注处理完、可以直接训练的钢材缺陷图像分割数据集第一反应不应该是“赶紧跑个模型看看”而是先搞清楚这批数据到底能解决产线上的什么问题。钢材表面缺陷检测在热轧、冷轧、涂镀等环节都是刚需划痕、夹杂、氧化铁皮压入、结疤、孔洞这几类缺陷如果只做分类或检测框定位精度往往不够——质检员需要知道缺陷的精确轮廓来判定等级和切除范围。图像分割就是干这个的给每个像素打上类别标签输出缺陷的精确边界。这个数据集的价值在于省掉了最耗时的标注环节4100张带标签的图直接决定了你能在多类别分割任务上快速验证方案而不是卡在数据准备上。适合谁用产线质检算法工程师、做工业视觉落地的团队、以及想用真实工业数据跑通分割全流程的开发者。下面按“数据怎么验→模型怎么选→训练怎么跑→坑怎么避”的顺序拆开讲。2. 数据集验收与格式转换4100张图到手后先做这四件事2.1 为什么不能直接开训先搞清标注格式和类别分布“已处理完可以直接训练”这句话要打个问号。不同来源的分割数据集标注格式可能是 COCO JSON、VOC XML、YOLO txt、或者直接是掩码 PNG。你得先确认三件事图像尺寸是否统一、标注是多边形还是掩码、类别数量和名称是否明确。常见做法是写个脚本统计每类缺陷的像素占比和图像数量如果某一类只占 2% 以下训练时大概率会被模型忽略。钢材缺陷里划痕和夹杂往往样本多孔洞和结疤样本少这种长尾分布直接决定你要不要做重采样或损失加权。import os, json from collections import Counter from PIL import Image # 统计掩码类别分布假设标签是单通道灰度掩码像素值即类别ID mask_dir labels/ counter Counter() for f in os.listdir(mask_dir): if f.endswith(.png): mask Image.open(os.path.join(mask_dir, f)) counter.update(mask.getdata()) total sum(counter.values()) for cls_id, cnt in sorted(counter.items()): print(f类别 {cls_id}: {cnt} 像素, 占比 {cnt/total*100:.2f}%)这段代码遍历所有掩码图统计每个像素值的出现次数。参数说明mask_dir指向标签文件夹getdata()返回所有像素值。如果输出里某个类别占比低于 1%训练时就要在损失函数里给它更高权重或者用过采样把含该类的图多复制几份。注意如果标签是彩色掩码需要先做颜色到类别 ID 的映射不能直接统计 RGB 值。2.2 转成训练框架认的格式以 YOLO 分割和 Mask2Former 为例不同框架吃不同格式。YOLOv8/v13 的分割任务需要每张图对应一个 txt每行是类别 x1 y1 x2 y2 ...的归一化多边形坐标。Mask2Former 和 mmrotate 这类则更习惯 COCO 格式的 JSON。如果你拿到的已经是掩码 PNG转 YOLO 格式需要先提取轮廓。import cv2 import numpy as np def mask_to_yolo_polygon(mask_path, img_w, img_h, class_id0): mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) contours, _ cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) lines [] for cnt in contours: if cv2.contourArea(cnt) 10: # 过滤噪点 continue epsilon 0.002 * cv2.arcLength(cnt, True) approx cv2.approxPolyDP(cnt, epsilon, True) points approx.reshape(-1, 2).astype(float) points[:, 0] / img_w points[:, 1] / img_h coords .join([f{p:.6f} for p in points.flatten()]) lines.append(f{class_id} {coords}) return lines逻辑说明findContours提取外轮廓approxPolyDP做多边形拟合减少点数epsilon控制拟合精度值越小越贴近原始轮廓但点数越多。归一化坐标是 YOLO 的硬要求。参数说明cv2.contourArea过滤小于 10 像素的噪点区域这个阈值根据你的图像分辨率调整4100 张图如果分辨率在 640 左右10 比较合适。转换完记得抽查几张可视化确认多边形和原掩码重合。2.3 划分训练集验证集别让同一张图的缺陷泄漏到两边随机划分在分割任务里有个隐蔽的坑如果同一张原图被裁剪成多个 patch随机划分会让训练集和验证集出现高度相似的样本验证指标虚高。钢材缺陷数据集如果是整图标注按 8:1:1 随机划分问题不大如果是切片数据必须按原图 ID 分组划分。常见做法是用sklearn的GroupShuffleSplit把原图文件名前缀作为 group。from sklearn.model_selection import GroupShuffleSplit import os files sorted(os.listdir(images/)) groups [f.split(_)[0] for f in files] # 假设文件名格式为 原图ID_切片号.jpg gss GroupShuffleSplit(n_splits1, test_size0.2, random_state42) train_idx, val_idx next(gss.split(files, groupsgroups)) print(f训练集 {len(train_idx)} 张, 验证集 {len(val_idx)} 张)参数说明test_size0.2表示验证集占 20%random_state固定后结果可复现。如果数据没有切片直接用train_test_split即可。注意划分后要检查两个集合的类别分布是否接近如果验证集里某类缺陷一张都没有指标会失真。3. 模型选型与训练环境从 YOLOv8-seg 到 Mask2Former 的取舍3.1 实时产线用 YOLOv8-seg离线高精度用 Mask2Former选模型先看推理延迟要求。产线在线检测通常要求单张图 50ms 以内YOLOv8-seg 的 nano 或 small 版本在 640 分辨率下能跑到 30-40ms取决于 GPU精度对钢材缺陷这种边界相对清晰的场景够用。如果做离线抽检或科研分析Mask2Former 的掩码质量明显更好尤其是小目标和重叠缺陷但推理速度慢一个数量级。mmrotate 更适合带旋转框的缺陷检测但分割任务用它属于绕路不推荐。我的建议先用 YOLOv8-seg 跑通全流程拿到 baseline再根据漏检情况决定要不要换重模型。模型输入尺寸推理延迟RTX 3060适合场景训练显存YOLOv8n-seg640~25ms产线实时4GBYOLOv8s-seg640~40ms产线实时6GBMask2Former (Swin-T)512~200ms离线高精度8GBDeepLabV3 (ResNet50)512~80ms折中方案6GB3.2 训练环境搭建CUDA、PyTorch 和 ultralytics 的版本对齐环境翻车是最高频的问题。YOLOv8 依赖 ultralytics 包它和 PyTorch 版本有绑定关系。常见做法是先用 conda 建一个干净环境装 PyTorch 时指定 CUDA 版本再装 ultralytics。不要混用 pip 和 conda 装同一个包否则 CUDA 库路径会乱。conda create -n steel_seg python3.10 -y conda activate steel_seg pip install torch2.1.0 torchvision0.16.0 --index-url https://download.pytorch.org/whl/cu118 pip install ultralytics8.1.0 python -c import torch; print(torch.cuda.is_available())参数说明cu118对应 CUDA 11.8如果你的驱动只支持 CUDA 12.x把 index-url 换成cu121。最后一行验证 GPU 是否可用输出True才算环境通了。注意ultralytics 版本不要追最新8.1.x 在分割任务上比较稳定新版有时会改数据加载逻辑导致旧格式报错。3.3 配置文件怎么写data.yaml 的五个关键字段YOLOv8-seg 训练需要一个 data.yaml里面五个字段必须写对path是数据集根目录train和val是图像文件夹的相对路径nc是类别数names是类别名列表。路径写错是最常见的翻车点建议用绝对路径先跑通再改相对路径。path: /data/steel_defect train: images/train val: images/val nc: 5 names: [scratch, inclusion, rolled_scale, pitted_surface, crazing]注意names的顺序必须和标注里的类别 ID 一致如果标注里 0 是划痕、1 是夹杂这里就不能写反。训练前用yolo checks命令检查数据集配置它会提示路径和标签格式问题。4. 训练参数调优与显存控制让 4100 张图跑出稳定指标4.1 批次大小和图像尺寸的显存博弈4100 张图不算多但分割任务的显存占用比检测高因为要保留掩码分支。RTX 3060 12GB 上YOLOv8s-seg 用 640 分辨率、batch8 大概占 6-7GB。如果显存不够优先降 batch 而不是降分辨率——分辨率降低对小缺陷的掩码质量影响很大。常见做法是先用 batch4 跑几个 epoch 看显存峰值再逐步往上加。yolo segment train datadata.yaml modelyolov8s-seg.pt epochs100 imgsz640 batch8 device0 patience20参数说明epochs100配合patience20表示 20 个 epoch 指标不升就早停避免过拟合。device0指定第一块 GPU。如果显存溢出把batch降到 4同时把workers从默认 8 降到 4减少数据加载进程的内存开销。4.2 学习率和损失权重钢材缺陷小目标需要更高分辨率YOLOv8-seg 默认学习率 0.01对 4100 张图的数据量偏大容易在后期震荡。我一般改成 0.001 配合余弦退火前 3 个 epoch 用 warmup 让模型稳定。如果某类缺陷像素占比极低可以在损失里给分割分支更高权重但 YOLOv8 不直接暴露这个参数替代方案是把imgsz提到 1024让小目标在特征图上占更多像素。from ultralytics import YOLO model YOLO(yolov8s-seg.pt) model.train( datadata.yaml, epochs100, imgsz1024, batch4, lr00.001, lrf0.01, warmup_epochs3, cos_lrTrue, patience20, device0 )参数说明lr0是初始学习率lrf是最终学习率比例cos_lrTrue启用余弦退火。imgsz1024会显著增加显存占用batch 要相应降到 4 或 2。注意提高分辨率后训练时间会翻倍先跑 10 个 epoch 看指标趋势再决定要不要跑满。4.3 数据增强钢材缺陷不能乱用翻转和色彩抖动通用增强策略在工业缺陷上不一定适用。水平翻转对划痕方向敏感的场景可能引入错误标签色彩抖动会改变氧化铁皮的颜色特征。我的经验是只开mosaic和scale关掉hsv_h、hsv_s、hsv_v的大幅抖动翻转概率设 0.5 以下。如果缺陷有明确方向性比如轧制方向的划痕翻转概率设 0。# 在 data.yaml 同级目录建 hyp.yaml mosaic: 0.5 scale: 0.3 fliplr: 0.3 flipud: 0.0 hsv_h: 0.0 hsv_s: 0.0 hsv_v: 0.0训练时用hyphyp.yaml加载。参数说明mosaic0.5表示 50% 概率做四图拼接scale0.3允许 30% 缩放。注意关掉 HSV 抖动后模型对光照变化的鲁棒性会下降如果产线光照不稳定保留hsv_v0.2左右的小幅亮度变化。5. 避坑与排查钢材缺陷分割训练中最容易翻车的五个点5.1 损失不下降mask loss 一直是 nan现象训练几个 epoch 后seg_loss变成 nan模型输出全零掩码。原因标注里有空掩码或全零标签YOLO 在计算分割损失时对空目标做除法导致 nan。解决训练前过滤掉没有标注的图或者把空标签图从训练集移除。检查方法遍历 labels 文件夹如果某个 txt 文件为空或所有行都是 0 像素面积直接删掉对应图。5.2 验证集指标很高但实际推理一塌糊涂现象验证集 mAP50 到 0.9拿产线图测试却什么都分不出来。原因训练集和验证集来自同一批数据分布一致但产线图的光照、背景、缺陷形态和数据集差异大。解决从产线另找 50-100 张图做独立测试集不要混入训练。如果差异确实大用增量训练在产线图上微调几个 epoch。5.3 显存溢出但 batch 已经降到 1现象CUDA out of memorybatch 降到 1 还是报错。原因imgsz太大或者workers太多导致数据加载进程占满内存。解决先把imgsz降到 640workers设 2关掉cacheTrue。如果还不行检查是不是加载了预训练模型后没有释放旧模型在训练脚本开头加torch.cuda.empty_cache()。5.4 类别不平衡导致小类完全漏检现象划痕和夹杂分得很好孔洞和结疤一个都检不出。原因小类样本太少模型倾向于预测大类。解决两个方向——一是过采样含小类的图把含孔洞的图复制 3-5 份二是在验证时单独看小类的混淆矩阵确认是漏检还是误分类。如果漏检严重考虑用 Mask2Former 的 focal loss 替代交叉熵。5.5 训练完导出 ONNX 后推理结果和 PyTorch 不一致现象PyTorch 下掩码正常导出 ONNX 后用 onnxruntime 推理掩码偏移或类别错乱。原因YOLOv8-seg 的导出需要指定opset12以上且后处理里的 sigmoid 和 argmax 在 ONNX 里可能被优化掉。解决导出时加simplifyTrue并用官方yolo export命令而不是手动转。yolo export modelruns/segment/train/weights/best.pt formatonnx opset12 simplifyTrue导出后用onnxruntime跑一张测试图和 PyTorch 输出做逐像素对比差异超过 1% 就要检查预处理是否一致归一化参数、通道顺序。6. 进阶技巧用测试时增强和增量训练把 mAP 再提 3-5 个点训练完 baseline 后如果指标卡在某个水平上不去有两个性价比最高的技巧。第一个是测试时增强TTA推理时对同一张图做水平翻转、多尺度缩放把多次预测的掩码做平均。YOLOv8 自带augmentTrue参数在model.predict里开启即可。代价是推理时间翻 3-4 倍适合离线抽检场景。results model.predict( sourcetest_images/, augmentTrue, # 开启 TTA conf0.25, iou0.5, imgsz1024, saveTrue )参数说明augmentTrue会做水平翻转和缩放增强conf和iou是后处理阈值。注意TTA 对小目标的提升最明显如果缺陷本身很大提升有限但耗时增加自己权衡。第二个技巧是增量训练。当产线新增了一批标注图不要从头训用model.train(resumeTrue)或者加载best.pt继续训 20-30 个 epoch学习率设小一个数量级。这样既能吸收新数据又不会遗忘旧类别。我一般会保留一个holdout集每次增量后在上面验证如果旧类别指标掉超过 2 个点就混入 20% 旧数据一起训。还有一个容易被忽略的点掩码后处理。YOLOv8-seg 输出的掩码是 160x160 的低分辨率上采样到原图后边缘会有锯齿。如果下游需要精确面积计算用cv2.morphologyEx做开闭运算平滑边缘或者用cv2.resize加双三次插值。这个操作在 CPU 上做一张图增加 5-10ms但对面积测量精度提升明显。最后说个血泪教训别在训练中途改data.yaml的类别顺序。我有一次把names里两个类别调换后继续训模型直接崩了损失从 0.3 跳到 2.0因为标签 ID 和模型输出层的对应关系乱了。改类别必须从头训或者至少重置分类头。希望帮到你。本文还有配套的精品资源点击获取