简介:YOLOv8钢材缺陷检测完整方案,面向工业质检算法学习者和视觉项目开发者,解决钢材表面多类型缺陷的自动识别与可视化检测问题。压缩包共2000个文件,包含346张标注图像、1403个txt标签及配套xml标注文件,以及训练好的pt权重、yaml配置、Python脚本、pyqt界面ui文件和训练曲线pdf等,整体约94.62MB,目录划分清晰。检测模型已训练完成,附带PR曲线、Loss曲线等评估结果,可直接用于图片、视频及摄像头实时检测场景。pyqt图形界面提供检测图片、视频和调用摄像头的选择项,便于二次开发与演示。数据集采用labelimg标注,jpg图片与xml、txt双格式标签分离存放,适合目标检测任务训练与验证。已有605人学习下载,适合具备一定深度学习基础、希望快速落地YOLOv8缺陷检测项目的工程师或研究者。
1. 从产线质检到桌面工具,YOLOv8钢材缺陷检测为什么需要一套完整链路
钢材表面缺陷检测是工业视觉里最典型的落地场景之一,热轧带钢、冷轧板卷、型钢表面在连续生产过程中会出现裂纹、麻点、夹杂、氧化皮压入等缺陷。传统方案靠人工目检,效率低且容易漏检,而通用目标检测模型直接套用在钢材表面场景,会遇到小目标占比高、缺陷形态长宽比极端、背景纹理干扰强这三个核心问题。YOLOv8在C2f结构和Anchor-Free检测头的设计上,相比YOLOv5有更灵活的梯度流和更简洁的解耦头,适合作为这个任务的基础骨架。
但模型训练只是其中一环。实际项目中,工程师需要一份整理好的数据集、一组能直接加载的缺陷检测权重,以及一个能让现场操作人员不看命令行也能完成单张图片和视频检测的图形界面。把这三件事拼起来,才是一条从模型训练到产线试用的完整链路。这篇文章假设你已经有Python基础,不需要GPU服务器,一台带有6GB以上显存的消费级显卡就能跑通全部流程。全文按数据集准备、权重训练、Qt界面封装、落地排错的顺序展开,每一步都有可直接执行的命令和代码。
2. 钢材缺陷检测数据集的选型与预处理,NEU-DET怎么变成YOLOv8能吃的格式
2.1 公开数据集和自采数据的取舍,缺样本时先做哪几件事
钢材缺陷检测领域最常用的公开数据集是NEU-DET,它来自东北大学,包含热轧带钢表面六大类缺陷:rolled-in scale(氧化铁皮压入)、patches(麻点)、crazing(裂纹)、pitted surface(凹坑)、inclusion(夹杂)、scratches(划痕)。每类样本180张,共1800张图片,分辨率200x200像素。这个数据集规模不大,直接训练容易过拟合,而且单张图片尺寸偏小,和产线相机采集的2000万像素工业图像分布差异明显。
如果你手头有现场采集数据,我一般会建议先用NEU-DET跑通流程,再用现场数据做微调。现场数据标注格式可能是VOC的XML,也可能是MS COCO的JSON。YOLOv8官方训练接口接受两种格式:一种是每张图片对应一个同名TXT文件,每一行是"class_id x_center y_center width height",坐标统一归一化到0到1;另一种是COCO格式的JSON,通过参数format="coco"指定。无论原始格式是什么,第一步都是统一转换成YOLO的TXT格式。
import os import xml.etree.ElementTree as ET def convert_voc_to_yolo(xml_path, out_dir, class_names): tree = ET.parse(xml_path) root = tree.getroot() img_w = int(root.find('size').find('width').text) img_h = int(root.find('size').find('height').text) txt_name = os.path.splitext(os.path.basename(xml_path))[0] + '.txt' out_path = os.path.join(out_dir, txt_name) with open(out_path, 'w') as f: for obj in root.findall('object'): class_name = obj.find('name').text if class_name not in class_names: continue class_id = class_names.index(class_name) bbox = obj.find('bndbox') x1 = float(bbox.find('xmin').text) y1 = float(bbox.find('ymin').text) x2 = float(bbox.find('xmax').text) y2 = float(bbox.find('ymax').text) x_center = (x1 + x2) / 2.0 / img_w y_center = (y1 + y2) / 2.0 / img_h w = (x2 - x1) / img_w h = (y2 - y1) / img_h f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}\n") class_names = ['crazing', 'inclusion', 'patches', 'pitted_surface', 'rolled-in_scale', 'scratches'] # 对NEU-DET的每一张XML标注执行转换 xml_dir = 'data/NEU-DET/ANNOTATIONS' txt_dir = 'data/NEU-DET/labels' os.makedirs(txt_dir, exist_ok=True) for xml_file in os.listdir(xml_dir): if xml_file.endswith('.xml'): convert_voc_to_yolo(os.path.join(xml_dir, xml_file), txt_dir, class_names)这段代码包含一个关键处理细节:NEU-DET原标注的类别名和YOLOv8的类别索引必须手动对齐,比如rolled-in_scale包含连字符,在类名列表里要严格一致。VOC坐标转归一化坐标的公式是固定的,但要注意某些标注工具的坐标原点在左下角,而图像数组的坐标原点在左上角,NEU-DET不存在这个问题,自采数据建议先画框验证一张再批量转换。
2.2 缺陷检测权重训练前的数据集划分与增强策略
数据集划分上,NEU-DET的标准做法是随机分配训练集和验证集,比如训练集占80%,验证集占20%。钢材缺陷检测和自然场景目标检测有个重要区别:相邻缺陷在空间上高度相关,同一块钢材上的缺陷类型往往有聚集性。如果按整图随机划分,可能出现同一批钢材的图像同时出现在训练集和验证集,评估指标虚高。更好的做法是按钢材母卷编号分组划分,这在NEU-DET上实现不了,因为官方没有提供图像间的关联信息,但对自采数据一定要按生产线批次划分。
import os import random from collections import defaultdict def split_dataset_by_group(image_dir, label_dir, output_dir, train_ratio=0.8): images = [f for f in os.listdir(image_dir) if f.endswith('.jpg')] groups = defaultdict(list) for img_name in images: # 假设文件名格式为 steel_20240115_batch01_0001.jpg # 取前三段标识一个母卷批次 parts = img_name.split('_') if len(parts) >= 3: group_key = '_'.join(parts[:3]) else: group_key = 'default' groups[group_key].append(img_name) train_imgs, val_imgs = [], [] for group_key, group_imgs in groups.items(): random.shuffle(group_imgs) split_idx = int(len(group_imgs) * train_ratio) train_imgs.extend(group_imgs[:split_idx]) val_imgs.extend(group_imgs[split_idx:]) os.makedirs(os.path.join(output_dir, 'images', 'train'), exist_ok=True) os.makedirs(os.path.join(output_dir, 'images', 'val'), exist_ok=True) os.makedirs(os.path.join(output_dir, 'labels', 'train'), exist_ok=True) os.makedirs(os.path.join(output_dir, 'labels', 'val'), exist_ok=True) for img_name in train_imgs: os.system(f"cp {os.path.join(image_dir, img_name)} {os.path.join(output_dir, 'images', 'train')}/") lbl = img_name.replace('.jpg', '.txt') if os.path.exists(os.path.join(label_dir, lbl)): os.system(f"cp {os.path.join(label_dir, lbl)} {os.path.join(output_dir, 'labels', 'train')}/") # 验证集同理 print(f"训练集 {len(train_imgs)} 张,验证集 {len(val_imgs)} 张")按批次划分这段代码用文件名提取分组键,这个思路在生产项目里很实用。NEU-DET文件名是NEU-CLS-XXXX形式,无法提取批次,但如果你的自采图像命名规则里包含了产线号、日期、轧制批次,这行代码可以直接复用。
数据增强策略上,钢材表面缺陷有三个特殊情况需要注意。第一,大多数缺陷是灰度纹理变化,色相饱和度增强对模型没有正向帮助,反而可能引入噪声,所以hsv_h、hsv_s参数应该调低。第二,缺陷长宽比极端,比如划痕可能横向跨整个图像宽度但高度只有几个像素,mosaic增强在拼接时会裁掉部分缺陷,严重降低小缺陷的可见性。第三,钢材表面有周期性纹理,fliplr水平翻转会让模型学习到左右对称的纹理模式,这对某些特定方向性缺陷可能是坏事。
基于这些分析,我通常使用如下增强参数:关闭hsv_h到0.015,hsv_s维持0.5但饱和度的作用本来就小,mosaic在最后10个epoch关闭,mixup设置为0.1,copy_paste开启到0.3。这套参数在NEU-DET上比默认参数mAP高约1.5到2个百分点,主要是mAP50-95的提升,因为小缺陷的定位更稳定了。实际效果会随数据分布变化,但不建议直接使用默认增强配置。
数据集整理完成后,还需要在YOLOv8的配置YAML里写入路径和类别名。train.txt和val.txt可以使用绝对路径列表,也可以只写图片目录路径让Ultralytics自动扫描。路径建议使用绝对路径而不是相对路径,因为Qt界面调用时工作目录可能被切换。
3. 训练钢材缺陷检测权重的完整流程,从C2f结构到损失函数曲线图
3.1 YOLOv8环境配置与训练参数的确定
YOLOv8的依赖环境搭建相对轻量,核心是ultralytics这个PyPI包,它会自动拉取PyTorch、OpenCV、Pandas等依赖。推荐使用conda创建独立环境,Python版本选3.9或3.10,PyTorch版本根据CUDA版本选择。
# 创建环境并安装依赖 conda create -n steel_yolov8 python=3.10 -y conda activate steel_yolov8 pip install ultralytics==8.2.0 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118这段命令里,--index-url指定了CUDA 11.8版本的PyTorch轮子。如果显卡驱动支持CUDA 12.1,可以换成cu121。安装完可以用python -c "import torch; print(torch.cuda.is_available())"检查GPU是否可用。GTX1660Ti这种6GB显存的卡也能训练YOLOv8s,但batch size只能设为8,输入分辨率调整到640。
训练时核心参数集中在模型尺寸、输入分辨率、批量大小和学习率四块。钢材缺陷检测场景,缺陷尺寸小,输入分辨率建议从默认的640提升到800或960。提高分辨率会带来显存压力,比如GTX1660Ti上6GB显存,imgsz=800时batch size只能开到4,imgsz=960时可能直接OOM。显存不够时优先降低batch而不是降低分辨率,因为小缺陷的检测效果对分辨率更敏感。
学习率权重上,YOLOv8默认使用AdamW优化器,lr0默认0.0001,加入weight_decay=0.0005。对于NEU-DET这类小数据集,用较小的lr0能避免前期震荡,我一般设0.0002。训练200个epoch后如果验证集loss还在下降,可以续训而不是从头开始。
3.2 model.train()训练与损失函数曲线图绘制
使用Ultralytics训练接口时,数据集YAML文件是唯一的配置入口。假设钢材缺陷数据集目录结构为steel_data/,下面有images/train、images/val、labels/train、labels/val,那么YAML可以写成这样:
# steel_dataset.yaml path: /home/user/steel_data train: images/train val: images/val nc: 6 names: 0: crazing 1: inclusion 2: patches 3: pitted_surface 4: rolled-in_scale 5: scratches启动训练的命令可以放在脚本文件里,方便记录实验参数:
# train_steel.py from ultralytics import YOLO model = YOLO('yolov8s.pt') results = model.train( data='steel_dataset.yaml', epochs=200, imgsz=800, batch=4, optimizer='AdamW', lr0=0.0002, weight_decay=0.0005, hsv_h=0.015, hsv_s=0.5, hsv_v=0.4, mosaic=0.8, mixup=0.1, copy_paste=0.3, device=0, project='steel_run', name='exp1' )每个训练参数都有明确的理由。imgsz=800是因为钢材缺陷里的小目标在640分辨率下可能只占几个像素;batch=4受限于6GB显存;mosaic=0.8而不是1.0,给最后10个epoch留出关闭余量,Ultralytics内部会在最后10个epoch自动降低mosaic到0,但保险起见显式设置;copy_paste=0.3对圆形或块状缺陷如麻点、夹杂有效,对长条状划痕意义不大。
训练完成后,在steel_run/exp1/weights/下会生成best.pt和last.pt两个权重文件。best.pt依据验证集mAP挑选,last.pt是最后一个epoch的状态。如果显存充足且训练稳定,best.pt基本能代表最优权重;如果训练后期过拟合,last.pt可能反而比best.pt泛化更好,建议两个都跑一遍测试集对比。
损失函数曲线图是判断训练是否健康的第一手资料。YOLOv8的results.csv里记录了train/box_loss、train/cls_loss、train/dfl_loss和对应的验证集损失。用Matplotlib画出来,主要看三点:第一,box_loss是否平滑下降,有震荡说明学习率太高;第二,cls_loss和val/cls_loss之间的缝隙是否变大,变大说明过拟合;第三,dfl_loss在最后50个epoch如果还在下降,说明分布焦点损失还在起作用,可以加训。
import pandas as pd import matplotlib.pyplot as plt # 读取训练日志 df = pd.read_csv('steel_run/exp1/results.csv') fig, axes = plt.subplots(1, 3, figsize=(15, 4)) # 边框损失 axes[0].plot(df['epoch'], df['train/box_loss'], label='train_box') axes[0].plot(df['epoch'], df['val/box_loss'], label='val_box') axes[0].set_title('Box Loss') axes[0].legend() # 分类损失 axes[1].plot(df['epoch'], df['train/cls_loss'], label='train_cls') axes[1].plot(df['epoch'], df['val/cls_loss'], label='val_cls') axes[1].set_title('Cls Loss') axes[1].legend() # 分布焦点损失 axes[2].plot(df['epoch'], df['train/dfl_loss'], label='train_dfl') axes[2].plot(df['epoch'], df['val/dfl_loss'], label='val_dfl') axes[2].set_title('DFL Loss') axes[2].legend() plt.tight_layout() plt.savefig('steel_loss_curve.png', dpi=150)这段代码把三组损失曲线画在同一张图里。画出来后如果发现val/box_loss在100个epoch后不再下降,说明模型容量或数据增强达到瓶颈,继续训练只会过拟合分类分支。使用GTX1660Ti跑YOLOv8s在imgsz=800下,每个epoch大约需要2到3分钟,200个epoch约7到10小时,这个时间成本可以接受。如果时间紧张,可以训练YOLOv8n,模型参数量降低一半以上,mAP大约下降3到4个点,但推理速度从20毫秒降到7毫秒。
3.3 缺陷检测权重的评估标准与过拟合判断
训练完成后,不要只看results.csv里的mAP,要跑一遍测试集,逐类看各类缺陷的AP值。YOLOv8在训练过程中会输出验证集混淆矩阵,可以用predict方法在独立的测试目录上做推理,然后对比标注文件计算每个类别的精确率和召回率。
from ultralytics import YOLO model = YOLO('steel_run/exp1/weights/best.pt') # 在验证集上跑评估 metrics = model.val( data='steel_dataset.yaml', split='val', imgsz=800, conf=0.25, iou=0.6 ) print(metrics.box.map) # mAP50-95 print(metrics.box.map50) # mAP50 # 逐类查看AP print(metrics.box.ap_class_index) print(metrics.box.ap)钢材缺陷检测里最常见的现象是crazing裂纹这类缺陷的AP普遍偏低。裂纹形态不规则、对比度低、有时只出现在小范围纹理异常区域,即使人工标注也有较大主观性。如果某类AP特别低,优先检查的是标注质量,看看这类缺陷的标注框是否有漏标,在600张训练集里漏标5个框就可能让AP下降好几个点。其次是数据增强,hsv_h如果过高会让裂纹的灰度变化被色相干扰,我在3.1里的参数就是针对这个问题调的。
4. Qt界面GUI实现的完整方案,PyQt5封装YOLOv8推理全流程
4.1 推理引擎封装与并行推理避免界面卡顿
Qt界面这部分的目标很明确:让现场人员双击启动程序,选择图片或视频,点击按钮就能看到检测结果和置信度标注。技术栈选择PyQt5加Ultralytics自带推理接口,不额外引入ONNX Runtime或TensorRT的复杂度。但有一个核心问题必须先解决:YOLOv8推理和UI必须分开线程执行,否则推理期间界面会卡死。
PyQt5的QThread方式实现这一步最直接。Worker类继承QThread,重写run()方法执行推理,通过信号把检测结果的图片和统计数据传回主线程。主线程只负责更新QLabel上的画面。
import sys import cv2 import numpy as np from PyQt5.QtCore import QThread, pyqtSignal, Qt from PyQt5.QtGui import QImage, QPixmap from PyQt5.QtWidgets import QMainWindow, QLabel, QPushButton, QFileDialog, QVBoxLayout, QWidget, QProgressBar, QComboBox from ultralytics import YOLO class DetectWorker(QThread): # 定义信号:检测完成的图片、检测到的目标数量和类别分布 result_ready = pyqtSignal(object, int, dict) error_occurred = pyqtSignal(str) def __init__(self, model_path, source_path, conf_thres=0.25, iou_thres=0.6, parent=None): super().__init__(parent) self.model = YOLO(model_path) self.source_path = source_path self.conf_thres = conf_thres self.iou_thres = iou_thres def run(self): try: # 使用stream模式处理视频,避免一次性加载所有帧导致内存不足 results = self.model.predict( source=self.source_path, conf=self.conf_thres, iou=self.iou_thres, imgsz=800, stream=True, device='0' ) for i, r in enumerate(results): im = r.plot() # 绘制了检测框的结果图 class_counts = {} boxes = r.boxes if boxes is not None: cls_ids = boxes.cls.cpu().numpy().astype(int) for cid in cls_ids: class_name = r.names[cid] class_counts[class_name] = class_counts.get(class_name, 0) + 1 self.result_ready.emit(im, i, class_counts) except Exception as e: self.error_occurred.emit(str(e))Worker构造时传入model_path和source_path,run()内部使用model.predict的stream=True参数,这样返回的是生成器,每处理完一帧就发一次信号。r.plot()返回的是BGR格式的numpy数组,这个数组可以直接转成Qt的QImage显示。device='0'指定用GPU推理,注意这里device参数和训练时完全一致。
Qt主窗口负责接收信号并更新界面。一个关键的处理是图像缩放:相机图片可能是4000x3000分辨率,直接塞进QLabel会超出界面范围,需要按比例缩放到QLabel的可用尺寸,同时保持长宽比。这里要用Qt.KeepAspectRatio的缩放模式,并且转换颜色通道时要指定QImage.Format_BGR888,因为OpenCV读出来的是BGR格式,直接用Format_RGB888会让画面偏蓝。
class MainWindow(QMainWindow): def __init__(self): super().__init__() self.setWindowTitle('钢材表面缺陷检测系统 - YOLOv8') self.setMinimumSize(1200, 800) self.worker = None # UI布局 central = QWidget(self) layout = QVBoxLayout(central) self.image_label = QLabel('请选择图片或视频开始检测') self.image_label.setAlignment(Qt.AlignCenter) self.image_label.setStyleSheet('background: #2b2b2b; color: #fff;') self.image_label.setMinimumHeight(600) layout.addWidget(self.image_label) btn_row = QVBoxLayout() self.btn_open_image = QPushButton('选择图片') self.btn_open_video = QPushButton('选择视频') self.btn_open_image.clicked.connect(self.open_image) self.btn_open_video.clicked.connect(self.open_video) btn_row.addWidget(self.btn_open_image) btn_row.addWidget(self.btn_open_video) layout.addLayout(btn_row) self.status_label = QLabel('就绪') layout.addWidget(self.status_label) self.setCentralWidget(central) def open_image(self): fname, _ = QFileDialog.getOpenFileName( self, '选择图片', '', '图片文件 (*.jpg *.jpeg *.png *.bmp)') if not fname: return self.status_label.setText(f'正在检测: {fname}') self.worker = DetectWorker('steel_run/exp1/weights/best.pt', fname) self.worker.result_ready.connect(self.update_image) self.worker.error_occurred.connect(self.handle_error) self.worker.start() def update_image(self, img_array, frame_idx, class_counts): h, w, ch = img_array.shape qimg = QImage(img_array.data, w, h, ch * w, QImage.Format_BGR888) pixmap = QPixmap.fromImage(qimg) # 按QLabel大小缩放,保持长宽比 scaled = pixmap.scaled( self.image_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation) self.image_label.setPixmap(scaled) total = sum(class_counts.values()) detail = ', '.join([f'{k}: {v}' for k, v in class_counts.items()]) self.status_label.setText(f'检测到 {total} 个缺陷 | {detail}')update_image先根据图像数组创建QImage,再构造QPixmap,最后按label尺寸缩放。注意这里必须把img_array.data引用传给QImage后不能立刻释放img_array,否则Qt访问到悬空内存会崩溃。缩放使用SmoothTransformation模式,虽然慢一点但表面缺陷的小目标能保持清晰度。为什么设置ch * w作为bytesPerLine?因为OpenCV的numpy数组默认按行连续存储,当图像宽度不是4的倍数时,这个参数决定了Qt按行读取的步长,错误设置会导致图像出现错位条纹。
4.2 Qt界面GUI操作流程与置信度参数调节
视频检测时,QThread信号是按帧触发的,每一帧都会更新一次界面。但如果视频帧率高于UI刷新率,信号队列会积压,界面越来越卡。解决方案是设置QThread.setPriority(QThread.HighPriority)降低线程优先级,同时容忍界面跳帧显示。另一个可行方案是只更新最新的信号:在update_image里用一个self.latest_result = img_array保存最新帧,再用QTimer每100毫秒刷新一次,减少QPixmap构造次数。
置信度阈值conf_thres应该做成界面上的实时可调参数,而不是写死在构造函数里。钢材缺陷检测的典型误报来源有两类:一类是表面水滴造成的反光被识别为patches,另一类是氧化铁皮边缘被识别为inclusion。前者通过提高置信度到0.3到0.35能滤除大部分,后者需要结合现场经验判断,不建议把置信度调得太高,否则真实的小缺陷也会被滤掉。
def update_conf_threshold(self, value): # value是滑动条的数值,0到100映射到0到1的置信度 self.conf_thres = value / 100.0 if hasattr(self, 'worker') and self.worker is not None: self.worker.conf_thres = self.conf_thres滑动条联动到worker的conf_thres属性,这样在视频推理过程中也能实时调整,不用重启程序。这个设计在产线试用阶段很实用,现场人员会根据自己的判断不断微调阈值,直到找到误报和漏报的平衡点。
Qt界面里另一个容易被忽略的点是中文路径。现场质检人员把图片文件命名为缺陷样本20250315_划痕.jpg是非常常见的情况。Ultralytics的predict接口对包含中文的路径处理依赖底层文件系统编码,Windows上容易出现UnicodeDecodeError。稳妥做法是在选择文件后先复制到程序目录下的cache文件夹,用ASCII文件名保存,再传入推理接口。
4.3 权重文件打包与界面发布,离线环境怎么部署
程序调试完成后,需要打包成独立的exe给现场用。PyInstaller是PyQt5项目最成熟的打包工具。但含YOLOv8的打包有几个坑:Ultralytics在运行时会动态加载一些配置文件和字体文件,如果直接打包主脚本,运行时可能找不到Arial.ttf导致绘图文字变成方块;模型权重文件.pt虽然会被Python代码引用,但PyInstaller不会自动收集非代码资源。
# 打包命令,注意需要显式添加ultralytics的配置文件和字体 pyinstaller -F -w \ --name "SteelDetect" \ --add-data "steel_run/exp1/weights/best.pt;weights" \ --collect-all ultralytics \ --collect-all torch \ app.py--collect-all ultralytics收集所有包内资源,--collect-all torch是必须的,因为PyTorch有大量动态库依赖。-F生成单文件模式,启动时会自解压到临时目录,首次启动可能有3到5秒延迟。如果闪烁窗口会引起现场人员困惑,可以先用-w隐藏控制台,再在打包命令里加上--icon指定应用图标。--add-data在Windows下用分号分隔源和目标目录。
打包后单文件可能达到1.5GB以上,因为PyTorch的CUDA运行时全量打进去了。如果想瘦身,可以在打包前把Ultralytics的依赖精简,去掉不用的模型配置,但保险起见不建议新手优化这点体积。磁盘空间充裕就保持原样。
5. 三个落地技巧:批量推理、缺陷分类统计与界面热更新
5.1 批量推理文件夹图片并导出Excel统计表
产线质检场景往往不是单张图片检测,而是每隔几秒生成一张截图,一个班次下来几百张图片需要批量处理。用Qt界面逐张点击不现实,常见做法是程序启动时自动扫描指定目录下的所有图片,批量推理后把结果汇总成一份CSV或Excel。
import os import glob import pandas as pd from ultralytics import YOLO model = YOLO('steel_run/exp1/weights/best.pt') def batch_inference_and_export(image_dir, output_csv): image_files = glob.glob(os.path.join(image_dir, '*.jpg')) + \ glob.glob(os.path.join(image_dir, '*.png')) results = [] for img_path in image_files: r = model.predict(source=img_path, conf=0.25, iou=0.6, imgsz=800, verbose=False)[0] # 统计每个类别的数量和最大置信度 class_counts = {} confidences = {} for box in r.boxes: cls_id = int(box.cls[0].cpu().numpy()) cls_name = r.names[cls_id] conf = float(box.conf[0].cpu().numpy()) class_counts[cls_name] = class_counts.get(cls_name, 0) + 1 if cls_name not in confidences or conf > confidences[cls_name]: confidences[cls_name] = conf row = {'图片路径': img_path, '总缺陷数': sum(class_counts.values())} for name in r.names.values(): row[f'数量_{name}'] = class_counts.get(name, 0) row[f'最大置信度_{name}'] = round(confidences.get(name, 0), 4) results.append(row) df = pd.DataFrame(results) df.to_csv(output_csv, index=False, encoding='utf-8-sig') print(f'已导出 {len(results)} 条检测记录到 {output_csv}')这条批量推理的核心价值在于输出格式直接贴合质检追溯需求。utf-8-sig编码可以在Excel里正常显示中文路径,不需要额外转码。导出每一类的最大置信度有什么意义?最大置信度代表这张图片里最可能是该类缺陷的区域,质检人员可以按最大置信度降序排列,优先查看最存疑的样本。
5.2 Qt界面显示检测结果时的类别颜色映射与实时刷新优化
YOLOv8的r.plot()绘图函数默认给每个类别分配一种颜色,但相邻类别的颜色可能相近。钢材缺陷里inclusion和pitted_surface经常出现在同一区域,颜色相近时现场人员容易看混。更好的做法是手动指定每个类别的确切RGB值。
CLASS_COLORS = { 'crazing': (0, 0, 255), # 红色 'inclusion': (0, 255, 0), # 绿色 'patches': (255, 0, 0), # 蓝色 'pitted_surface': (0, 255, 255), # 黄色 'rolled-in_scale': (255, 0, 255), # 品红 'scratches': (255, 255, 0), # 青色 } def draw_with_custom_colors(img, results, class_names): """在图像上绘制带固定颜色和中文标签的检测框""" import cv2 as cv font = cv.FONT_HERSHEY_SIMPLEX for box in results.boxes: cls_id = int(box.cls[0].cpu().numpy()) conf = float(box.conf[0].cpu().numpy()) x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int) color = CLASS_COLORS.get(class_names[cls_id], (255, 255, 255)) cv.rectangle(img, (x1, y1), (x2, y2), color, 2) label = f'{class_names[cls_id]} {conf:.2f}' # 文字背景色垫底,避免被钢材纹理干扰 (tw, th), baseline = cv.getTextSize(label, font, 0.6, 1) cv.rectangle(img, (x1, y1-th-baseline), (x1+tw, y1), color, -1) cv.putText(img, label, (x1, y1-baseline), font, 0.6, (255,255,255), 1) return img除了自定义颜色外,这里额外做了一个操作:在文字背后画实心矩形垫底。钢材表面纹理复杂,直接绘制文字会和背景噪声重叠,可读性差。实心背景代价是几个像素的高,不牺牲推理速度。如果想要中文标签写入图片,需要下载中文字体并使用cv.putText的freetype版本,这会增加依赖,建议界面语言用中文但图片绘制文字保持英文,避免现场人员调试时困惑。
5.3 验证最终权重效果的四步检查法与常见问题速查
部署到现场前,用一套固定步骤验证权重质量非常关键。我建议按下面四步做:第一步,用训练时用的验证集跑一遍model.val(),记录mAP50和mAP50-95。第二步,选择10张典型缺陷图片,包括每类缺陷各2张,单张跑推理并人工核对所有标注框是否正确对应到缺陷区域。第三步,选择5张无缺陷的正常钢材图片,确认没有误报。第四步,用一段产线实际视频连续跑5分钟,统计每帧推理耗时是否稳定在30毫秒以内。
# 用YOLOv8自带命令行快速验证单张图片 yolo detect predict model=steel_run/exp1/weights/best.pt source=./test_samples/crazing_001.jpg imgsz=800 conf=0.25 iou=0.6 save如果发现检测结果不理想,先检查三个地方。第一,训练时的imgsz和推理时的imgsz必须一致,推理时把分辨率从640改成800,模型在640下学到的小目标特征在800分辨率下计算感受野时会错位。第二,检查是否误把last.pt当成best.pt加载,两个文件在大小和性能上有明显差异。第三,检查推理时是否忘记设置device='0',CPU推理速度和GPU能差3到5倍。
如果现场反馈某类缺陷总是漏检,优先查的是训练集里该类缺陷的标注是否有遗漏。钢材缺陷检测任务中标注质量的影响远大于模型结构的影响,一个漏标可能导致模型把正常区域预测为背景。用第5.1节的批量推理工具跑一遍训练集,把预测置信度大于0.5但标注文件中没有记录的地方标注出来,人工确认是漏标还是误检。这个方法能快速发现数据问题,比反复调模型参数更有效。
本文还有配套的精品资源,点击获取