ARTICLE DETAIL

资讯详情

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

YOLOv5钢材表面缺陷检测:NEU-DET数据集训练与部署实战

YOLOv5钢材表面缺陷检测:NEU-DET数据集训练与部署实战 简介面向工业质检场景的YOLOv5钢材表面缺陷检测系统源码及配套数据集适合深度学习和机器视觉学习者、算法工程师用于模型训练与调优。压缩包共5528个文件约28.28MB包含1800张钢材表面jpg图像及对应1800个xml标注文件另配1806个txt标签与PyTorch、YOLOv5所需的py脚本、yaml配置文件可覆盖数据标注、模型训练、推理检测全流程。已有2514人学习下载资源内还提供标签含义说明、参数配置等辅助文件便于快速理解缺陷类别与训练参数。使用者可直接基于该数据集开展YOLOv5训练实验验证数据增强、迁移学习等策略在钢材缺陷识别上的效果也可借鉴源码结构构建自己的检测系统。1. YOLOv5钢材表面缺陷检测六类缺陷一条训练推理链路全齐无论是钢带还是热轧钢板表面缺陷检测都是出厂前必过的一关。人工质检靠肉眼盯着传送带漏检率随着疲劳度直线上升而且裂纹、麻点、氧化铁皮这些缺陷形态接近新老检验员的判定经常打架。这套YOLOv5钢材表面缺陷检测系统源码包把“数据集 标注 训练 推理”整条链路封装在一起模型基于YOLOv5数据采用NEU-DET公开钢材缺陷集覆盖六类典型缺陷训练和检测代码拿到就能用。适合两类人——做工业视觉项目需要快速出demo的工程师以及想拿一个真实场景练手目标检测的学生。拿到源码后不需要从零搭网络重点是把数据吃透、把参数调明白。2. 把 NEU-DET 转成 YOLO 格式六类缺陷、标签转换与训练集划分2.1 NEU-DET 的构成六类缺陷各 300 张NEU-DET 是东北大学发布的钢材表面缺陷公开数据集。图像是灰度图尺寸约200×200总共1800张缺陷分为六类每类正好300张类别分布非常均匀缺陷类别英文标签数量形态特征裂纹crazing300表面细密网状裂纹对比度低夹杂物inclusion300点状暗色颗粒局部发黑斑点patches300片状亮斑边界模糊麻点pitted_surface300细小凹坑群呈点状分布氧化铁皮rolled-in_scale300条状压入的氧化皮层划痕scratches300连续细线方向随机对目标检测来说1800张图属于偏小的数据量六类分布均匀算是个安慰。钢材表面的缺陷有个显著特点——同类之间外观差异很大。比如“裂纹”可能是几根细短线也可能是一片密集网格“划痕”有时清晰连贯有时时断时续。这意味着模型要学的是“纹理异常”这种抽象概念而不是固定形状。所以用这个数据集做出来的模型mAP50在0.75到0.85算正常范围。如果看到有人报0.95先怀疑他是不是用训练集图像做了评估。2.2 标注转换从 XML 到 YOLO TXT归一化是关键公开的NEU-DET目标检测版本标注格式常见的有两种XMLVOC格式和MAT格式。YOLOv5需要的是与每张图同名的txt文件每一行代表一个目标class_id center_x center_y width height其中center_x、center_y是目标中心点相对图像宽高的比例width、height是目标宽高相对图像尺寸的比例四个值都在0到1之间。拿到XML标注后我习惯用下面这段脚本统一转import os import cv2 import xml.etree.ElementTree as ET CLASS_MAP { crazing: 0, inclusion: 1, patches: 2, pitted_surface: 3, rolled-in_scale: 4, scratches: 5 } def voc_to_yolo(xml_dir, img_dir, out_dir): os.makedirs(out_dir, exist_okTrue) error_files [] for xml_name in sorted(os.listdir(xml_dir)): if not xml_name.endswith(.xml): continue base xml_name.replace(.xml, ) xml_path os.path.join(xml_dir, xml_name) img_path os.path.join(img_dir, base .jpg) # 用cv2读图拿真实宽高做归一化 img cv2.imread(img_path) if img is None: error_files.append(base) continue img_h, img_w img.shape[:2] tree ET.parse(xml_path) root tree.getroot() lines [] for obj in root.findall(object): cls obj.find(name).text if cls not in CLASS_MAP: continue bnd obj.find(bndbox) x1 float(bnd.find(xmin).text) y1 float(bnd.find(ymin).text) x2 float(bnd.find(xmax).text) y2 float(bnd.find(ymax).text) cx (x1 x2) / 2.0 / img_w cy (y1 y2) / 2.0 / img_h w (x2 - x1) / img_w h (y2 - y1) / img_h lines.append(f{CLASS_MAP[cls]} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}) with open(os.path.join(out_dir, base .txt), w) as f: f.write(\n.join(lines)) print(f转换完成失败 {len(error_files)} 个文件: {error_files}) voc_to_yolo(path/to/xml, path/to/images, path/to/labels)这段代码的逻辑很简单循环读XML把每个目标的左上右下坐标换算成中心点加宽高再除以图像宽高做归一化。为什么要归一化因为YOLOv5训练时会做letterbox缩放把图从200×200垫成640×640中间加灰边。如果标注存的是绝对像素缩放后框就错位了归一化坐标在不同分辨率下通用这也是YOLO格式能在所有目标检测框架里传递的原因。注意img.shape返回的顺序是高度在前、宽度在后接的时候别把img_w和img_h接反。脚本末尾打印失败文件列表用来排查损坏图像。2.3 训练集划分与数据增强小数据集最怕划不干净NEU-DET总共就1800张划分方式直接影响最终评估的可信度。这里有一个容易忽略的点NEU-DET的图片是按区域采集的同一块钢板表面的多张截图像如果同时出现在训练集和验证集里验证集mAP会虚高换到真实产线上性能立刻掉一截。import os import random img_dir images_all val_ratio 0.2 random.seed(42) imgs [f for f in os.listdir(img_dir) if f.endswith(.jpg)] random.shuffle(imgs) val_count int(len(imgs) * val_ratio) val_files set(imgs[:val_count]) # 生成 YOLOv5 需要的 images 与 labels 目录结构 for split, subset in [(train, imgs[val_count:]), (val, val_files)]: os.makedirs(fdataset/steel/images/{split}, exist_okTrue) os.makedirs(fdataset/steel/labels/{split}, exist_okTrue) moved 0 for name in subset: src os.path.join(img_dir, name) dst os.path.join(fdataset/steel/images/{split}, name) os.rename(src, dst) # 实际项目里建议用copy保证原始数据完整 label_name name.replace(.jpg, .txt) if os.path.exists(os.path.join(labels_all, label_name)): os.rename( os.path.join(labels_all, label_name), os.path.join(fdataset/steel/labels/{split}, label_name) ) moved 1 print(f{split}: {moved} 张)random.seed(42)固定随机种子别人复现的时候拿到的划分和你一致。20%做验证集也就是1440张训练、360张验证对YOLOv5s来说够用。真实项目里建议改成copy而不是rename原始NEU-DET数据保留一份后面想重新划分或补数据时还能复用。数据增强方面YOLOv5默认开启mosaic增强把四张图拼成一张训练对小目标召回有好处。钢材缺陷是灰度纹理hsv色调增强作用不大但旋转和缩放要克制划痕旋转90度后特征含义会变翻转建议只开水平翻转角度增强控制在±10度以内。3. YOLOv5训练实操conda 环境、data.yaml 与收敛判断3.1 环境版本怎么锁torch、CUDA 与 YOLOv5 的搭配YOLOv5源码对Python版本不算挑剔但torch、torchvision、CUDA之间必须配套版本错一个就可能出现训练时找不到GPU或者前向推理结果全为nan。我反复用过一套稳定组合Python 3.9 torch 1.13.1 torchvision 0.14.1 CUDA 11.730系和40系显卡都能跑。conda create -n steel_yolo python3.9 -y conda activate steel_yolo pip install torch1.13.1 torchvision0.14.1 --index-url https://download.pytorch.org/whl/cu117 pip install -r requirements.txt--index-url参数指定的是PyTorch官方预编译轮子的下载源cu117后缀对应CUDA 11.7。如果不指定pip可能装成CPU版本train.py启动后一直提示找不到GPU。requirements.txt里的opencv-python、numpy、matplotlib、pyyaml、tqdm、seaborn按YOLOv5官方清单装就能过。装完先验证一下环境python -c import torch; print(torch.__version__, torch.cuda.is_available())输出类似“1.13.1cu117 True”。如果输出False说明装成了CPU版回到上一步换--index-url重新装。这一步看似多余但环境问题占训练翻车原因的三成以上花10秒确认很划算。3.2 data.yaml 与训练命令目录结构是最大隐藏坑数据集目录建议直接按YOLOv5约定组织然后写一个对应的yaml# steel_data.yaml train: ./dataset/steel/images/train val: ./dataset/steel/images/val nc: 6 names: [crazing, inclusion, patches, pitted_surface, rolled-in_scale, scratches]YOLOv5会自动根据train的路径推导labels路径images/train对应labels/trainimages/val对应labels/val。也就是说Labels目录必须和images目录放在同一级父目录下且子目录名保持一致。很多人在这里翻车——图片路径配对了但labels目录结构不对启动训练时日志刷出大量“No labels found in ...”训练等于白跑。训练命令建议这样写python train.py \ --weights yolov5s.pt \ --data steel_data.yaml \ --epochs 100 \ --batch-size 16 \ --imgsz 640 \ --device 0 \ --name steel_s \ --patience 20每次训练会在runs/train/下生成一个steel_s目录保存权重和日志。epoch设为100对NEU-DET这种规模足够YOLOv5s在单张1080Ti上从头训练大约一到两小时主要瓶颈在数据集读取速度。--weights yolov5s.pt是COCO预训练权重让模型先学到通用特征再迁移到钢材缺陷比从零初始化好很多。--batch-size 16看你显卡显存不够就降到8或4不用硬凑。--imgsz 640是输入分辨率NEU-DET原图才200×200letterbox会把它垫到640×640如果显存紧张降到416但裂纹、划痕本来就细分辨率太低容易漏。3.3 训练日志与收敛判断mAP50 是首要指标训练过程中YOLOv5实时打印三类lossbox_loss是框回归误差cls_loss是分类误差obj_loss是置信度误差。此外每轮结束后会输出验证集的mAP50和mAP50-95。我最关心的是mAP50。NEU-DET六类缺陷的检测难度差异大mAP50能到0.8模型就算可用。判断是否过拟合有一个简单方法如果训练集loss还在降但验证集mAP50连续10个epoch不涨甚至下跌就说明模型开始背训练集了。此时回到数据增强把hyp.scratch-low.yaml里的degrees调低、hsv_h和hsv_s调低或者干脆把mosaic增强的比例加大。YOLOv5里开启mosaic的参数是mosaic1.0默认就是开的。另一个判断依据是训练前几个epoch的loss下降速度。如果前5个epoch的box_loss几乎不动大概率是anchor设置跟目标尺寸不匹配或者是标签文件本身有问题。下一章讲具体的坑。4. 避坑指南钢材表面缺陷训练最常见的五个翻车点4.1 灰度图读入变单通道训练直接崩现象train.py启动时报错提示“Expected 3 input channels, got 1”或者loss输出为nan。原因NEU-DET的图是灰度图单通道。YOLOv5官方datasets.py里用了cv2.imread会自动转成三通道没问题。但如果你自己改了数据读取逻辑或者用PIL打开图片再转成numpy数组就很容易拿到(H, W)的单通道数组喂给模型自然报错。解决统一用cv2.imread读图并在读图后检查维度import cv2 img cv2.imread(img_path) # 灰度图cv2也会复制为3通道 if img is None: print(f图片读取失败: {img_path}) continue if len(img.shape) 2: img cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)cv2.imread对灰度图会自动复制成三通道这行代码本身就能规避大部分通道问题。新手容易在这里直接翻车老手则在自定义数据集时犯同样的错。4.2 标签目录结构不对训练刷满屏警告现象训练启动正常但日志里大量出现“WARNING: No labels found in ... skipping”。原因YOLOv5根据train和val的images路径自动推导labels路径你的目录层级一旦错位它就找不到txt文件。常见的是把labels放在dataset/steel/labels/train但images放在dataset/steel/images/train而yaml里train写的却是dataset/steel/images导致labels的实际路径与推导路径不匹配。解决写个小脚本核对目录结构python -c import os for split in [train, val]: img_d fdataset/steel/images/{split} lab_d fdataset/steel/labels/{split} if not os.path.isdir(lab_d): print(f缺失目录: {lab_d}) else: imgs set(os.listdir(img_d)) labels set(os.listdir(lab_d)) miss sum(1 for f in imgs if f.replace(.jpg, .txt) not in labels) print(f{split}: 图片{len(imgs)}张, 标签{len(labels)}个, 缺失{miss}个) YOLOv5的路径推导机制是取train路径的父目录把images替换成labels再拼上同名的子目录后缀。所以把images/train、labels/train、images/val、labels/val建在同一父目录下就不会出问题。4.3 数据量小过拟合mAP50 虚低现象训练集loss一路降到0.02但验证集mAP50始终在0.5以下。原因1440张训练图对YOLOv5s来说偏少。模型容量大很快把训练集细节背了下来验证集上泛化不足。解决优先用预训练权重迁移yolov5s.pt的COCO通用特征能显著缓解数据量不足。其次把数据增强强度提上来开启mosaic把hsv_h调到0.02以下degrees限制在5到10之间。钢材缺陷是纹理问题旋转幅度过大反而会让模型学到错误的纹理方向尤其是scratches。如果还想更保守换yolov5n权重模型更小过拟合风险更低。4.4 anchor 尺寸与缺陷形态不匹配现象前10个epoch的box_loss几乎不动mAP50一直为0。原因YOLOv5默认anchor是在COCO数据集上聚类得到的COCO里大目标占比高而NEU-DET的缺陷在200×200原图里往往只有几十个像素。默认anchor尺度偏大模型预测的框根本够不到小而细长的缺陷。解决YOLOv5在训练时会自动重新聚类anchor启动日志里如果看到“Analyzing anchors...”说明它正在算。如果算完anchor还是不合适手动指定也行。有一个不推荐但常见的做法是训练时加--noautoanchor强制用COCO的默认anchor结果就是loss不降。不要用这个参数自动anchor是YOLOv5专门给这类小目标场景准备的机制让它跑完就行。4.5 六类缺陷难度不均个别类 mAP 偏低现象训练结束后rolled_in_scale的mAP能到0.9crazing却只有0.3。原因rolled_in_scale是块状目标边界清晰特征稳定crazing是网状细纹对比度低同类之间差异极大。六类样本数量相同但学习难度完全不同模型优化时自然会偏向更容易学、loss下降更快的类别。解决对困难类做过采样。具体操作是从训练集中找出所有含crazing的图额外做一次随机裁剪放大2倍把裁剪后的图加入训练集。这样crazing的有效训练样本数变成原来的1.5倍左右模型有更多机会学习它的纹理模式。这个逻辑同样适用于scratches效果通常能拉高5个百分点以上。5. 效果验证与部署从 best.pt 到漏检率统计的最后一公里5.1 用 best.pt 跑推理阈值怎么设训练结束后权重在runs/train/steel_s/weights/下。best.pt是验证集mAP最高的那一次保存的权重last.pt是最后一个epoch的权重推理和部署一律用best.pt。推理命令python detect.py \ --weights runs/train/steel_s/weights/best.pt \ --source data/test \ --imgsz 640 \ --conf-thres 0.35 \ --save-txt--conf-thres是置信度阈值0.35对钢材缺陷是起步值。缺陷目标小、衬度低如果可视化结果里正常缺陷没画全就把阈值降到0.25代价是误检框会变多。建议先出0.35的结果看漏检还是误检占主导再决定往哪边调。--save-txt会把检测框坐标也保存下来后面做漏检率统计要用。5.2 漏检率统计别只看 mAPmAP是全局指标但工业质检关心的是“这批钢板里到底漏了几个缺陷”。我常用一个简单的匹配脚本把检测结果和真值标注做IoU匹配import os def cxcywh2xyxy(box): cx, cy, w, h box return [cx - w / 2, cy - h / 2, cx w / 2, cy h / 2] def iou(a, b): a cxcywh2xyxy(a) b cxcywh2xyxy(b) ix1, iy1 max(a[0], b[0]), max(a[1], b[1]) ix2, iy2 min(a[2], b[2]), min(a[3], b[3]) inter max(0, ix2 - ix1) * max(0, iy2 - iy1) area_a (a[2] - a[0]) * (a[3] - a[1]) area_b (b[2] - b[0]) * (b[3] - b[1]) return inter / (area_a area_b - inter 1e-6) def eval_miss(pred_file, gt_file, iou_thresh0.5): preds, gts [], [] for line in open(pred_file).read().strip().splitlines(): parts line.split() preds.append([int(parts[0])] [float(x) for x in parts[1:5]]) for line in open(gt_file).read().strip().splitlines(): parts line.split() gts.append([int(parts[0])] [float(x) for x in parts[1:5]]) if not gts: return 0, 0 miss 0 for gt in gts: hit False for pred in preds: if pred[0] gt[0] and iou(pred[1:], gt[1:]) iou_thresh: hit True break if not hit: miss 1 return miss, miss / len(gts) miss, rate eval_miss(pred.txt, gt.txt) print(f漏检 {miss} 个目标, 漏检率 {rate:.2%})IoU阈值为0.5类别编号一致才算命中。漏检率计算是“没被命中的真值框数除以总真值框数”。生产场景通常要求漏检率低于3%如果达不到就把漏检案例导出来看是哪类缺陷、面积多大针对性补数据或用第4章讲过的过采样策略。这套脚本虽然短但比全局mAP更能暴露模型在真实工况下的短板。5.3 部署到 ONNX注意输入分辨率一致性模型验证通过之后部署导出是下一步。YOLOv5导出ONNX的成本很低python export.py \ --weights runs/train/steel_s/weights/best.pt \ --include onnx \ --imgsz 640一个关键约束导出时的--imgsz必须跟训练时一致。训练用的640导出也得是640。如果训练用的是416导出时改成640检测框的尺度会在NMS阶段错位明明模型没问题部署端结果却全线漂移。从那以后我每训练一个工业检测模型验收清单里都会强制加两项不同conf-thres下的漏检率曲线以及ONNX与PyTorch推理结果的一致性比对。这两步走完模型才算真正能上产线。希望这套流程对你做钢材缺陷检测也有帮助。本文还有配套的精品资源点击获取
返回列表