基于YOLOv5的钢材表面缺陷检测系统实战：从NEU-DET训练到部署

简介：目标检测是计算机视觉的核心任务之一，旨在同时完成目标定位与分类。在工业质检场景中，深度学习目标检测技术正逐步替代传统人工目检和图像处理方法，解决钢材表面裂纹、夹杂、麻点等缺陷难以稳定识别的问题。YOLOv5作为经典的单阶段检测算法，在速度与精度之间取得了良好平衡，配合公开数据集NEU-DET，可快速构建一套从数据准备、模型训练到推理部署的完整检测系统。本文围绕一套可运行的源码，系统讲解了VOC标注转YOLO格式、数据增强策略、训练超参数调优、模型评估指标及检测器封装等关键环节，并总结了针对小目标缺陷的优化思路和TensorRT加速部署经验，为工业视觉检测及毕业设计项目提供具有工程价值的落地参考。 钢材表面缺陷检测这类项目，最近聊的人越来越多，前阵子我整理了一套基于YOLOv5的钢材表面缺陷数据集检测系统源码，从数据准备、模型训练到推理检测全链路跑通，用NEU-DET这个公开数据集来做的。这篇文章就围绕这套源码，把整个系统的设计思路、关键代码、训练细节和踩过的坑完整梳理一遍，给正在做工业质检、毕业设计或者想入门目标检测实战的朋友一个可以直接参考的完整方案。

先说清楚这套东西解决什么问题。传统钢材表面质检主要靠人工肉眼，连续盯产线几个小时，疲劳造成的漏检率根本压不住；换到传统图像处理，阈值分割、边缘检测那套方法，光照一变、表面带点油污水渍，算法就翻车。深度学习目标检测尤其是YOLOv5这一类单阶段检测器，能同时兼顾速度和精度，在GPU上跑实时推理完全没问题，所以才成了当前钢企质检智能化改造的主流方向。

1. 项目背景与整体设计思路

1.1 为什么必须换掉人工眼检

钢材表面缺陷不是一个抽象概念，它是轧制过程中各种工艺问题留下的具体痕迹。NEU-DET数据集里定义了六类典型缺陷：裂纹（crazing）、夹杂（inclusion）、斑块（patches）、麻点（pitted_surface）、氧化铁皮（rolled-in_scale）、划痕（scratches）。

这六类缺陷形态差异很大。裂纹是线状的，亮度忽高忽低，和背景对比度低到怀疑人生；夹杂是块状暗斑，尺寸变化大；斑块和麻点容易混淆；氧化铁皮表面纹理复杂，本身就有鳞片状背景干扰。人工目检疲劳后最容易漏掉的就是那种低对比度的细小裂纹，而这恰恰是影响钢材等级的关键缺陷。传统CV方案很难用一个固定模板去覆盖六类形态差异这么大的目标，这也是我选深度学习来做这块的根本原因。

1.2 技术选型：YOLOv5 为什么还是这么能打

YOLOv8、YOLOv9、YOLOv10都迭代到新版本了，为什么我还要回头用YOLOv5做这个项目？看起来有点过时，但实际选型不是追新，是看性价比。

YOLOv5目前生态最成熟。GitHub上star数量大，教程遍地，遇到问题几乎都能搜到答案。工业项目最怕那种网上资料少、报错没人踩过的坑，YOLOv5基本不存在这个问题。而且它的代码结构非常清晰，模型定义在models下，数据集配置在data下，超参数在data/hyps下，训练和推理入口分别是train.py和detect.py，拿到源码改改就能跑。相比之下YOLOv8的代码抽象度高一些，很多细节被封装了，二次修改反而费劲。

从精度角度说，NEU-DET数据集每类只有300张图，属于典型的小样本工业场景。YOLOv5s作为基准模型，mAP@0.5做到0.75左右问题不大，换YOLOv5m或YOLOv5l能再往上提4到6个百分点，这个水平已经能覆盖实际产线初筛需求了。再往上换更大的模型，FPS掉得厉害，工业现场部署还得掂量算力成本。

1.3 系统整体架构与源码目录规划

这套系统的源码组织结构，从压缩包解开之后是下面这个样子：

steel-defect-detection/ ├── yolov5/ # YOLOv5 官方主框架 │ ├── train.py # 训练入口 │ ├── detect.py # 推理入口 │ ├── models/ # 网络结构定义 │ ├── data/ # 数据集配置与超参数 │ └── utils/ # 工具函数 ├── tools/ │ ├── voc2yolo.py # VOC格式转YOLO格式 │ └── split_dataset.py # 数据集划分脚本 ├── datasets/ │ └── NEU-DET/ # 钢材缺陷数据集 ├── runs/ │ ├── train/ # 训练输出 │ └── detect/ # 推理输出 ├── inference/ │ ├── detector.py # 封装的检测器 │ └── batch_detect.py # 批量检测脚本 └── requirements.txt

把yolov5官方仓库原封不动放进来，自己写的代码放在tools和inference里，这样做的好处是后续YOLOv5升级时，只需要替换官方部分，自己的代码完全不用动。这是这种工程项目的经典做法，绝不修改底层源码，而是通过外层封装来扩展功能。

2. 数据集准备与预处理

2.1 NEU-DET数据集的规格和六类缺陷

NEU-DET是公开的钢材表面缺陷标准数据集，一共1800张灰度图像，六类缺陷每类300张，每张图是200x200像素。

有个细节值得注意：图像是灰度图，但YOLOv5训练默认走RGB三通道。处理方式有两个，一是在数据加载时把灰度图复制成三通道，这符合大多数人的操作习惯；二是自定义单通道数据集，但需要改模型结构和数据加载代码，工作量大了不少。我实际做的时候选择复制成三通道，因为200x200的图分辨率本身不高，复制通道不会引入额外噪声，三通道输入对预训练权重也友好。

另一个现实问题是：1800张图做训练偏少，而且六类缺陷数量均衡但形态差异大，类别内部不一致性高。比如裂纹类，细小的长度十几像素，宽的能有几十像素，aspect ratio差别很大。这种情况下的预处理策略就很重要。

2.2 标注格式转换：从VOC XML到YOLO TXT

NEU-DET原始标注是VOC格式的XML文件，一个缺陷目标对应一个bndbox。YOLOv5训练需要的是每个图像对应一个同名txt文件，每一行是class_id x_center y_center width height，坐标全部归一化到0到1。

转换脚本核心逻辑很直白，用Python的xml.etree.ElementTree解析XML，再结合opencv读图拿宽高，把绝对像素坐标归一化。工具脚本大概长这样：

import xml.etree.ElementTree as ET import os import cv2 def convert_xml_to_yolo(xml_path, img_dir, out_dir): tree = ET.parse(xml_path) root = tree.getroot() img_name = root.find('filename').text img_path = os.path.join(img_dir, img_name) img = cv2.imread(img_path) h, w = img.shape[:2] out_txt = os.path.join(out_dir, os.path.splitext(img_name)[0] + '.txt') with open(out_txt, 'w') as f: for obj in root.iter('object'): cls = obj.find('name').text if cls not in class_ids: continue box = obj.find('bndbox') x1 = float(box.find('xmin').text) y1 = float(box.find('ymin').text) x2 = float(box.find('xmax').text) y2 = float(box.find('ymax').text) x_center = (x1 + x2) / 2 / w y_center = (y1 + y2) / 2 / h box_w = (x2 - x1) / w box_h = (y2 - y1) / h f.write(f"{class_ids[cls]} {x_center:.6f} {y_center:.6f} {box_w:.6f} {box_h:.6f}\n")

这里有个很隐蔽的坑：一部分XML文件里的filename字段和实际文件名大小写不一致，比如XML里写的是Image_001.jpg，实际文件名是image_001.jpg。如果只依赖XML里的filename合并路径，会直接报图不存在。稳一点的做法是直接遍历图像目录，拿文件名前缀去匹配XML文件，而不是反过来。

2.3 数据划分与增强策略

数据集划分上，1800张图按8:1:1切成训练集1440张、验证集180张、测试集180张。有一个原则必须遵守：同一张原始图像不能同时出现在训练和验证集里，这个数据集本身不存在这种问题，但如果自己采集数据时，同一块缺陷的不同裁剪图就要注意了，否则验证集分数虚高，毫无参考意义。

增强策略上，我保留了YOLOv5默认的Mosaic增强。Mosaic会把4张训练图随机缩放后拼成一张，这在小目标检测场景里相当于变相扩大了小目标数量。NEU-DET的裂纹、麻点很多都是小目标，Mosaic对这类目标帮助很大。

HSV增强里，把饱和度增益调低了一点，因为钢材表面颜色本来就单调，饱和度变化太大会让模型学到不真实的颜色分布。翻转增强我开了水平翻转，垂直翻转没有开，因为钢材表面缺陷在轧制方向上有物理意义，垂直翻转会让缺陷形态失真。

3. YOLOv5模型训练的完整流程

3.1 环境搭建与依赖安装

环境是这套系统最容易卡住新手的地方。YOLOv5对PyTorch版本比较敏感，这里给出我实测稳定的组合：

Python 3.9 PyTorch 2.0.0 torchvision 0.15.0 CUDA 11.7 opencv-python 4.8.0

安装的时候优先用conda创建独立环境，不要直接装在base环境里：

conda create -n yolo python=3.9 conda activate yolo pip install torch==2.0.0 torchvision==0.15.0 --index-url https://download.pytorch.org/whl/cu117 git clone https://github.com/ultralytics/yolov5.git cd yolov5 pip install -r requirements.txt

注意requirements.txt里的opencv版本如果装不上，单独装opencv-python-headless代替也可以，只是服务端跑的时候省得和GUI库打架。遇到numpy版本冲突就强制降级到1.23.5，这是YOLOv5老版本最经典的兼容性问题。

3.2 修改配置文件：data.yaml和模型结构

数据集配置写在data/NEU-DET.yaml里：

path: ../datasets/NEU-DET train: images/train val: images/val test: images/test nc: 6 names: ['crazing', 'inclusion', 'patches', 'pitted_surface', 'rolled-in_scale', 'scratches']

这里的关键是path字段，我一开始写绝对路径，换机器就报错；改成相对路径之后，只要脚本从项目根目录启动就不会出问题。训练前跑一次python train.py --data NEU-DET.yaml --weights yolov5s.pt，如果数据路径配置有问题，会提示找不到标签文件，排查起来很快。

模型结构默认用yolov5s.yaml，anchor是COCO数据集聚出来的。NEU-DET的缺陷目标普遍比COCO小，直接沿用默认anchor不是最优解，理论上能用自动anchor调整，YOLOv5有个--noautoanchor参数可以关闭自动计算。我测试下来，不关自动anchor会先计算当前数据集聚类结果，计算出的anchor往往比默认的更小，对NEU-DET这类小目标场景更友好，所以不用额外改。

3.3 训练命令与超参数解读

训练命令我最终用的是：

python train.py \ --data data/NEU-DET.yaml \ --weights yolov5s.pt \ --img 640 \ --batch 16 \ --epochs 300 \ --hyp data/hyps/hyp.scratch-low.yaml \ --project ../runs/train \ --name neu_defect \ --device 0

几个关键参数解释一下。

img=640，输入尺寸。NEU-DET原始图是200x200，如果直接按原始尺寸训练，模型对微小缺陷的检测确实更直接，但YOLOv5的训练流程会先把图resize到640，200x200放大到640后纹理细节会有一定损失。所以这里要测试不同输入尺寸，我用640比用320的mAP高了大约5个点，因为大输入尺寸保留了更多小目标细节。

batch=16，这个值取决于显存大小。YOLOv5默认是16，如果12GB以下的显卡拉低到8或者4，配合梯度累积也可以。

epochs=300。NEU-DET数据集小，训练到200个epoch左右就能收敛，但300更保险，配合YOLOv5自带的早停机制best.pt和last.pt会自动保存最优权重。

hyp文件里，我手动调整了这两个超参数：

mosaic: 0.8 fliplr: 0.5

mosaic默认是1.0，我降到0.8，因为如果训练全程都用Mosaic，最后几个epoch的loss波动会比较大，到了后期Mosaic带来的过拟合风险要超过它带来的数据多样性收益。

3.4 训练过程监控与结果评估

训练过程中，我主要通过runs/train/neu_defect目录下的三个曲线来判断模型状态：box_loss、obj_loss、cls_loss。

box_loss前50个epoch下降很快，100个epoch以后进入平台期，这是正常现象。obj_loss如果出现先降后升，大概率是过拟合或数据集标签噪声问题。cls_loss对NEU-DET这种类别均衡的数据集一般不担心，但如果是自己采集的类别不平衡数据，就要重点看cls_loss里每个类别的差距。

训练结束后还有几个更关键的指标：P（精确率）、R（召回率）、mAP@0.5、mAP@0.5:0.95。

我实际用yolov5s训出来的结果大概是：

召回率低于精确率说明还是有一批缺陷没被模型找出来，主要漏在小尺寸裂纹和低对比度斑块上。后面优化方向很明确：提高输入分辨率、做更激进的小目标增强、或者换更高精度的模型。

4. 检测系统的工程实现

4.1 系统模块划分与检测流程

训练完模型，接下来就是把模型包成一个能被实际业务调用的检测系统。这套系统我分成四个模块：图像输入模块、模型推理模块、结果解析模块、统计输出模块。

图像输入模块支持单张图片、批量文件夹、摄像头视频流三种方式。模型推理模块封装了YOLOv5的模型加载和推理过程，统一输出检测结果对象。结果解析模块把每个检测框映射到六类缺陷名称，同时计算缺陷面积占比。统计输出模块把批次检测结果汇总成报表，标记有缺陷和无缺陷的图，并按缺陷类型做数量统计。

核心检测流程可以概括成：读图 -> 预处理（resize、归一化） -> 模型推理 -> NMS后处理 -> 结果可视化/输出。

4.2 推理核心代码：Detector类封装

这是inference/detector.py的核心部分，我把YOLOv5的推理逻辑封装成了一个类：

import torch import cv2 import numpy as np class SteelDefectDetector: def __init__(self, weights_path, device='0'): self.device = device if torch.cuda.is_available() else 'cpu' self.model = torch.hub.load('ultralytics/yolov5', 'custom', path=weights_path, force_reload=False) self.model.conf = 0.25 self.model.iou = 0.45 self.model.classes = None self.names = self.model.names def detect_image(self, img_path): img = cv2.imread(img_path) if img is None: return None results = self.model(img) detections = [] for *xyxy, conf, cls in results.xyxy[0]: detections.append({ 'bbox': [int(x) for x in xyxy], 'confidence': float(conf), 'class_id': int(cls), 'class_name': self.names[int(cls)] }) return detections def visualize(self, img, detections): colors = {0: (0, 0, 255), 1: (255, 0, 0), 2: (0, 255, 0), 3: (255, 255, 0), 4: (255, 0, 255), 5: (0, 255, 255)} for det in detections: x1, y1, x2, y2 = det['bbox'] label = f"{det['class_name']} {det['confidence']:.2f}" cv2.rectangle(img, (x1, y1), (x2, y2), colors[det['class_id']], 2) cv2.putText(img, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, colors[det['class_id']], 2) return img

torch.hub.load方式最大的优点是代码量小，加载本地权重文件非常方便。不过有个坑：第一次调用torch.hub.load会联网下载YOLOv5仓库，如果环境离线，会直接失败。解决方法是直接把yolov5克隆到本地，然后用--model参数指定本地路径，或者干脆用repo_or_dir参数指向本地：

self.model = torch.hub.load('../yolov5', 'custom', path=weights_path, source='local')

推理时的conf阈值，我试过0.1到0.5的范围，最终选0.25是综合考虑了误检和漏检。太低了会把背景纹理当成缺陷，太高了又放过低置信度的真实缺陷。iou阈值0.45是YOLOv5默认的NMS参数，重叠目标很少，不太用动。

4.3 批量检测、视频流与结果统计

批量检测脚本batch_detect.py负责遍历整个文件夹的图片，逐张送入Detector并保存结果，同时生成一份CSV统计表。

import os import csv import cv2 def batch_detect(detector, input_dir, output_dir): os.makedirs(output_dir, exist_ok=True) rows = [] for img_name in sorted(os.listdir(input_dir)): img_path = os.path.join(input_dir, img_name) detections = detector.detect_image(img_path) if detections is None: continue img = cv2.imread(img_path) vis_img = detector.visualize(img, detections) cv2.imwrite(os.path.join(output_dir, img_name), vis_img) defect_types = [d['class_name'] for d in detections] rows.append({ 'image': img_name, 'defect_count': len(detections), 'defect_types': ';'.join(defect_types), 'has_defect': len(detections) > 0 }) with open('result.csv', 'w', newline='') as f: writer = csv.DictWriter(f, fieldnames=rows[0].keys()) writer.writeheader() writer.writerows(rows) print(f"Processing done. {len(rows)} images checked.")

摄像头实时检测就没有太多代码上的特殊要求，直接循环读帧，每帧调用detector.detect_image，然后imshow显示。实际跑起来的时候，如果FPS很低，优先把img尺寸降下来，或者换一个轻量模型比如YOLOv5n、YOLOv5s，在实时场景下比追求极致精度的模型实用得多。

5. 常见问题与排查经验

5.1 训练阶段高频问题速查

这一块直接上表格，都是实际运行中碰到的问题：

有一个排查小技巧：YOLOv5每次训练都会在runs/train下生成一张labels.jpg，展示所有训练图的标注框分布。如果这张图里的框大小分布非常不均匀，比如全部集中在几个像素级别，说明标注或预处理出问题了，这时候先别急着训练，把数据检查清楚再跑，否则白烧时间。

5.2 样本数量与分布带来的坑

NEU-DET只有1800张图，训练集1440张，属于很小的数据规模。这种规模下最明显的两个问题就是过拟合和分布外泛化能力弱。

过拟合表现为验证集mAP还行，但换个实际产线的打光环境，检测效果崩了。缓解手段除了数据增强，还能用迁移学习，YOLOv5从COCO预训练权重起步，本身就已经把通用特征学好了，接下来微调只需学钢材缺陷的特定特征即可。我实测直接从零训练比用预训练权重低了差不多8个点的mAP，所以不要为了省下载流量而不用预训练权重。

另一个问题是类别间难度不均衡。裂纹类缺陷因为对比度低，AP普遍比其他类别低10个点左右。这种时候可以单独看类别AP曲线，针对难例类别的图像做更强的增强，或者给难分类别增加样本权重。YOLOv5没有现成的类别权重参数，我自己的做法是写一个简单的数据采样器，让模型在每个epoch里多“见到”难分类别的样本，效果有提升但不明显，这也说明如果真想提高难类别AP，核心还是要在数据层面补充更真实的缺陷样本。

6. 性能优化与后续扩展方向

6.1 针对小目标的专项优化

NEU-DET的缺陷目标尺寸大多小于32x32像素，对小目标不敏感是YOLOv5的已知短板，因为高层特征图下采样倍数大，小目标的特征早就被“压扁”了。针对这个问题，我做了三个方向优化。

第一是提高输入分辨率。img从640提升到1280，小目标的像素占比变大，mAP@0.5能提升四五个点。缺点也很明显，推理速度直接下降，需要权衡产线要求的FPS。

第二是增加高分辨率特征图输出。修改yolov5s.yaml里的anchors和检测头，让模型在P3（80x80）之外再输出一个P2（160x160）特征图检测层。这个改动需要同步调整损失函数和后处理代码，改动量不小但效果最好，尤其是对裂纹这种细长目标。

第三是融合注意力机制。在C3模块后面加一个轻量的SE模块或者CBAM模块，让模型更关注缺陷区域。我用过CBAM，训练出来的模型对斑块和麻点的区分度确实更好，但推理时间多了约15%，部署时可以考虑。

6.2 推理加速与工程部署

模型训练出来是PyTorch格式，直接部署效率不高。工业现场常规做法是把PyTorch模型转成ONNX，再转成TensorRT引擎，充分利用显卡的Tensor Core算力。

用YOLOv5自带的导出脚本可以一键转ONNX：

python export.py --weights best.pt --img 640 --batch 1 --include onnx

转ONNX时要注意，YOLOv5的检测头里有几个算子对ONNX不友好，比如grid生成和anchor解码部分。官方工具链已经处理好了动态shape和batch问题，但如果你手动改过模型结构，就要用--dynamic参数开启动态轴，否则转出来的模型只能固定输入尺寸。

转TensorRT时我遇到最典型的问题是量化精度下降。如果用FP16精度，NEU-DET小目标检测的mAP大概损失2到3个点，换成INT8量化直接掉8个点以上。所以我在部署上优先选FP16，mAP损失小，推理速度还能提升2到3倍。如果你的实际业务对精度要求不那么高，可以用INT8换更大的吞吐。

另外，如果没有NVIDIA GPU，也可以考虑OpenVINO部署到Intel CPU上。YOLOv5s转OpenVINO后，在i5-12400上推理耗时大概15到20毫秒，虽然比不上GPU，但胜在成本低、部署简单，很多工厂现有的工控机就能跑。

7. 个人经验与落地心得

做了这个项目之后，我最大的体会是：工业视觉项目的难点往往不在模型结构，而在数据质量和工程细节。

YOLOv5本身是个通用工具箱，它不会因为换了数据集就自动变好，但只要你把数据处理好、环境配好、超参数调对，它就能稳定输出一个不错的检测结果。很多人一上来就想改模型结构或者换新算法，结果把时间和精力都耗在调试上，最后效果还不如老老实实把数据标注和增强做扎实。

几个具体心得再强调一下：标注格式转换脚本一定要先跑小批量验证，不要一口气转换全部数据之后才发现坐标归一化错了；训练前先看一眼labels.jpg，确认标注框没问题再启动训练；训练结束后不要只盯着mAP看，按类别拆分AP，找出模型真正的薄弱环节。

这个项目后续要扩展也很方便。比如把NEU-DET换成你们自己产线采集的真实缺陷数据，只要按照同样的目录结构放到datasets下，改一下yaml文件，训练和部署流程完全不用动。配合在线标注工具，半天就能完成数据流转。再往产品化走，可以加一个Web前端，把检测结果实时呈现到大屏上，后台自动生成质检报告，这样整个系统就能从“能跑模型”变成“真正能用的质检工具”了。

本文还有配套的精品资源，点击获取