简介:针对钢材表面缺陷检测需求,这份资料以YOLOv7算法为基础,提供了一套完整可用的检测工程。内容包含训练好的模型权重、LabelImg标注的数据集、精确率-召回率曲线与损失曲线等评估结果,适合工业视觉算法工程师和深度学习入门者进行快速验证与二次开发。压缩包为RAR格式,共一百六十六个文件,整体大小约二百零四兆字节;文件构成上以模型权重文件、Python脚本、配置文件、Jupyter示例、训练图片以及XML/TXT两种格式的标注文件为主,另含环境配置脚本,便于不同框架下复现。目前已有1287人学习下载,属于同类资源中关注度较高的一份;权重文件可直接用于推理或微调,数据集按目录区分标注格式,能显著降低数据准备和模型部署的时间成本。
1. 钢材缺陷检测不是模型难,是数据和权重难
热轧带钢表面的裂纹、麻点、氧化皮、划痕,一直是产线质检里最磨人的环节。人工盯屏幕漏检率随班次下降明显,机器视觉方案又卡在同一个地方:不是YOLOv7的原理有多深,而是「标注好的数据集」和「训练好的权重」这两样东西,大多数团队根本没有。这套资源把YOLOv7钢材缺陷检测的完整链路一次性打包了:已经训练好的检测权重、配套的钢材缺陷数据集、PR曲线和loss曲线的训练日志,以及TensorRT和ONNX Runtime部署用的notebook。拿到手可以直接推理复现,也可以自己接着训练。适合正在做工业质检横向课题、毕业设计,或者想快速验证YOLOv7在自己产线数据上能不能用的人。
2. 资源盘点:先搞清楚压缩包里每一份文件是干什么的
很多人在网上下载项目压缩包后第一件事就是找best.pt,结果翻遍目录也没找到,其实是没看懂目录结构。先别急着跑代码,花两分钟把文件分组看清楚,后面能少踩一半的坑。这套资源里真正值得关注的东西分三类:训练产物、数据集、部署文件,剩下全是官方YOLOv7仓库自带的工程文件。
2.1 从目录里挑出三样关键的东西
我把根目录文件按用途拆了一下,见下表:
| 文件 / 目录 | 类型 | 干什么用的 |
|---|---|---|
events.out.tfevents.1679279440.DESKTOP-AJP7QI2.64992.0 | TensorBoard日志 | 训练全过程的loss、mAP、学习率曲线,时间戳对应2023年3月左右的一次完整训练 |
Dockerfile | 环境文件 | 构建训练和推理镜像用,里面锁定了CUDA、PyTorch等基础环境 |
YOLOv7-Dynamic-Batch-TENSORRT.ipynb | 部署notebook | TensorRT动态batch推理的完整流程,从导出engine到调用推理 |
YOLOv7-Dynamic-Batch-ONNXRUNTIME.ipynb | 部署notebook | ONNX Runtime动态batch推理,不依赖TensorRT时用这个 |
compare_YOLOv7_vs_YOLOv5m6_half.ipynb | 对比实验 | YOLOv7和YOLOv5m6在half精度下的速度和精度对比 |
compare_YOLOv7e6_vs_YOLOv5x6_half.ipynb | 对比实验 | YOLOv7e6和YOLOv5x6在half精度下的对比 |
compare_YOLOv7e6_vs_YOLOv5x6.ipynb | 对比实验 | 全精度下e6和5x6的对比 |
.gitignore/yolov7-main.iml | 工程文件 | 不用管,IDE和Git的辅助文件 |
权重和数据集的存放位置不在根目录,而是在runs/train/和独立的数据文件夹里。runs/train/下面每个子目录对应一次训练,weights/里有best.pt和last.pt,前者是按验证集mAP保存的最优权重,后者是最后一轮epoch的权重。数据集则分为images(jpg图片)、annotations/xml(VOC格式标注)、labels(YOLO格式txt标注)三块,图片和两种标注文件通过同名文件对应。
2.2 训练产物别放着吃灰,里面有完整的质检报告
那个events.out.tfevents开头的长文件名,很多人以为是缓存文件直接删了,其实是TensorBoard的日志。作者把整个训练过程的曲线都存下来了,用一行命令就能看:
tensorboard --logdir runs --port 6006启动后浏览器打开http://localhost:6006,你能看到box_loss、obj_loss、cls_loss三条损失曲线,以及mAP、precision、recall随epoch的变化。我看训练日志的习惯是先看val曲线的抖动程度:如果val loss有规律的锯齿状下降,说明batch size和learning rate配合得还行;如果val曲线在高位平着走,说明模型欠拟合,这时候直接看PR曲线文件会更直观。
PR曲线在runs/train/对应的训练目录下,文件名是PR_curve.png,每类缺陷一条曲线。钢材缺陷检测里我最关心的是小目标那几条曲线,比如麻点和划痕通常框小、对比度低,曲线会往左下角沉。如果这条曲线离右上角近,说明precision和recall能同时保持在较高水平;如果曲线中间有个明显的拐弯下坠,说明存在大量低置信度的误检框,推理时可以把置信度阈值往上调。
2.3 Dockerfile和notebook:环境与部署一套带走
Dockerfile的作用是把环境固化成镜像,避免在别人的机器上复现时被CUDA版本、PyTorch版本折磨。我一般会先看里面的基础镜像版本,确认CUDA版本,再决定本地用裸环境还是走Docker。
三个compare_开头的notebook是官方仓库自带的对比实验,用来判断不同模型尺寸在同等条件下的精度和推理速度差异。做钢材缺陷检测时我建议看一眼compare_YOLOv7e6_vs_YOLOv5x6.ipynb,如果产线对准确率要求高但没GPU部署条件,这里的数据能帮你决定是不是要换大模型;如果只是常规质检,YOLOv7原始尺寸已经够用,部署成本低很多。
3. 数据准备:xml与txt双标注格式怎么对应,怎么校验
钢材缺陷检测的模型训练全过程里,数据格式是第一个坑。这套资源里同一个数据集给了一份xml和一份txt,分别放在两个文件夹。这就意味着作者用labelimg标注后,又做了一次格式转换。搞清楚这两种格式的对应关系,你才能正确修改训练配置里的标注路径。
3.1 同一张图,两种标注格式的差异
xml是VOC Pascal格式,标签载体是标签名,比如<name>crazed</name>,坐标是绝对值:
<annotation> <filename>0a9be0c3.jpg</filename> <size> <width>200</width> <height>200</height> </size> <object> <name>crazed</name> <bndbox> <xmin>42</xmin> <ymin>83</ymin> <xmax>156</xmax> <ymax>171</ymax> </bndbox> </object> </annotation>txt是YOLO格式,没有标签名,只有类别序号和归一化坐标,一行一个框:
0 0.495000 0.635000 0.570000 0.440000五个数字依次是:类别id、中心点x、中心点y、框宽、框高,全部除以图片宽高做了归一化。类别id从0开始,和你在训练配置文件里写的names列表顺序一一对应。例如names写成['crazed', 'inclusion'],那么txt里第一列为0就代表crazed,为1代表inclusion。
注意xml和txt分别存放在annotations/xml和labels两个文件夹中,文件主名与jpg图片完全一致。训练框架读取txt时,只看图片同名txt是否存在,并不关心同目录下有没有xml,所以你在yaml里配置的labels路径一旦指向xml文件夹,训练不会报错,但会提示找不到标签文件。
3.2 标注一致性核查:一个脚本找出脏数据
拿到数据集后别直接开训,先跑一遍标注核查脚本。我曾遇到txt里类别id最大值为6,而数据yaml只定义了6类,结果训练时所有框都被忽略,模型学了个寂寞。这个脚本把同名的xml和txt解析出来做对比,能查出坐标不一致、类别越界、文件缺失三类问题:
import os import glob import xml.etree.ElementTree as ET img_w, img_h = 200, 200 # 改成你数据集的真实分辨率 xml_dir = "steel_data/annotations/xml" txt_dir = "steel_data/labels/train" errors = 0 for xml_path in glob.glob(os.path.join(xml_dir, "*.xml")): base = os.path.splitext(os.path.basename(xml_path))[0] txt_path = os.path.join(txt_dir, base + ".txt") if not os.path.exists(txt_path): print(f"[缺失] {base} 没有对应的txt") errors += 1 continue root = ET.parse(xml_path).getroot() xml_boxes = [] for obj in root.iter("object"): box = obj.find("bndbox") xmin = float(box.find("xmin").text) xmax = float(box.find("xmax").text) ymin = float(box.find("ymin").text) ymax = float(box.find("ymax").text) cx = ((xmin + xmax) / 2) / img_w cy = ((ymin + ymax) / 2) / img_h w = (xmax - xmin) / img_w h = (ymax - ymin) / img_h xml_boxes.append([cx, cy, w, h]) with open(txt_path) as f: txt_boxes = [ [float(v) for v in line.strip().split()[1:]] for line in f if line.strip() ] if len(xml_boxes) != len(txt_boxes): print(f"[数量不一致] {base}: xml={len(xml_boxes)}, txt={len(txt_boxes)}") errors += 1 continue for xb, tb in zip(xml_boxes, txt_boxes): if max(abs(a - b) for a, b in zip(xb, tb)) > 0.01: print(f"[坐标差异] {base}: xml={xb}, txt={tb}") errors += 1 print(f"共检查出 {errors} 个问题") if errors == 0: print("标注一致性校验通过,可以开始训练")这段脚本的核心逻辑是把xml里的绝对坐标换算成和txt一致的归一化中心点格式,然后逐框比对。容差设为0.01,因为浮点精度不同可能导致微小差异。类别id是否越界也要查,单独对txt文件多跑一步:读取每行第一个数字,如果大于等于你配置的类别总数,直接删掉这行或人工复查。
3.3 从xml批量转txt:留着备用
如果哪天你想把这份数据拿到其他框架里用,或者想自己标注一份新数据,批量转换脚本是刚需。labelimg支持保存VOC格式xml,但不会同时生成YOLO格式txt,所以转换这一步基本绕不开:
import os import xml.etree.ElementTree as ET xml_dir = "annotations/xml" out_dir = "labels/train" os.makedirs(out_dir, exist_ok=True) class_names = ["crazed", "inclusion", "patches", "pitted_surface", "rolled_in_scale", "scratches"] for xml_name in os.listdir(xml_dir): if not xml_name.endswith(".xml"): continue root = ET.parse(os.path.join(xml_dir, xml_name)).getroot() img_w = float(root.find("size").find("width").text) img_h = float(root.find("size").find("height").text) lines = [] for obj in root.iter("object"): name = obj.find("name").text if name not in class_names: print(f"跳过未知类别 {name} 在 {xml_name}") continue cls_id = class_names.index(name) box = obj.find("bndbox") xmin = float(box.find("xmin").text) ymin = float(box.find("ymin").text) xmax = float(box.find("xmax").text) ymax = float(box.find("ymax").text) # 中心点坐标归一化,并把越界的值夹到 [0, 1] cx = max(0.0, min(1.0, (xmin + xmax) / 2 / img_w)) cy = max(0.0, min(1.0, (ymin + ymax) / 2 / img_h)) w = max(0.0, min(1.0, (xmax - xmin) / img_w)) h = max(0.0, min(1.0, (ymax - ymin) / img_h)) lines.append(f"{cls_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}") base = os.path.splitext(xml_name)[0] with open(os.path.join(out_dir, base + ".txt"), "w") as f: f.write("\n".join(lines))class_names的列表顺序就是最终txt里的类别编号,务必和后面训练yaml里的names顺序保持一致,否则会出现「类别id变了但含义没变」的坑。坐标夹到[0,1]是为了防止个别标注框边缘超出图片范围导致归一化值大于1,训练时YOLO会警告甚至忽略这些框。
这套数据格式的处理流程,同样适用于yolov8训练自己的数据集。YOLOv8对标注格式的要求和YOLOv7基本一致,都是每行「类别id cx cy w h」,所以把labels文件夹和图片目录按yolov8的目录规范放好,就能直接复用。
4. 模型训练与权重使用:从data yaml到读懂loss曲线
训练是自己复现权重必须过的一关,也是看清楚这份资源到底「有没有水分」的关键。很多项目只给权重不给训练过程,这套资源连TensorBoard日志都给了,意味着你可以完整复盘作者当时的训练策略,也可以在这个基础上继续微调。
4.1 data yaml与类别顺序:训练前唯一必须改的东西
YOLOv7的训练入口是train.py,它通过--data参数读一个yaml文件。这个yaml定义了训练集、验证集路径和类别信息。钢材缺陷检测场景下,yaml长这样:
train: steel_data/images/train val: steel_data/images/val nc: 6 names: 0: crazed 1: inclusion 2: patches 3: pitted_surface 4: rolled_in_scale 5: scratchestrain和val指向图片目录,框架会自动在同级目录下找同名txt标签。比如图片在steel_data/images/train/0a9be0c3.jpg,标签就在steel_data/labels/train/0a9be0c3.txt,注意不是images的兄弟目录,而是labels这个独立目录。
nc是类别总数,names是类别名列表。这里有个很容易翻车的细节:txt里第一列的数字是按这个names顺序编码的,不是按字母序,也不是按你在labelimg里看到的顺序。我习惯把训练前检查「txt类别id是否小于nc」列为固定动作,跑一条命令就清楚:
python - <<'EOF' import glob for f in glob.glob("steel_data/labels/train/*.txt"): for line in open(f): cid = int(line.strip().split()[0]) if cid >= 6: print(f, cid) EOF这个检查十秒钟跑完,能挡住后续一整天的无效训练。如果你拿到的数据集里类别分布偏斜,比如patches占了80%、scratches不到5%,建议在yaml里先不做处理,用脚本统计后决定要不要加类权重,而不是盲目开始。
4.2 训练命令与超参数:batch size、img size、epochs怎么配
确认数据路径没问题后,启动训练的标准命令如下:
python train.py \ --data data/steel.yaml \ --cfg cfg/training/yolov7.yaml \ --weights '' \ --batch-size 16 \ --img 640 \ --epochs 150 \ --workers 8 \ --device 0 \ --name steel_defect--weights ''表示从零开始训练,而不是加载COCO预训练权重。钢材缺陷数据集规模不大,从头训练150个epoch完全够;如果你想借COCO预训练模型做迁移学习,把--weights换成yolov7.pt即可,收敛会更快,但要注意类别数不同会导致输出层shape不匹配,YOLOv7会自动调整,不需要手动改网络结构。
--batch-size是显存敏感参数。钢材缺陷图片通常是200x200或300x300的小图,但训练时会被resize到640x640,显存占用不可小觑。我的经验是:8G显存用batch-size 8或16,如果训练中途报CUDA out of memory,先把batch降到8,再把img降到512,优先保batch。--workers在Windows下建议设成0,在Linux下可以设成CPU核心数的一半,否则数据加载会成为瓶颈。
训练开始后,终端会周期性打印每轮epoch的box_loss、obj_loss、cls_loss和mAP@0.5。正常情况是三行loss稳步下降,mAP在前期有波动但整体抬头。如果发现loss降着降着突然变成nan,十有八九是学习率太大或batch里混进了损坏图片,先把--hyp data/hyp.scratch.p5.yaml里的lr0从0.01改成0.001重新来。
4.3 读懂权重和曲线:best.pt、last.pt、PR曲线和mAP的对应关系
训练结束的目录结构在runs/train/steel_defect/下:
runs/train/steel_defect/ ├── weights/ │ ├── best.pt # 验证集mAP最高的一轮 │ └── last.pt # 最后一轮 ├── PR_curve.png # 每类缺陷的PR曲线 ├── confusion_matrix.png ├── results.png # loss和mAP总览 └── events.out.tfevents.* # TensorBoard原始日志best.pt是训练过程中保存的验证集mAP最优权重,也是后续推理和部署的首选。last.pt有时候反而比best好用,因为有些任务在训练后期验证集mAP不再上升但泛化能力还在变好,我一般两个都保留。有人直接下载yolov8的官方预训练权重来跑钢材检测,但COCO的80类和缺陷类别完全对不上,输出层也要改,远不如直接用这份训练好的yolov7权重来得快。
看训练效果,我推荐直接扫一眼results.png。这张图里左上角是train/val的box_loss曲线,val曲线整体贴着train曲线下降说明没有过拟合;如果val曲线在某个epoch后反弹,说明开始过拟合了,回调早停即可。PR曲线则回答了另一个问题:在这个数据集上,每一类缺陷能在什么置信度下保持多高的精确率。钢材缺陷里那种细长条裂纹,PR曲线往往带一段快速下滑,实际部署时就得把置信度阈值压到0.25以下才不漏检。
TensorBoard日志文件也在这个目录里,想看细节执行:
tensorboard --logdir runs/train/steel_defect这个命令会输出一个本地地址,浏览器打开后能逐epoch查看每个指标。我习惯把val mAP@0.5达到0.9以上的轮次记下来,用那个epoch的权重做部署,而不是机械地用最后一个best。
5. 避坑与排查:五个翻车现场和处理办法
这部分内容来自我自己跑钢材缺陷检测项目时的教训,也包含和同行交流时经常被问到的共性问题。每一条都按现象、原因、解决的顺序写,遇到类似情况可以直接对照处理。
5.1 标签目录指错:loss在降但mAP几乎为0
现象:训练能正常启动,loss数值在下降,训练日志里也没有报错,但跑了30个epoch后mAP始终在0.05以下,基本等于瞎猜。
原因:data yaml里配置的标签路径指向了xml文件夹,或者指向了一个空目录。YOLOv7的数据加载器对「该图片没有标签」是静默容忍的,它会跳过这张图,但不会报「找不到标签文件」。结果就是模型一直在学空背景,loss自然下降,但没有目标可学。
解决:训练前先检查yaml路径下是否存在txt。最稳的办法是第4.1节里的类别id检查脚本,顺手把「txt文件是否存在」也打印出来。我现在的习惯是:启动训练后等前两个epoch输出,看一眼每个epoch的图片数量和标签框数量,如果num_samples比训练集总数小很多,立刻中断去查路径。
5.2 类别顺序错位:标签没丢,但预测框全是背景
现象:数据检查一切正常,txt文件数量和图片对得上,类别id也没有越界。训练完跑推理,检测框全都打在随机位置,置信度还特别高。
原因:txt里的类别id和data yaml里的names顺序不一致。常见于从网上下载数据集时,原作者txt里0代表crazed,但你yaml里0写的是inclusion,或者反过来。坐标没变,框位置是对的,但标签语义错位,模型学到的映射关系就是错的。
解决:用第3.2节的核查脚本,把xml里的name和txt第一列对应起来逐条验证。具体做法是解析xml得到(name, cx, cy, w, h),再解析txt得到(cls_id, cx, cy, w, h),如果xml里name为crazed的框对应txt第一列为1,而yaml里1是inclusion,那就是顺序错了。修正names列表后重新生成txt,不要手动改txt内容,很容易改漏。
5.3 显存翻车:CUDA out of memory与训练中断
现象:训练到第5轮,突然报CUDA out of memory,前面几轮跑得好好的,然后训练进程直接退出。
原因:YOLOv7的显存占用会随训练推进变化,尤其是开启了--cache-images或数据增强时,输入图尺寸虽然固定,但中间feature map的缓存峰值可能出现在训练中段。更常见的原因是验证阶段也要跑一次完整的前向推理,显存本来就吃紧的情况下,验证那一瞬间爆掉。
解决:优先把batch-size从16降到8,这是最直接的手段。其次是--img 512,对小目标召回率影响不大但显存占用能砍掉近一半。还不行就在命令里加--cache-images False,关闭图像缓存。如果是偶发性的OOM,把--workers调低或改成0,避免数据预加载和训练竞争显存。这个坑本质上是显存预算问题,没有玄学,就是按上述顺序逐步降资源。
5.4 换了自己的缺陷图片,效果明显变差
现象:拿训练好的权重检测训练集里的图片,效果很好,一换成产线实拍的图片,或者换成燃气管道、水下管道的裂缝图像,漏检率大幅上升。
原因:这套权重的训练数据来自热轧带钢表面,光照均匀、背景固定、缺陷形态相对标准。换到别的表面缺陷场景,背景纹理、光照条件、缺陷外观分布全都变了,模型学到的特征不再匹配。这不是模型坏了,而是数据分布变了。用mmrotate那类旋转框检测方案也不会解决这个问题,因为问题不在框的形式,而在域差异。
解决:把这份权重当作预训练模型,用你自己的缺陷图片做迁移学习。做法是冻结backbone,只训练head,前30个epoch用较低学习率,30个epoch后解冻全部层再微调50个epoch。我在第4.2节提过,加载best.pt而不是COCO权重,收敛更快,几十张标注图就能见效。
5.5 TensorRT和ONNX转换时的shape不匹配
现象:用官方export脚本导出onnx成功,但转TensorRT engine时报input shape mismatch,或者推理时输入尺寸一变就报错。
原因:多数是导出时没有明确指定动态batch维度。YOLOv7的export.py默认按静态batch导出,--batch 1导出的onnx只接受batch=1的输入,你在推理时换成batch=4就炸了。
解决:导出时显式加--dynamic,让onnx的输出层保留batch维度的动态属性。如果是转TensorRT,trtexec构建engine时要为input设置三个shape:--minShapes=input:1x3x640x640、--optShapes=input:4x3x640x640、--maxShapes=input:8x3x640x640。这里的maxShape要小于等于你的显存能承受的上限,否则构建成功但运行时爆显存。资源里那两个Dynamic-Batch开头的notebook就是干这个的,直接照着跑,别自己重新造轮子。
6. 部署进阶:用ONNX Runtime跑动态batch推理
验证一个训练好的权重到底能不能上产线,最快的方式是走一遍ONNX Runtime推理链路。不依赖TensorRT也中间层依赖少,CPU和GPU都能跑。先导出onnx:
python export.py \ --weights runs/train/steel_defect/weights/best.pt \ --img 640 \ --batch 1 \ --dynamic导出的best.onnx就是部署文件。如果导出时报opset相关错误,在命令里加--simplify或用onnx-simplifier工具先简化一次,多半是某些算子在onnx里表达冗余导致的。这一步我每次都会做:导出后用onnx.checker.check_model验证一次模型结构,再去谈推理。
ONNX Runtime推理代码骨架如下:
import onnxruntime as ort import numpy as np sess = ort.InferenceSession( "best.onnx", providers=["CUDAExecutionProvider", "CPUExecutionProvider"] ) input_name = sess.get_inputs()[0].name img = np.zeros((1, 3, 640, 640), dtype=np.float32) # 实际换成预处理后的图 outs = sess.run(None, {input_name: img}) # 输出shape: (1, 25200, 4+1+6),6是钢材缺陷的类别数 # 25200 = (80*80 + 40*40 + 20*20) * 3,三个尺度共25200个anchor框 scores = outs[0][0, :, 4] # objectness boxes = outs[0][0, :, :4]输出维度里4+1+6的含义是:前4个值是中心点坐标和宽高,第5个是objectness置信度,后6个是六类缺陷的类别得分。25200是yolov7在640x640输入下三个检测尺度生成的anchor总数,固定不变。实际部署时先把objectness低于0.25的框过滤掉,再做NMS合并重叠框,最后按类别得分取最大值作为最终类别。
动态batch正确验证的方法是:第一次输入batch=1跑一遍,第二次输入batch=4再跑一遍,两次推理都能正常返回结果说明dynamic axis工作正常。从那以后,我每次拿到别人给的yolov7权重,都会强制走一遍「txt类别id检查 → export导出 → onnx动态batch推理」这三步流程,权重本身没有问题才是真正能落地的前提,这三个动作能挡掉七成以上的部署事故。希望帮到你。
本文还有配套的精品资源,点击获取