基于YOLOv8与Grad-CAM的铝合金表面瑕疵识别系统:从弱监督标注到PyQt5可视化界面的完整实战
基于YOLOv8与Grad-CAM的铝合金表面瑕疵识别系统:从弱监督标注到PyQt5可视化界面的完整实战
源码获取:https://mbd.pub/o/bread/YZaTmZtuZg==
摘要: 本文详细介绍了一套完整的铝合金表面瑕疵识别系统的开发过程。系统采用弱监督学习思路,利用ResNet18分类模型结合Grad-CAM热力图自动生成伪边界框标注,解决了工业场景下目标检测数据集缺乏精细标注的难题。基于伪标注数据训练YOLOv8目标检测模型,在11类瑕疵检测任务上达到了**mAP@0.5=92.55%**的精度。同时,本文完整讲解了基于PyQt5的图形用户界面开发,涵盖瑕疵检测、训练可视化、数据集浏览和模型训练控制四大功能模块。所有代码均已开源,适合作为计算机视觉、工业质检方向的毕业设计或项目实战参考。
一、前言:工业质检智能化的迫切需求
1.1 项目背景
铝合金作为现代工业中应用最广泛的金属材料之一,被大量应用于航空航天、汽车制造、建筑装饰、电子产品外壳等领域。在铝合金型材的生产过程中,由于原材料质量、加工工艺、环境因素等多方面原因,型材表面不可避免地会出现各种瑕疵,如擦花、漏底、桔皮、碰伤、涂层开裂等。这些表面缺陷不仅影响产品的美观性,更严重的是会降低产品的力学性能、耐腐蚀性能和使用寿命,直接关系到产品的质量安全和市场竞争力。
传统的铝合金表面瑕疵检测主要依赖人工目视检查。然而,人工检测存在诸多难以克服的弊端:首先,检测效率低下,一名熟练质检员每分钟只能检查有限数量的型材;其次,长时间的高强度工作容易导致视觉疲劳,漏检率和误检率随工作时间显著上升;再次,人工检测标准难以统一,不同质检员对同一瑕疵的判定可能存在差异;最后,人工成本持续攀升,尤其是在劳动力密集型的制造业中,质检环节的人力成本已成为企业的重要负担。
随着深度学习技术的快速发展,基于计算机视觉的自动化瑕疵检测成为解决上述问题的有效途径。卷积神经网络(CNN)具有强大的特征提取能力,能够从海量图像数据中自动学习瑕疵的视觉特征,实现高效、准确、稳定的自动化检测。然而,将深度学习技术应用于工业瑕疵检测场景,面临着几个核心挑战:
第一,标注成本高昂。 目标检测任务需要精确的边界框标注(Bounding Box Annotation),即不仅要判断图像中存在何种瑕疵,还要精确框出瑕疵的位置。工业图像的分辨率通常很高,瑕疵区域往往较小且形态不规则,人工标注工作量大、专业性强、成本极高。
第二,瑕疵类型多样。 铝合金表面的瑕疵类型多达十余种,不同类型的瑕疵在颜色、纹理、形状、大小上差异显著,甚至同一类型的瑕疵在不同光照条件、拍摄角度下也会呈现出不同的视觉表现,这对模型的泛化能力提出了很高要求。
第三,实时性要求。 工业生产线上通常要求检测系统能够达到实时或近实时的处理速度,以匹配生产节拍。这意味着模型不仅要有高准确率,还要有低延迟、高吞吐量的推理性能。
1.2 项目核心思路
针对上述挑战,本项目设计了一套"分类先行、弱监督标注、检测跟进、界面封装"的完整技术路线:
第一步,图像分类。 利用ResNet18预训练模型,在只有图像级别类别标签(Image-level Label)的数据集上训练分类模型。这一步骤的目标是建立一个能够理解"每种瑕疵长什么样"的基础模型,分类准确率达到94.39%。
第二步,弱监督伪标注。 这是本项目的核心创新点。我们利用Grad-CAM(Gradient-weighted Class Activation Mapping)技术,从训练好的分类模型中提取类别激活热力图。Grad-CAM能够高亮显示图像中对分类决策贡献最大的区域,这些区域往往就是瑕疵所在的位置。通过对热力图进行二值化、轮廓检测和最小外接矩形提取,我们自动生成了边界框标注,并将其转换为YOLO格式。整个过程无需任何人工画框操作,极大地降低了标注成本。
第三步,目标检测训练。 将自动生成的伪标注数据输入YOLOv8模型进行训练。YOLOv8作为当前业界领先的实时目标检测算法,具有速度快、精度高、易部署的特点。我们选用轻量级的YOLOv8n作为基础模型,在保持较高检测精度的同时确保推理速度满足实时性要求。
第四步,PyQt5界面封装。 为了让系统具备良好的可用性和演示效果,我们使用PyQt5开发了一套完整的图形用户界面,包含瑕疵检测、训练结果可视化、数据集浏览和模型训练控制四大功能模块。界面支持中文显示,操作直观便捷,可直接部署于生产环境或用于学术答辩演示。
1.3 技术栈概览
| 层级 | 技术/工具 | 版本/说明 |
|---|---|---|
| 深度学习框架 | PyTorch + torchvision | 2.5.1 + CUDA 12.1 |
| 目标检测 | Ultralytics YOLOv8 | yolov8n预训练模型 |
| 图像分类 | torchvision ResNet18 | ImageNet预训练权重 |
| 弱监督标注 | pytorch-grad-cam | Grad-CAM热力图提取 |
| 图形界面 | PyQt5 | 5.15+ |
| 图像处理 | OpenCV + Pillow | 4.5+ / 9.0+ |
| 数据可视化 | matplotlib + pandas | 3.5+ / 1.3+ |
| 开发语言 | Python | >= 3.8 |
本文将按上述技术路线,逐模块、逐文件地讲解系统的完整实现过程。所有代码均基于真实项目,建议读者边阅读边动手实践。
二、环境配置与项目搭建
2.1 硬件与系统环境
本项目在以下环境中开发和测试:
- 操作系统: Windows 10/11(64位)
- CPU: Intel Core i7/i9 或 AMD Ryzen 7/9 系列
- GPU: NVIDIA GeForce RTX 3060/4060 或更高(显存 >= 8GB)
- 内存: >= 16GB DDR4/DDR5
- 硬盘: SSD,预留至少 10GB 空间
虽然项目支持CPU模式运行,但深度学习训练强烈建议使用NVIDIA GPU以加速计算。如果没有独立显卡,可以将batch size适当调小,使用CPU进行训练,只是训练时间会大幅延长。
2.2 Python环境配置
建议使用Anaconda或Miniconda创建独立的虚拟环境,避免与系统Python或其他项目产生依赖冲突。
# 创建虚拟环境(Python 3.10)
conda create -n aluminum_defect python=3.10 -y
# 激活环境
conda activate aluminum_defect
2.3 依赖包安装
项目的依赖包列表记录在 requirements.txt 中,内容如下:
# 铝合金瑕疵识别系统依赖
# Python >= 3.8
# 深度学习框架(CUDA 12.1)
torch>=2.5.1
torchvision==0.20.1+cu121
--extra-index-url https://download.pytorch.org/whl/cu121
# YOLOv8 目标检测
ultralytics>=8.0
# 图像处理
opencv-python>=4.5
Pillow>=9.0
# 数据处理与可视化
numpy<2.0
matplotlib>=3.5
pandas>=1.3
# 弱监督学习(Grad-CAM)
grad-cam>=1.4
# GUI
PyQt5>=5.15
安装命令:
pip install -r requirements.txt
特别注意CUDA版本匹配问题。 PyTorch的GPU版本需要与系统安装的CUDA Toolkit版本严格匹配。如果你的CUDA版本不是12.1,需要访问PyTorch官网查询对应的安装命令。例如,对于CUDA 11.8:
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
安装完成后,验证PyTorch是否能正确调用GPU:
import torch
print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA是否可用: {torch.cuda.is_available()}")
if torch.cuda.is_available():
print(f"CUDA版本: {torch.version.cuda}")
print(f"GPU型号: {torch.cuda.get_device_name(0)}")
2.4 项目目录结构
建议按照以下结构组织项目文件:
aluminum-defect-detection/
├── configs/
│ └── config.yaml # 项目全局配置文件
├── dataset/
│ ├── guangdong_round1_train2_20180916/ # 原始数据集
│ │ └── 瑕疵样本/ # 按类别组织的瑕疵图片
│ ├── annotations/ # Grad-CAM生成的伪标注文件
│ ├── splits/ # 数据集划分文件
│ │ ├── train.txt
│ │ └── val.txt
│ └── yolo/ # YOLO格式数据集
│ ├── dataset.yaml
│ ├── images/train/
│ ├── images/val/
│ ├── labels/train/
│ └── labels/val/
├── output/ # 训练输出目录
│ ├── classifier_best.pth # 最佳分类模型
│ ├── class_to_idx.json # 类别映射文件
│ ├── yolov8_train/ # YOLO训练结果
│ │ ├── weights/best.pt
│ │ ├── results.csv
│ │ ├── confusion_matrix.png
│ │ └── ...
│ └── ultralytics_config/ # Ultralytics配置缓存
├── src/
│ ├── data/ # 数据加载模块
│ │ ├── classification_dataset.py
│ │ └── transforms.py
│ ├── models/ # 模型定义模块
│ │ └── classifier.py
│ └── gui/ # PyQt5图形界面模块
│ ├── main_window.py
│ ├── detector_tab.py
│ ├── training_results_tab.py
│ ├── dataset_tab.py
│ ├── training_control_tab.py
│ ├── detector_worker.py
│ └── utils.py
├── scripts/ # 可执行脚本
│ ├── train_classifier.py
│ ├── generate_pseudo_labels.py
│ ├── train_detector.py
│ ├── prepare_yolo_dataset.py
│ ├── split_dataset.py
│ ├── analyze_dataset.py
│ └── run_gui.py
├── requirements.txt
└── README.md
2.5 配置文件详解
configs/config.yaml 是项目的全局配置文件,集中管理数据路径、模型参数和训练超参:
dataset:
train_dir: "dataset/guangdong_round1_train2_20180916/瑕疵样本"
test_a_dir: "dataset/guangdong_round1_test_a_20180916"
test_b_dir: "dataset/guangdong_round1_test_b_20181009"
defect_types_file: "dataset/初赛瑕疵类型说明.xlsx"
model:
classifier:
backbone: "resnet18"
pretrained: true
detector:
model: "yolov8n"
pretrained: true
training:
classifier:
epochs: 30
batch_size: 32
learning_rate: 0.001
image_size: 224
detector:
epochs: 100
batch_size: 16
learning_rate: 0.01
image_size: 640
output:
dir: "output"
classifier_weights: "output/classifier_best.pth"
detector_weights: "output/detector_best.pt"
annotations_dir: "dataset/annotations"
配置文件采用YAML格式,层次清晰,便于修改。在实际使用中,只需修改此文件中的参数,无需改动代码,即可适配不同的硬件环境和实验需求。
三、数据集介绍与预处理
3.1 数据集来源与结构
本项目使用的数据集来自广东某铝材厂表面瑕疵检测竞赛初赛数据集(简称"广东铝材数据集")。该数据集是工业瑕疵检测领域较为经典的公开数据集之一,被众多学术研究和工程项目所采用。
原始数据集的目录结构如下:
guangdong_round1_train2_20180916/
├── 瑕疵样本/
│ ├── 不导电/
│ │ ├── 不导电20180830131536对照样本.jpg
│ │ ├── 不导电20180830131551对照样本.jpg
│ │ └── ...
│ ├── 凸粉/
│ ├── 擦花/
│ ├── 桔皮/
│ ├── 横条压凹/
│ ├── 涂层开裂/
│ ├── 漏底/
│ ├── 碰伤/
│ ├── 脏点/
│ ├── 起坑/
│ └── 其他/
└── 无瑕疵样本/
├── xxx.jpg
└── ...
数据集包含两个主要部分:瑕疵样本和无瑕疵样本。瑕疵样本按照瑕疵类型分为11个类别文件夹,每个文件夹内存放属于该类别的瑕疵图片。图片命名通常包含瑕疵类型和时间戳信息,便于追溯和管理。
3.2 瑕疵类别详解
系统需要识别的11种铝合金表面瑕疵及其视觉特征如下:
| 类别ID | 类别名称 | 视觉特征描述 |
|---|---|---|
| 0 | 不导电 | 表面导电性异常区域,通常表现为颜色或光泽与周围明显不同的斑块 |
| 1 | 其他 | 不属于标准类别的其他类型瑕疵,作为兜底类别 |
| 2 | 凸粉 | 表面凸起的粉末状缺陷,呈颗粒状或团状分布 |
| 3 | 擦花 | 表面擦伤或划痕,通常呈线状或带状,方向随机 |
| 4 | 桔皮 | 表面纹理类似桔皮的波纹状缺陷,光泽不均匀 |
| 5 | 横条压凹 | 横向条状的压痕或凹陷,通常沿型材宽度方向分布 |
| 6 | 涂层开裂 | 表面保护涂层出现的裂纹,呈不规则网状或线状 |
| 7 | 漏底 | 涂层缺失导致底色露出,呈块状或片状,颜色对比明显 |
| 8 | 碰伤 | 碰撞造成的局部损伤,通常呈不规则凹坑或划痕组合 |
| 9 | 脏点 | 表面附着的污渍或斑点,与基材结合不紧密 |
| 10 | 起坑 | 表面出现的凹坑或气泡破裂痕迹,呈圆形或椭圆形 |
这些瑕疵在颜色、纹理、形状和尺度上差异很大。例如,"擦花"通常表现为细长的线条,而"漏底"则是不规则的块状区域;"脏点"通常较小且分散,而"横条压凹"则可能贯穿整个型材表面。这种高度的类内差异和类间相似性,给自动检测带来了不小的挑战。
3.3 数据集统计分析
在正式训练之前,我们首先对数据集进行统计分析,了解各类别的样本分布情况。scripts/analyze_dataset.py 脚本实现了这一功能:
def analyze_dataset(dataset_path):
dataset_path = Path(dataset_path)
stats = defaultdict(list)
total_images = 0
for category_dir in sorted(dataset_path.iterdir()):
if category_dir.is_dir():
category_name = category_dir.name
image_files = []
for ext in ['*.jpg', '*.jpeg', '*.png', '*.bmp']:
image_files.extend(category_dir.glob(ext))
image_files = list(set(image_files))
count = len(image_files)
stats[category_name] = image_files
total_images += count
print(f"{category_name:20s}: {count:4d} 张图片")
return stats, total_images
运行该脚本后,可以得到类似如下的统计输出:
============================================================
数据集分析: dataset/guangdong_round1_train2_20180916/瑕疵样本
============================================================
不导电 : 120 张图片
凸粉 : 95 张图片
擦花 : 310 张图片
桔皮 : 150 张图片
横条压凹 : 80 张图片
涂层开裂 : 65 张图片
漏底 : 280 张图片
碰伤 : 110 张图片
脏点 : 200 张图片
起坑 : 90 张图片
其他 : 45 张图片
============================================================
类别总数: 11
图片总数: 1545
============================================================
从统计结果可以看出,数据集存在明显的类别不平衡问题。"擦花"和"漏底"的样本数量最多(超过250张),而"其他"和"涂层开裂"的样本数量较少(不足70张)。类别不平衡是工业瑕疵检测数据集的普遍现象,因为某些瑕疵在实际生产中确实发生频率较低。在后续的训练中,我们需要通过数据增强、类别权重调整等策略来缓解这一问题。
3.4 数据集划分——分层采样
为了公平地评估模型性能,我们需要将数据集划分为训练集和验证集。考虑到类别不平衡的问题,简单的随机划分可能导致某些类别在验证集中完全没有样本或样本极少。因此,我们采用**分层采样(Stratified Sampling)**策略,确保训练集和验证集中各类别的比例与原始数据集一致。
scripts/split_dataset.py 实现了基于scikit-learn的 train_test_split 函数的分层划分:
from sklearn.model_selection import train_test_split
def split_dataset(image_data, test_size=0.2, random_seed=42):
paths = [item[0] for item in image_data]
labels = [item[1] for item in image_data]
train_paths, val_paths, train_labels, val_labels = train_test_split(
paths, labels,
test_size=test_size,
random_state=random_seed,
stratify=labels
)
return list(zip(train_paths, train_labels)), list(zip(val_paths, val_labels))
划分比例为 8:2(训练集80%,验证集20%),随机种子设为42以保证结果可复现。分层采样确保了每个类别在训练集和验证集中都有合理的代表性,避免了因数据划分不当导致的评估偏差。
划分完成后,会在 dataset/splits/ 目录下生成 train.txt 和 val.txt 两个文件,每行记录一张图片的相对路径(如 瑕疵样本/擦花/擦花20180830164545对照样本.jpg)。后续的数据加载器将读取这两个文件,只加载对应划分的样本。
3.5 数据集加载器设计
src/data/classification_dataset.py 中定义了 DefectClassificationDataset 类,这是PyTorch标准的 Dataset 子类,用于加载分类任务的数据:
class DefectClassificationDataset(Dataset):
def __init__(self, root_dir, split_file=None, transform=None, class_to_idx=None):
self.root_dir = Path(root_dir)
self.transform = transform
self.samples = []
self.targets = []
self.parent_dir = self.root_dir.parent
# 加载划分文件,筛选有效样本
if split_file and os.path.exists(split_file):
with open(split_file, 'r', encoding='utf-8') as f:
valid_files = set(line.strip().replace('\\', '/') for line in f if line.strip())
else:
valid_files = None
# 构建类别映射
if class_to_idx is None:
self.class_to_idx = {}
idx = 0
for class_dir in sorted(self.root_dir.iterdir()):
if class_dir.is_dir():
self.class_to_idx[class_dir.name] = idx
idx += 1
else:
self.class_to_idx = class_to_idx
self.idx_to_class = {v: k for k, v in self.class_to_idx.items()}
# 加载样本
for class_name, class_idx in self.class_to_idx.items():
class_dir = self.root_dir / class_name
if not class_dir.is_dir():
continue
for img_path in class_dir.glob('*.jpg'):
if valid_files is not None:
rel_path = str(img_path.relative_to(self.parent_dir)).replace('\\', '/')
if rel_path not in valid_files:
continue
self.samples.append(str(img_path))
self.targets.append(class_idx)
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
img_path = self.samples[idx]
label = self.targets[idx]
image = Image.open(img_path).convert('RGB')
if self.transform:
image = self.transform(image)
return image, label
该数据集类的设计考虑了以下关键点:
- 划分文件支持: 通过
split_file参数,可以只加载训练集或验证集的样本,避免数据泄露。 - 类别映射复用: 训练集自动构建
class_to_idx映射,验证集通过传入相同的映射保证标签一致性。 - 跨平台兼容: 统一将Windows反斜杠路径转换为正斜杠,确保在不同操作系统上都能正确匹配。
- 类别权重计算: 提供了
get_class_weights()方法,用于计算逆频率权重,后续可用于加权损失函数缓解类别不平衡。
数据增强变换分为训练集和验证集两种:
train_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transforms.RandomRotation(15),
transforms.ColorJitter(brightness=0.2, contrast=0.2),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
val_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
训练集采用了丰富的数据增强策略:随机水平翻转、随机垂直翻转、随机旋转(±15度)、颜色抖动(亮度/对比度变化)。这些增强操作有效地扩充了训练样本的多样性,提高了模型的泛化能力。验证集只进行必要的尺寸调整和归一化,保证评估结果的一致性。归一化使用的均值和标准差是ImageNet数据集的统计值,因为我们将使用ImageNet预训练权重进行迁移学习。
四、图像分类模型——ResNet18的设计与训练
4.1 为什么选择ResNet18
在深度学习图像分类领域,ResNet(Residual Network,残差网络)是一个具有里程碑意义的架构。ResNet由何恺明等人在2015年提出,通过引入**残差连接(Skip Connection)**解决了深层网络中的梯度消失和退化问题,使得训练上百层甚至上千层的深度网络成为可能。
本项目选择 ResNet18 作为分类骨干网络,主要基于以下考虑:
第一,模型复杂度适中。 ResNet18是ResNet家族中最轻量的版本之一,总参数量约1100万。相比于ResNet50(2500万参数)或ResNet101(4400万参数),ResNet18的训练速度更快,显存占用更少,更适合在个人工作站或有限计算资源的环境下进行实验和迭代。
第二,特征提取能力足够。 虽然ResNet18较浅,但对于工业瑕疵分类这种目标明确、背景相对单一的任务,其特征提取能力已经完全足够。事实上,在我们的实验中,ResNet18的验证准确率达到94.39%,证明了其有效性。
第三,便于Grad-CAM应用。 Grad-CAM需要从网络的深层卷积层提取梯度信息生成热力图。ResNet18的结构简洁清晰,其 layer4 输出的特征图尺寸为 [B, 512, 7, 7],既保留了足够的空间分辨率用于定位,又具有足够的语义深度用于分类,非常适合后续的热力图生成。
第四,预训练权重丰富。 ResNet18在PyTorch的 torchvision.models 中提供了官方ImageNet预训练权重,可以直接加载使用。预训练模型已经学习到了大量通用的低级视觉特征(如边缘、纹理、颜色)和高级语义特征,通过迁移学习可以大幅减少训练时间和数据需求。
4.2 模型架构设计
src/models/classifier.py 中定义了 DefectClassifier 类:
class DefectClassifier(nn.Module):
def __init__(self, num_classes, pretrained=True):
super(DefectClassifier, self).__init__()
# 加载预训练ResNet18
self.backbone = models.resnet18(pretrained=pretrained)
# 替换最后的全连接层
in_features = self.backbone.fc.in_features
self.backbone.fc = nn.Linear(in_features, num_classes)
self.num_classes = num_classes
def forward(self, x):
return self.backbone(x)
def get_features(self, x):
"""提取layer4特征图,用于Grad-CAM"""
x = self.backbone.conv1(x)
x = self.backbone.bn1(x)
x = self.backbone.relu(x)
x = self.backbone.maxpool(x)
x = self.backbone.layer1(x)
x = self.backbone.layer2(x)
x = self.backbone.layer3(x)
x = self.backbone.layer4(x)
return x
模型设计的核心操作是替换全连接层。原始ResNet18的最后一层 fc 输出1000维向量(对应ImageNet的1000个类别),我们需要将其替换为输出11维向量的新层(对应11种瑕疵类别)。in_features 为512,是 layer4 输出的通道数。
此外,我们额外实现了 get_features 方法。该方法前向传播到 layer4 结束,不经过全局平均池化层和全连接层,输出的特征图尺寸为 [B, 512, 7, 7]。这个特征图既保留了空间结构信息(7x7的空间分辨率),又包含了深层语义信息(512个通道),是Grad-CAM生成热力图的关键输入。
4.3 训练流程详解
scripts/train_classifier.py 实现了完整的分类模型训练流程。下面逐段解析关键代码。
训练一个epoch:
def train_epoch(model, train_loader, criterion, optimizer, device):
model.train()
running_loss = 0.0
correct = 0
total = 0
for batch_idx, (inputs, targets) in enumerate(train_loader):
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
if batch_idx % 10 == 0:
print(f' Batch {batch_idx}/{len(train_loader)} | '
f'Loss: {running_loss/(batch_idx+1):.4f} | '
f'Acc: {100.*correct/total:.2f}%')
epoch_loss = running_loss / len(train_loader)
epoch_acc = 100. * correct / total
return epoch_loss, epoch_acc
训练函数遵循标准的PyTorch训练范式:设置模型为训练模式(model.train(),启用Dropout和BatchNorm的统计更新)、遍历DataLoader、将数据移至GPU、梯度清零、前向传播、计算损失、反向传播、参数更新。同时统计累计损失和准确率,每10个batch打印一次进度信息,便于监控训练状态。
验证函数:
def validate(model, val_loader, criterion, device):
model.eval()
running_loss = 0.0
correct = 0
total = 0
with torch.no_grad():
for inputs, targets in val_loader:
inputs, targets = inputs.to(device), targets.to(device)
outputs = model(inputs)
loss = criterion(outputs, targets)
running_loss += loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
epoch_loss = running_loss / len(val_loader)
epoch_acc = 100. * correct / total
return epoch_loss, epoch_acc
验证时设置模型为评估模式(model.eval()),并使用 torch.no_grad() 上下文管理器禁用梯度计算,可以显著减少显存占用并加速推理。
主训练循环:
def main():
# 加载配置
with open(config_path, 'r', encoding='utf-8') as f:
config = yaml.safe_load(f)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 创建数据加载器
train_loader, val_loader, class_to_idx = create_dataloaders(
config_path=str(config_path),
batch_size=config['training']['classifier']['batch_size'],
num_workers=0 # Windows下设置为0避免多进程问题
)
# 创建模型
num_classes = len(class_to_idx)
model = create_classifier(num_classes=num_classes, pretrained=True, device=device)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=config['training']['classifier']['learning_rate'])
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer, mode='min', factor=0.5, patience=5, verbose=True
)
# 训练参数
epochs = config['training']['classifier']['epochs']
best_val_acc = 0.0
# 保存类别映射
import json
with open(class_mapping_path, 'w', encoding='utf-8') as f:
json.dump(class_to_idx, f, ensure_ascii=False, indent=2)
for epoch in range(epochs):
train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device)
val_loss, val_acc = validate(model, val_loader, criterion, device)
scheduler.step(val_loss)
if val_acc > best_val_acc:
best_val_acc = val_acc
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'val_acc': val_acc,
'class_to_idx': class_to_idx,
}, best_model_path)
plot_training_curves(train_losses, train_accs, val_losses, val_accs, plot_path)
训练配置的关键参数如下:
- 优化器: Adam,学习率0.001。Adam结合了Momentum和RMSProp的优点,对学习率的缩放不敏感,通常作为首选优化器。
- 学习率调度:
ReduceLROnPlateau,当验证损失连续5个epoch不下降时,学习率减半(factor=0.5)。这种自适应学习率调整策略可以帮助模型在训练后期更精细地收敛。 - 损失函数: 交叉熵损失(CrossEntropyLoss),这是多分类任务的标准选择。
- Epoch数: 30轮。由于使用了预训练权重,模型不需要从零开始训练,30轮通常已经足够收敛。
- Batch Size: 32。在8GB显存的GPU上可以稳定运行。
- Num Workers: 0。在Windows环境下,PyTorch的多进程DataLoader有时会与CUDA初始化产生冲突,设置为0可以避免这类问题,代价是数据加载略慢。
训练曲线绘制:
def plot_training_curves(train_losses, train_accs, val_losses, val_accs, save_path):
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
ax1.plot(train_losses, label='Train Loss', marker='o')
ax1.plot(val_losses, label='Val Loss', marker='s')
ax1.set_xlabel('Epoch')
ax1.set_ylabel('Loss')
ax1.set_title('Training and Validation Loss')
ax1.legend()
ax1.grid(True)
ax2.plot(train_accs, label='Train Acc', marker='o')
ax2.plot(val_accs, label='Val Acc', marker='s')
ax2.set_xlabel('Epoch')
ax2.set_ylabel('Accuracy (%)')
ax2.set_title('Training and Validation Accuracy')
ax2.legend()
ax2.grid(True)
plt.tight_layout()
plt.savefig(save_path, dpi=150)
训练完成后,会在 output/ 目录下保存最佳模型权重 classifier_best.pth、训练曲线图 training_curves.png 以及类别映射文件 class_to_idx.json。在我们的实验中,验证集最佳准确率达到了 94.39%,为后续的Grad-CAM伪标注生成奠定了坚实基础。
五、弱监督标注——Grad-CAM伪边界框生成
5.1 为什么需要弱监督学习
目标检测是计算机视觉中最重要、最具挑战性的任务之一。与图像分类只需判断"图中有什么"不同,目标检测还需要回答"目标在哪里",即输出目标的精确位置(边界框坐标)和类别。这种精细的位置标注通常需要大量的人工劳动:标注员需要仔细查看每张图片,用矩形框精确框出每个目标物体的位置,并标注其类别。
在工业瑕疵检测场景中,标注困难尤为突出:
- 工业图像分辨率通常很高(2000x2000甚至更高),瑕疵区域可能只占图像的很小一部分,定位和框选需要极高的专注度和专业知识。
- 瑕疵的边界往往是模糊、不规则的,不同标注员对"框应该多大"的理解可能存在差异,导致标注一致性差。
- 数据量通常较大,数百甚至数千张图像的逐张标注需要数周甚至数月的时间成本。
弱监督学习(Weakly Supervised Learning)正是在这种背景下应运而生的一种解决方案。弱监督学习的核心思想是:利用容易获取的、粗粒度的标注信息(如图像级标签、点标注、涂鸦标注等),通过算法自动推断或生成细粒度的标注信息(如边界框、像素级掩码)。本项目采用的 Grad-CAM + 热力图后处理 方案,就是一种典型的弱监督标注方法。
5.2 Grad-CAM原理深入讲解
Grad-CAM(Gradient-weighted Class Activation Mapping) 由Selvaraju等人在2017年提出,是一种用于可视化卷积神经网络决策依据的技术。Grad-CAM能够生成一张"类别判别热力图",高亮显示图像中对某个特定类别预测贡献最大的区域。
Grad-CAM的核心数学原理如下:
假设我们要解释分类模型对类别 c c c 的预测。设模型最后一个卷积层输出的特征图为 A ∈ R K × H × W A \in \mathbb{R}^{K \times H \times W} A∈RK×H×W,其中 K K K 是通道数, H H H 和 W W W 是空间尺寸。对于ResNet18,layer4 的输出就是 A A A,其尺寸为 [ 512 , 7 , 7 ] [512, 7, 7] [512,7,7]。
Grad-CAM首先计算类别 c c c 的预测分数 y c y^c yc 对特征图 A A A 每个通道每个空间位置的梯度:
∂ y c ∂ A i j k \frac{\partial y^c}{\partial A_{ij}^k} ∂Aijk∂yc
然后,对每个通道 k k k,计算梯度在空间维度上的全局平均,得到该通道的重要性权重 α k c \alpha_k^c αkc:
α k c = 1 Z ∑ i = 1 H ∑ j = 1 W ∂ y c ∂ A i j k \alpha_k^c = \frac{1}{Z} \sum_{i=1}^{H} \sum_{j=1}^{W} \frac{\partial y^c}{\partial A_{ij}^k} αkc=Z1i=1∑Hj=1∑W∂Aijk∂yc
其中 Z = H × W Z = H \times W Z=H×W 是空间位置总数。这个权重 α k c \alpha_k^c αkc 量化了第 k k k 个特征图通道对于类别 c c c 的重要性。
最后,将所有通道的特征图按照权重进行线性组合,并通过ReLU激活函数得到Grad-CAM热力图:
L G r a d − C A M c = ReLU ( ∑ k α k c A k ) L_{Grad-CAM}^c = \text{ReLU}\left(\sum_{k} \alpha_k^c A^k\right) LGrad−CAMc=ReLU(k∑αkcAk)
ReLU的作用是只保留对类别 c c c 有正向贡献的特征,抑制负向贡献(即可能属于背景或其他类别的区域)。
直观理解,Grad-CAM回答了这样一个问题:模型在判断这张图属于"擦花"时,它的"注意力"集中在图像的哪些区域? 这些注意力高度集中的区域,极有可能就是"擦花"瑕疵实际出现的位置。
5.3 Grad-CAM热力图生成实现
scripts/generate_pseudo_labels.py 中定义了 GradCAMExtractor 类,封装了Grad-CAM热力图的生成逻辑:
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
class GradCAMExtractor:
def __init__(self, model, target_layer_name='backbone.layer4'):
self.model = model
self.model.eval()
# 获取目标层
target_layer = dict(model.named_modules())[target_layer_name]
# 创建Grad-CAM对象
self.grad_cam = GradCAM(
model=model,
target_layers=[target_layer]
)
def generate_cam(self, image_tensor, class_idx):
targets = [ClassifierOutputTarget(class_idx)]
cam = self.grad_cam(input_tensor=image_tensor, targets=targets)
# 归一化到[0, 1]
cam = cam.squeeze()
cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
return cam
代码解析:
-
目标层选择:
target_layer_name='backbone.layer4'指定了ResNet18的最后一个残差块作为Grad-CAM的目标层。这一层包含最深层的语义特征,同时仍保留了一定的空间分辨率(7x7),是生成定位热力图的最佳选择。如果选择更浅的层(如layer3),空间分辨率更高但语义信息不足;如果选择全局池化后的层,则完全失去了空间信息。 -
ClassifierOutputTarget:
pytorch-grad-cam库提供了ClassifierOutputTarget类,用于指定我们关注的类别索引。在传入grad_cam()函数时,库会自动计算该类别输出对目标层的梯度。 -
归一化: 生成的原始CAM值范围不确定,通过
(cam - min) / (max - min + 1e-8)将其线性映射到[0, 1]区间,便于后续统一阈值处理。1e-8是为了防止除零。
5.4 从热力图到边界框
Grad-CAM生成的是一张连续值的热力图,我们需要将其转换为离散的边界框标注。这一过程分为三个步骤:二值化、轮廓检测、YOLO格式转换。
第一步:热力图二值化
def cam_to_bbox(cam, threshold=0.5, min_area_ratio=0.01):
binary = (cam > threshold).astype(np.uint8)
contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if len(contours) == 0:
return None
largest_contour = max(contours, key=cv2.contourArea)
area = cv2.contourArea(largest_contour)
h, w = cam.shape
if area < min_area_ratio * h * w:
return None
x, y, w_box, h_box = cv2.boundingRect(largest_contour)
return [x, y, x + w_box, y + h_box]
二值化使用固定阈值 threshold=0.5,即只保留热力值高于0.5的区域作为候选瑕疵区域。然后使用OpenCV的 findContours 函数寻找连通域。由于热力图中可能存在多个分散的响应区域,我们取面积最大的连通域作为瑕疵主体。
min_area_ratio=0.01 是一个过滤条件,要求瑕疵区域的面积至少占整图的1%。这个条件可以排除极小的噪声响应,避免生成过小的无意义边界框。
第二步:最小外接矩形
对最大的连通域调用 cv2.boundingRect,得到包围该连通域的最小轴对齐矩形。返回的 x, y, w, h 转换为 [x_min, y_min, x_max, y_max] 格式。
第三步:YOLO格式转换
YOLO系列模型要求的标注格式为:每个目标一行,包含 class_id x_center y_center width height,其中四个坐标值均为相对于图像尺寸的归一化值(0~1之间)。
def bbox_to_yolo_format(bbox, image_width, image_height, class_id):
x_min, y_min, x_max, y_max = bbox
x_center = (x_min + x_max) / 2.0 / image_width
y_center = (y_min + y_max) / 2.0 / image_height
width = (x_max - x_min) / image_width
height = (y_max - y_min) / image_height
# 裁剪到[0, 1]范围
x_center = max(0, min(1, x_center))
y_center = max(0, min(1, y_center))
width = max(0, min(1, width))
height = max(0, min(1, height))
return [class_id, x_center, y_center, width, height]
边界情况处理(裁剪到 [0, 1])非常重要。由于Grad-CAM生成的边界框可能存在微小的越界(如x_max略大于图像宽度),如果不做裁剪,YOLO训练时会报错。
5.5 伪标注质量评估与阈值选择
Grad-CAM的 threshold 参数直接影响伪标注的质量。阈值太高会导致边界框过小甚至漏检;阈值太低会引入大量背景噪声,边界框过于松散。
在我们的实验中,尝试了0.3、0.5、0.7等多个阈值。最终选择 0.5 作为默认阈值,基于以下观察:
- threshold=0.3时,边界框普遍偏大,背景区域被大量包含,后续YOLO训练的mAP较低(约85%)。
- threshold=0.5时,边界框能够较好地覆盖瑕疵主体,背景噪声较少,YOLO训练的mAP达到最优(92%+)。
- threshold=0.7时,部分瑕疵的边界框过小甚至消失( especially for low-contrast defects like “桔皮”),导致漏检率上升。
对于少数Grad-CAM完全无法生成有效边界框的情况(如热力图全区域响应都很弱),我们采用 fallback 策略:将整张图像作为边界框。虽然这种标注质量很差,但总比没有标注要好,而且YOLO对少量噪声标注具有一定鲁棒性。
5.6 批量伪标注生成
generate_pseudo_labels.py 的主函数遍历训练集中的所有瑕疵图片,逐张生成伪标注:
def main():
# 加载配置和模型
model = DefectClassifier(num_classes=num_classes)
checkpoint = torch.load(model_path, map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
grad_cam_extractor = GradCAMExtractor(model)
# 遍历所有瑕疵类别
for defect_dir in defect_dirs:
class_name = defect_dir.name
class_id = class_to_idx[class_name]
for img_path in defect_dir.glob('*.jpg'):
image = Image.open(img_path).convert('RGB')
orig_width, orig_height = image.size
# 预处理并生成CAM
image_tensor = transform(image).unsqueeze(0)
cam = grad_cam_extractor.generate_cam(image_tensor, class_id)
# 调整CAM尺寸到原图大小
cam_resized = cv2.resize(cam, (orig_width, orig_height))
# 转换为边界框
bbox = cam_to_bbox(cam_resized, threshold=0.5, min_area_ratio=0.01)
if bbox is None:
bbox = [0, 0, orig_width, orig_height]
yolo_bbox = bbox_to_yolo_format(bbox, orig_width, orig_height, class_id)
# 保存标注
annotation_path = annotations_dir / f"{img_path.stem}.txt"
save_yolo_annotation(yolo_bbox, annotation_path)
运行完成后,dataset/annotations/ 目录下会生成与图片同名的 .txt 标注文件。这些标注文件随后会被整理到YOLO格式的目录结构中,供YOLOv8训练使用。
值得注意的是,Grad-CAM伪标注的精度虽然不如人工标注,但对于后续的YOLOv8训练已经足够。 YOLO系列算法本身就对标注噪声具有较好的鲁棒性,其损失函数设计(如CIoU Loss)能够容忍边界框的一定偏移。在我们的实验中,基于伪标注训练的YOLOv8n达到了92.55%的mAP,充分证明了弱监督标注方案的可行性和有效性。
六、YOLOv8目标检测模型训练
6.1 YOLOv8架构简介
YOLO(You Only Look Once)系列是目标检测领域最具影响力的算法家族之一。与传统两阶段检测器(如Faster R-CNN)不同,YOLO采用单阶段(One-Stage)设计,将目标检测任务建模为一个端到端的回归问题:直接从输入图像预测边界框坐标和类别概率,无需区域提议(Region Proposal)步骤。这种设计使得YOLO具有极高的推理速度,能够满足实时检测的需求。
YOLOv8 由Ultralytics公司于2023年发布,是YOLO系列的最新版本之一。相比前代YOLOv5,YOLOv8在架构设计和训练策略上进行了多项重要改进:
- 骨干网络优化: 采用C2f模块(CSPLayer with 2 Convolutions and Fast fusion)替代原有的C3模块,在保持相近计算量的同时增强了梯度流和特征融合能力。
- 解耦头(Decoupled Head): 将分类和回归任务分离到不同的分支,避免了两个任务之间的特征耦合,提升了检测精度。
- Anchor-Free设计: 不再依赖预定义的锚框(Anchor),而是直接预测边界框的中心点坐标和宽高,简化了模型结构并减少了对先验知识的依赖。
- Task-Aligned Assigner: 改进了正负样本分配策略,使得训练过程中的样本分配更加合理,尤其有利于小目标的检测。
YOLOv8提供了多个不同规模的模型变体(n、s、m、l、x),参数量和计算量依次递增,检测精度也随之提升。本项目中我们选择 YOLOv8n(nano),这是最小、最轻量的版本,其优势在于:
- 速度快: 单张图片推理时间仅数毫秒,轻松满足实时检测需求。
- 显存占用小: 训练时batch size可以设得更大,或者在更低配的GPU上运行。
- 易于部署: 模型体积小,便于集成到边缘设备或嵌入式系统中。
对于毕业设计或原型验证来说,YOLOv8n的精度(在我们的任务上mAP达到92%+)已经完全足够,同时其轻量特性也使得训练和演示更加便捷。
6.2 训练脚本详解
scripts/train_detector.py 实现了YOLOv8的训练流程。由于YOLOv8的Ultralytics库封装程度很高,训练代码本身非常简洁,但我们在脚本中增加了一些重要的工程化处理。
import os
import sys
import io
import logging
from pathlib import Path
# 将Ultralytics配置重定向到项目目录(解决沙箱权限问题)
_ultralytics_dir = Path('output/ultralytics_config')
_ultralytics_dir.mkdir(parents=True, exist_ok=True)
os.environ['YOLO_CONFIG_DIR'] = str(_ultralytics_dir)
os.environ['ULTRALYTICS_SETTINGS_DIR'] = str(_ultralytics_dir)
# 禁用Ultralytics的rich进度条(解决特殊Unicode字符乱码)
os.environ['YOLO_VERBOSE'] = 'False'
os.environ['ULTRALYTICS_PROGRESS_BAR'] = 'False'
# 修复Windows控制台中文乱码问题
log_file = Path('output/train_log.txt')
log_file.parent.mkdir(parents=True, exist_ok=True)
logging.basicConfig(
level=logging.INFO,
format='%(message)s',
handlers=[
logging.FileHandler(log_file, mode='a', encoding='utf-8'),
]
)
logger = logging.getLogger(__name__)
# 重定向 print 中的中文字符到日志文件
_original_print = print
def print(*args, **kwargs):
msg = ' '.join(str(a) for a in args)
if any(ord(c) > 127 for c in msg):
logger.info(msg)
try:
_original_print(*args, **kwargs)
except UnicodeEncodeError:
pass
from ultralytics import YOLO
工程化处理解析:
-
Ultralytics配置目录重定向: Ultralytics库默认会在用户主目录下创建配置文件夹。在某些受限环境(如沙箱、容器)中,这可能导致权限错误。通过设置
YOLO_CONFIG_DIR和ULTRALYTICS_SETTINGS_DIR环境变量,我们将配置缓存重定向到项目内部的output/ultralytics_config/目录,避免了权限问题。 -
禁用Rich进度条: Ultralytics默认使用
rich库绘制精美的进度条和表格,但这些输出中包含大量Unicode特殊字符(如方框、线条符号)。在Windows默认编码(GBK)的控制台中,这些字符会导致严重的乱码甚至程序崩溃。我们通过设置YOLO_VERBOSE='False'和ULTRALYTICS_PROGRESS_BAR='False'完全禁用了这些输出。 -
中文乱码修复: Windows控制台的编码问题是一个经典难题。我们的解决方案是双重输出:包含中文字符的日志写入UTF-8编码的日志文件
output/train_log.txt,纯英文信息输出到控制台。这样既能保留完整的中文训练记录,又不会在控制台出现乱码。
主训练函数:
def main():
dataset_yaml = Path('dataset/yolo/dataset.yaml')
if not dataset_yaml.exists():
print(f"错误: 数据集配置文件不存在: {dataset_yaml}")
print("请先运行 prepare_yolo_dataset.py")
sys.exit(1)
# 加载预训练YOLOv8n模型
model = YOLO('yolov8n.pt')
# 开始训练
results = model.train(
data=str(dataset_yaml),
epochs=100,
batch=16,
imgsz=640,
patience=20,
device=0,
project='output',
name='yolov8_train',
exist_ok=True,
pretrained=True,
optimizer='Adam',
lr0=0.001,
lrf=0.01,
momentum=0.9,
weight_decay=0.0005,
warmup_epochs=3,
warmup_momentum=0.8,
warmup_bias_lr=0.1,
seed=42,
workers=0,
verbose=True
)
# 验证模型
metrics = model.val()
print(f"mAP@0.5: {metrics.box.map50:.4f}")
print(f"mAP@0.5:0.95: {metrics.box.map:.4f}")
print(f"Precision: {metrics.box.mp:.4f}")
print(f"Recall: {metrics.box.mr:.4f}")
训练参数详解:
| 参数 | 值 | 说明 |
|---|---|---|
epochs |
100 | 最大训练轮数 |
batch |
16 | 每批样本数 |
imgsz |
640 | 输入图像尺寸(640x640) |
patience |
20 | 早停耐心值:验证mAP连续20轮不提升则停止训练 |
device |
0 | 使用第0块GPU |
optimizer |
Adam | 优化器类型 |
lr0 |
0.001 | 初始学习率 |
lrf |
0.01 | 最终学习率 = lr0 * lrf = 0.00001 |
momentum |
0.9 | SGD动量(Adam中用于梯度一阶矩估计) |
weight_decay |
0.0005 | L2正则化系数,防止过拟合 |
warmup_epochs |
3 | 学习率预热轮数,防止训练初期震荡 |
seed |
42 | 随机种子,保证实验可复现 |
workers |
0 | 数据加载线程数(Windows下设为0) |
早停机制(Early Stopping): patience=20 表示如果验证集的mAP连续20个epoch没有提升,训练将自动停止并恢复最佳模型权重。在我们的实验中,训练在大约73轮时因早停触发而结束(手动观察CSV记录约46轮数据),此时模型已经充分收敛,继续训练不会带来明显的精度提升,反而可能因过拟合导致泛化性能下降。
学习率预热(Warmup): 在训练的前3个epoch,学习率从一个很小的值线性增加到设定的初始学习率0.001。预热策略可以有效防止训练初期的梯度爆炸和损失震荡,尤其在使用较大batch size时效果显著。
6.3 YOLO格式数据集准备
在运行训练脚本之前,需要先将图片和伪标注文件整理成YOLO格式。scripts/prepare_yolo_dataset.py 完成了这一工作:
def prepare_yolo_dataset():
# 创建YOLO目录结构
yolo_dir = Path('dataset/yolo')
(yolo_dir / 'images/train').mkdir(parents=True, exist_ok=True)
(yolo_dir / 'images/val').mkdir(parents=True, exist_ok=True)
(yolo_dir / 'labels/train').mkdir(parents=True, exist_ok=True)
(yolo_dir / 'labels/val').mkdir(parents=True, exist_ok=True)
# 根据train.txt/val.txt复制图片和标注
for rel_path in train_files:
img_path = Path('dataset/guangdong_round1_train2_20180916') / rel_path
label_path = Path('dataset/annotations') / (img_path.stem + '.txt')
if label_path.exists():
shutil.copy(img_path, yolo_dir / 'images/train' / img_path.name)
shutil.copy(label_path, yolo_dir / 'labels/train' / label_path.name)
# 创建dataset.yaml
dataset_config = {
'path': str(yolo_dir.absolute()),
'train': 'images/train',
'val': 'images/val',
'nc': len(class_to_idx),
'names': list(class_to_idx.keys())
}
with open(yolo_dir / 'dataset.yaml', 'w', encoding='utf-8') as f:
yaml.dump(dataset_config, f, allow_unicode=True, default_flow_style=False)
YOLO格式要求目录结构如下:
dataset/yolo/
├── dataset.yaml # 数据集配置文件
├── images/
│ ├── train/ # 训练集图片
│ └── val/ # 验证集图片
└── labels/
├── train/ # 训练集标注(.txt)
└── val/ # 验证集标注(.txt)
dataset.yaml 的内容示例:
path: D:/aluminum-defect-detection/dataset/yolo
train: images/train
val: images/val
nc: 11
names:
- 不导电
- 其他
- 凸粉
- 擦花
- 桔皮
- 横条压凹
- 涂层开裂
- 漏底
- 碰伤
- 脏点
- 起坑
6.4 训练结果与性能分析
训练完成后,output/yolov8_train/ 目录下会生成丰富的训练产物:
| 文件/目录 | 说明 |
|---|---|
weights/best.pt |
验证mAP最高的模型权重 |
weights/last.pt |
最后一个epoch的模型权重 |
results.csv |
每个epoch的训练/验证指标记录 |
results.png |
训练曲线汇总图 |
confusion_matrix.png |
混淆矩阵 |
PR_curve.png |
精确率-召回率曲线 |
F1_curve.png |
F1分数曲线 |
labels.jpg |
数据集聚类可视化 |
train_batch*.jpg |
训练批次增强效果预览 |
在我们的实验中,模型在验证集上取得了以下性能:
| 指标 | 数值 | 说明 |
|---|---|---|
| mAP@0.5 | 0.9255 | IoU阈值0.5时的平均精度,衡量检测框定位的准确性 |
| mAP@0.5:0.95 | 0.6910 | COCO标准mAP,在IoU 0.5~0.95范围内取平均,更严格 |
| Precision | 0.9090 | 精确率:预测为正的样本中实际为正的比例 |
| Recall | 0.8933 | 召回率:实际为正的样本中被正确检出的比例 |
mAP@0.5 = 92.55% 的结果在工业瑕疵检测领域属于优秀水平。考虑到我们的训练数据是基于Grad-CAM自动生成的伪标注,而非昂贵的人工精确标注,这一结果尤其令人鼓舞。它充分证明了弱监督学习在降低标注成本的同时,仍能保持较高的检测精度。
从 results.csv 中可以观察到训练过程的动态变化:box_loss(边界框回归损失)和cls_loss(分类损失)在前20个epoch快速下降,随后趋于平稳;mAP在30个epoch后进入平台期,在70个epoch附近达到峰值。早停机制在合适的时机终止了训练,避免了过拟合。
混淆矩阵(confusion_matrix.png)显示,模型在"擦花"、“漏底”、“脏点"等样本较多的类别上表现优异,而在"涂层开裂”、"其他"等样本稀缺的类别上存在一定程度的混淆。这与数据集的类别不平衡分布一致,提示我们可以通过增加稀缺类别的样本或调整类别权重来进一步提升整体性能。
七、PyQt5图形界面开发——打造工业级可视化系统
7.1 GUI整体架构设计
一个优秀的深度学习项目,不仅需要强大的算法内核,还需要友好的人机交互界面。尤其在工业质检和学术答辩场景中,图形用户界面(GUI)能够将抽象的模型推理过程转化为直观的可视化结果,大幅降低系统的使用门槛。
本项目采用 PyQt5 作为GUI开发框架。PyQt5是Qt框架的Python绑定,具有以下优势:
- 跨平台: 同一套代码可在Windows、Linux、macOS上运行。
- 功能丰富: 提供了从基础控件到高级图形视图、多媒体、网络等完整的功能模块。
- 信号与槽机制: Qt独特的信号-槽(Signal-Slot)机制实现了对象间的松耦合通信,非常适合多线程GUI开发。
- 界面美观: 支持QSS(Qt StyleSheets)样式表,可以像CSS一样美化界面。
- 与Python生态兼容: 可以无缝集成OpenCV、PIL、Matplotlib等Python库。
GUI的整体架构采用 “主窗口 + 多标签页” 的设计模式:
QMainWindow(主窗口)
├── QTabWidget(标签页容器)
│ ├── Tab 1: DetectorTab(瑕疵检测)
│ ├── Tab 2: TrainingResultsTab(训练结果)
│ ├── Tab 3: DatasetTab(数据集浏览)
│ └── Tab 4: TrainingControlTab(训练控制)
├── QMenuBar(菜单栏)
└── QStatusBar(状态栏)
每个标签页都是一个独立的 QWidget 子类,负责特定的功能模块。这种模块化设计使得代码结构清晰,便于独立开发和维护。
7.2 GUI启动入口与环境初始化
scripts/run_gui.py 是GUI的启动入口。在创建 QApplication 之前,脚本完成了一系列关键的环境初始化工作:
import os
import sys
import io
from pathlib import Path
PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
# 解决 Windows 控制台中文乱码
if sys.platform == 'win32':
os.system('chcp 65001 > nul 2>&1')
try:
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace', line_buffering=True)
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace', line_buffering=True)
except Exception:
pass
os.environ['PYTHONIOENCODING'] = 'utf-8'
os.environ['PYTHONUTF8'] = '1'
# 配置 Ultralytics 目录
ultralytics_dir = PROJECT_ROOT / 'output' / 'ultralytics_config'
ultralytics_dir.mkdir(parents=True, exist_ok=True)
os.environ['YOLO_CONFIG_DIR'] = str(ultralytics_dir)
os.environ['ULTRALYTICS_SETTINGS_DIR'] = str(ultralytics_dir)
os.environ['YOLO_VERBOSE'] = 'False'
# 切换工作目录到项目根
os.chdir(PROJECT_ROOT)
# PyQt5 DPI 缩放
os.environ['QT_AUTO_SCREEN_SCALE_FACTOR'] = '1'
from PyQt5.QtWidgets import QApplication
from PyQt5.QtCore import Qt
QApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True)
QApplication.setAttribute(Qt.AA_UseHighDpiPixmaps, True)
环境初始化解析:
-
控制台编码修复: Windows默认使用GBK编码,而Python源码和日志文件通常使用UTF-8。脚本通过
chcp 65001切换代码页到UTF-8,并用TextIOWrapper重定向标准输出流的编码。这是解决Windows下中文乱码最彻底的方式之一。 -
工作目录切换:
os.chdir(PROJECT_ROOT)确保所有相对路径(如dataset/yolo/dataset.yaml、output/yolov8_train)都能正确解析,无论用户从哪个目录启动脚本。 -
高DPI支持: 在高分辨率屏幕(如4K显示器)上,Qt默认不会自动缩放,导致界面元素过小。通过设置
QT_AUTO_SCREEN_SCALE_FACTOR和AA_EnableHighDpiScaling,Qt会根据屏幕DPI自动调整界面缩放比例,保证在高分辨率屏幕上的可用性。
7.3 主窗口 MainWindow
src/gui/main_window.py 定义了应用程序的主窗口,继承自 QMainWindow:
class MainWindow(QMainWindow):
def __init__(self, model_path='output/yolov8_train/weights/best.pt'):
super().__init__()
self.model_path = model_path
self.setWindowTitle("铝合金瑕疵识别系统 v1.0")
self.resize(1400, 900)
self._build_menu()
self._build_central()
self._build_status_bar()
主窗口的默认尺寸设为1400x900,这个尺寸在1080p显示器上能够完整显示所有内容,在更高分辨率屏幕上借助DPI缩放也能保持良好的可用性。
菜单栏设计:
def _build_menu(self):
menubar = self.menuBar()
file_menu = menubar.addMenu("文件(&F)")
exit_action = QAction("退出(&Q)", self)
exit_action.setShortcut(QKeySequence("Ctrl+Q"))
exit_action.triggered.connect(self.close)
file_menu.addAction(exit_action)
view_menu = menubar.addMenu("视图(&V)")
refresh_action = QAction("刷新训练结果(&R)", self)
refresh_action.setShortcut(QKeySequence("F5"))
refresh_action.triggered.connect(self._on_refresh)
view_menu.addAction(refresh_action)
help_menu = menubar.addMenu("帮助(&H)")
about_action = QAction("关于(&A)", self)
about_action.triggered.connect(self._on_about)
help_menu.addAction(about_action)
菜单栏遵循标准的桌面应用设计规范,包含文件、视图、帮助三个菜单。快捷键设置(如 Ctrl+Q 退出、F5 刷新)提升了操作效率。&F 这样的带下划线字符定义了Alt快捷键,按住Alt键后按对应字母即可打开菜单,这是Windows桌面应用的常见交互模式。
关于对话框:
def _on_about(self):
QMessageBox.about(
self, "关于",
"<h2>铝合金瑕疵识别系统 v1.0</h2>"
"<p>基于 YOLOv8 + PyQt5 的工业质检系统</p>"
"<p>支持功能:</p>"
"<ul>"
"<li>🔍 单张/批量图片瑕疵检测与可视化</li>"
"<li>📈 训练曲线、混淆矩阵、PR曲线查看</li>"
"<li>📂 数据集浏览与统计</li>"
"<li>⚙ 模型训练控制与参数调节</li>"
"</ul>"
"<p>模型性能: mAP@0.5 = 92.55%</p>"
)
关于对话框使用HTML格式文本,支持粗体、列表等富文本样式,使信息展示更加美观专业。
状态栏设计:
状态栏左侧显示当前加载的模型文件名,右侧显示GPU信息(通过 torch.cuda.get_device_name 动态获取)。这让用户在打开系统时就能一目了然地了解当前运行环境。
关闭事件处理:
def closeEvent(self, event):
try:
if hasattr(self, 'detector_tab') and self.detector_tab:
self.detector_tab.shutdown()
if hasattr(self, 'control_tab') and self.control_tab:
self.control_tab.shutdown()
except Exception as e:
print(f"清理资源时出错: {e}")
event.accept()
重写 closeEvent 方法,在窗口关闭时显式停止后台线程(检测线程和训练进程),避免程序退出后仍有子进程在后台运行,造成资源泄漏或僵尸进程。
7.4 瑕疵检测标签页 DetectorTab
DetectorTab 是系统的核心功能页面,负责加载图片、执行YOLOv8推理、显示检测结果和保存结果。其UI采用 左右分栏 布局,左侧为控制面板,右侧为图像显示区。
UI布局代码:
def _build_ui(self):
root = QHBoxLayout(self)
splitter = QSplitter(Qt.Horizontal)
root.addWidget(splitter)
# === 左侧:控制面板 ===
left = QWidget()
left_layout = QVBoxLayout(left)
# 检测参数组
params_group = QGroupBox("检测参数")
params_form = QFormLayout(params_group)
self.conf_spin = QDoubleSpinBox()
self.conf_spin.setRange(0.05, 1.0)
self.conf_spin.setSingleStep(0.05)
self.conf_spin.setValue(0.25)
params_form.addRow("置信度阈值:", self.conf_spin)
self.iou_spin = QDoubleSpinBox()
self.iou_spin.setRange(0.05, 1.0)
self.iou_spin.setSingleStep(0.05)
self.iou_spin.setValue(0.45)
params_form.addRow("IoU 阈值:", self.iou_spin)
left_layout.addWidget(params_group)
# 操作按钮组
btn_group = QGroupBox("操作")
btn_layout = QVBoxLayout(btn_group)
self.btn_load_image = QPushButton("加载单张图片")
self.btn_load_image.clicked.connect(self.on_load_image)
btn_layout.addWidget(self.btn_load_image)
self.btn_load_folder = QPushButton("加载文件夹")
btn_layout.addWidget(self.btn_load_folder)
self.btn_detect = QPushButton("开始检测")
self.btn_detect.setEnabled(False)
btn_layout.addWidget(self.btn_detect)
self.btn_save = QPushButton("保存结果")
self.btn_save.setEnabled(False)
btn_layout.addWidget(self.btn_save)
left_layout.addWidget(btn_group)
# 进度条
self.progress = QProgressBar()
self.progress.setVisible(False)
left_layout.addWidget(self.progress)
# 批量结果列表
batch_group = QGroupBox("检测列表")
self.batch_list = QListWidget()
left_layout.addWidget(batch_group, stretch=1)
splitter.addWidget(left)
# === 右侧:显示区 ===
right = QWidget()
right_layout = QVBoxLayout(right)
info_group = QGroupBox("检测详情")
self.info_text = QTextEdit()
self.info_text.setReadOnly(True)
right_layout.addWidget(info_group)
self.image_label = QLabel("请加载图片")
self.image_label.setAlignment(Qt.AlignCenter)
self.image_label.setStyleSheet(
"QLabel { background-color: #2b2b2b; color: #cccccc; "
"border: 1px solid #555; min-height: 500px; }"
)
right_layout.addWidget(self.image_label, stretch=1)
splitter.addWidget(right)
splitter.setSizes([400, 1000])
布局设计要点:
-
QSplitter分割器: 使用
QSplitter而非固定布局,用户可以通过拖动分割线自由调整左右面板宽度,适应不同屏幕尺寸和个人偏好。 -
QGroupBox分组: 将相关控件放入
QGroupBox中,并添加标题(如"检测参数"、“操作”),使界面层次清晰,用户能够快速找到所需功能。 -
QFormLayout表单布局: 在"检测参数"组中使用
QFormLayout,自动将标签和输入框成对排列,标签右对齐、输入框左对齐,形成整齐的视觉对齐线。 -
状态管理: "开始检测"和"保存结果"按钮默认禁用(
setEnabled(False)),只有在加载了模型或完成了检测后才启用,防止用户进行无效操作。 -
暗色主题图像区: 图像显示Label设置了深色背景(
#2b2b2b)和灰色边框,这在未加载图片时提供了良好的视觉占位效果,同时与检测框的鲜艳颜色形成对比。
单张图片加载与检测:
def on_load_image(self):
path, _ = QFileDialog.getOpenFileName(
self, "选择图片", "",
"图片文件 (*.jpg *.jpeg *.png *.bmp *.tif *.tiff *.webp)"
)
if not path:
return
self._load_single_image(path)
def on_detect(self):
if not self.current_image_path:
QMessageBox.information(self, "提示", "请先加载图片")
return
self.worker.detect_single(self.current_image_path)
self.progress.setVisible(True)
self.progress.setRange(0, 0) # 不确定进度
self._set_buttons_enabled(detect=False)
QFileDialog.getOpenFileName 弹出标准的文件选择对话框,自动过滤支持的图片格式。检测时进度条设为不确定模式(setRange(0, 0)),显示循环动画,提示用户后台正在处理。
批量检测功能:
def on_load_folder(self):
folder = QFileDialog.getExistingDirectory(self, "选择图片文件夹")
if not folder:
return
folder_p = Path(folder)
image_files = sorted([
p for p in folder_p.iterdir()
if p.is_file() and p.suffix.lower() in self.SUPPORTED_EXTS
])
self.batch_results = []
self.batch_list.clear()
for p in image_files:
item = QListWidgetItem(f"⏳ {p.name}")
item.setData(Qt.UserRole, str(p))
self.batch_list.addItem(item)
self.worker.detect_batch([str(p) for p in image_files])
self.progress.setVisible(True)
self.progress.setRange(0, len(image_files))
self.progress.setValue(0)
批量检测是工业场景中非常实用的功能。用户选择一个包含多张图片的文件夹后,系统会为每张图片创建一个列表项,初始状态显示 ⏳ 图标。检测完成后,列表项更新为 ✓ 图标并显示检测到的瑕疵数量,方便用户快速浏览哪些图片存在问题。
检测结果显示:
def _display_result(self, image_path, results, annotated):
self.current_image_path = image_path
self.current_results = results
if annotated is None and results is not None:
annotated = draw_detections_on_image(
img, results, DEFECT_CLASSES,
conf_threshold=self.conf_spin.value()
)
self.current_annotated = annotated
if annotated is not None:
self._show_image(annotated)
self.info_text.clear()
self.info_text.append(f"图片: {Path(image_path).name}")
if results is not None and hasattr(results, 'boxes'):
boxes = results.boxes
n = len(boxes)
self.info_text.append(f"检测到 {n} 个瑕疵:")
for idx, box in enumerate(boxes, 1):
cls = int(box.cls[0].item())
conf = float(box.conf[0].item())
name = DEFECT_CLASSES[cls]
self.info_text.append(f" {idx}. {name} - 置信度: {conf:.3f}")
右侧的 info_text(只读QTextEdit)显示详细的检测结果:图片名称、瑕疵总数、每个瑕疵的类别名称和置信度。这些信息以结构化的文本形式呈现,便于用户快速了解检测结果的全貌。
结果保存:
保存功能不仅导出带检测框的图片,还同时生成JSON格式的结构化数据文件,包含每个检测框的坐标、类别和置信度,便于后续的数据分析和系统集成。
def on_save_result(self):
# 保存PNG图片
cv2.imwrite(path, self.current_annotated)
# 同时保存 JSON 标注
json_path = Path(path).with_suffix('.json')
detections = []
for box in self.current_results.boxes:
xyxy = box.xyxy[0].cpu().numpy().tolist()
cls = int(box.cls[0].item())
conf = float(box.conf[0].item())
detections.append({
"class_id": cls,
"class_name": DEFECT_CLASSES[cls],
"confidence": conf,
"bbox_xyxy": xyxy
})
with open(json_path, 'w', encoding='utf-8') as f:
json.dump({...}, f, ensure_ascii=False, indent=2)
7.5 后台推理线程 DetectorWorker
在GUI中进行深度学习推理时,绝对不能在主线程(UI线程)中执行耗时操作。如果直接在按钮点击的槽函数中调用 model.predict(),模型推理期间GUI会完全卡死,无法响应用户的任何操作,甚至会被操作系统判定为"未响应"而强制关闭。
解决方案是使用 QThread 工作线程。DetectorWorker 继承自 QThread,在独立的线程中加载YOLO模型并执行推理:
class DetectorWorker(QThread):
detection_finished = pyqtSignal(str, object, object)
batch_progress = pyqtSignal(int, int, str)
batch_finished = pyqtSignal(list)
error_occurred = pyqtSignal(str)
model_loaded = pyqtSignal(str)
def __init__(self, model_path='output/yolov8_train/weights/best.pt',
device='', conf=0.25, iou=0.45, parent=None):
super().__init__(parent)
self.model_path = model_path
self.device = device if device else None
self.conf = conf
self.iou = iou
self._stop_flag = False
self.model = None
self.task_queue = []
self.current_task = None
信号(Signal)设计:
DetectorWorker 定义了5个信号,用于向主线程汇报状态:
model_loaded:模型加载完成,通知主线程启用"开始检测"按钮。detection_finished:单张图片检测完成,传递图片路径和检测结果。batch_progress:批量检测进度更新,传递当前进度、总数和图片路径。batch_finished:批量检测全部完成,传递所有结果列表。error_occurred:发生错误,传递错误信息,主线程弹出警告框。
信号-槽机制是Qt多线程编程的精髓。工作线程只负责执行计算任务,通过发射信号将结果传递给主线程;主线程的槽函数接收信号后更新UI。这种设计保证了线程安全——只有主线程能够操作GUI控件。
任务队列模式:
def detect_single(self, image_path):
self.task_queue.append(('single', str(image_path)))
def detect_batch(self, image_paths):
self.task_queue.append(('batch', [str(p) for p in image_paths]))
def run(self):
if self.model is None:
self.load_model()
if self.model is None:
return
while not self._stop_flag:
if not self.task_queue:
self.msleep(100)
continue
self.current_task = self.task_queue.pop(0)
task_type, payload = self.current_task
try:
if task_type == 'single':
self._run_single(payload)
elif task_type == 'batch':
self._run_batch(payload)
except Exception as e:
self.error_occurred.emit(f"检测出错: {e}")
Worker采用 任务队列模式 而非直接执行。主线程通过 detect_single 或 detect_batch 将任务加入队列,Worker线程在 run() 方法中循环检查队列并依次执行。这种模式的好处是:
- 支持任务的异步提交,主线程无需等待。
- 避免了重复创建线程的开销,一个Worker线程可以处理多个检测请求。
- 便于实现取消功能:通过设置
_stop_flag可以优雅地停止后续任务。
模型加载策略:
def load_model(self):
try:
from ultralytics import YOLO
self.model = YOLO(self.model_path)
self.model_loaded.emit(f"模型已加载: {self.model_path}")
except Exception as e:
self.error_occurred.emit(f"模型加载失败: {e}")
self.model = None
模型加载被延迟到Worker线程启动后执行,而不是在 __init__ 中。这是因为YOLO模型的初始化过程(下载预训练权重、构建计算图)较为耗时,如果在主线程中初始化会导致GUI启动卡顿。延迟加载使得主窗口能够先快速显示出来,同时后台静默加载模型。
7.6 中文标签绘制与图像工具
目标检测的可视化通常需要在图像上绘制边界框和类别标签。OpenCV自带的 cv2.putText 函数不支持中文,如果直接用英文标签,在工业场景中会严重影响可读性。因此,我们采用 PIL(Pillow)绘制中文 + OpenCV图像处理 的混合方案。
src/gui/utils.py 中的 draw_detections_on_image 函数实现了带中文标签的检测框绘制:
def draw_detections_on_image(img_bgr, detections, class_names=None, conf_threshold=0.25):
if class_names is None:
class_names = DEFECT_CLASSES
h, w = img_bgr.shape[:2]
colors = get_class_colors(len(class_names))
# 转换 BGR -> RGB for PIL
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
pil_img = Image.fromarray(img_rgb)
draw = ImageDraw.Draw(pil_img)
font = get_pil_font(max(14, min(24, w // 50)))
det_list = _normalize_detections(detections, conf_threshold)
for det in det_list:
box = det['box']
cls_idx = int(det['class'])
conf = float(det['conf'])
cls_name = class_names[cls_idx]
bgr_color = colors[cls_idx % len(colors)]
rgb_color = (bgr_color[2], bgr_color[1], bgr_color[0])
# 绘制矩形框
draw.rectangle(box, outline=rgb_color, width=2)
# 绘制标签背景
label = f"{cls_name} {conf:.2f}"
bbox = draw.textbbox((0, 0), label, font=font)
text_w = bbox[2] - bbox[0]
text_h = bbox[3] - bbox[1]
tx1, ty1 = int(box[0]), int(box[1])
tx2, ty2 = tx1 + text_w + 6, ty1 + text_h + 4
draw.rectangle([tx1, ty1, tx2, ty2], fill=rgb_color)
draw.text((tx1 + 3, ty1 + 2), label, fill=(255, 255, 255), font=font)
# RGB -> BGR 转回OpenCV格式
return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
技术要点:
-
颜色空间转换: OpenCV默认使用BGR格式,而PIL使用RGB格式。绘制前
cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB),绘制后cv2.cvtColor(..., cv2.COLOR_RGB2BGR),确保颜色正确。 -
动态字体大小:
get_pil_font(max(14, min(24, w // 50)))根据图像宽度动态计算字体大小。对于高分辨率图像使用更大的字体,低分辨率图像使用较小的字体,保证标签的可读性。 -
自适应标签位置: 如果检测框位于图像顶部,标签可能会超出图像边界。代码通过判断
ty1 < text_h + 4检测这种情况,并将标签移至框内显示。 -
类别颜色映射:
get_class_colors使用HSV色彩空间均匀分布生成颜色,确保11个类别的颜色既鲜艳又容易区分。 -
检测结果标准化:
_normalize_detections函数统一处理Ultralytics Results对象和自定义字典列表两种输入格式,提高了函数的通用性。
字体加载:
def get_pil_font(size=16):
font_paths = [
"C:/Windows/Fonts/msyh.ttc",
"C:/Windows/Fonts/msyh.ttf",
"C:/Windows/Fonts/simhei.ttf",
"C:/Windows/Fonts/simsun.ttc",
]
for fp in font_paths:
if os.path.exists(fp):
try:
return ImageFont.truetype(fp, size)
except Exception:
continue
return ImageFont.load_default()
字体加载按优先级依次尝试微软雅黑、黑体、宋体,如果都失败则回退到默认字体。微软雅黑是Windows系统中最美观的中文字体之一,支持ClearType抗锯齿,在屏幕上显示效果清晰锐利。
图像格式转换:
def numpy_to_qpixmap(img_bgr):
from PyQt5.QtGui import QPixmap, QImage
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
h, w, ch = img_rgb.shape
bytes_per_line = ch * w
q_img = QImage(img_rgb.data, w, h, bytes_per_line, QImage.Format_RGB888).copy()
return QPixmap.fromImage(q_img)
numpy_to_qpixmap 实现了OpenCV图像到Qt图像的无缝转换。关键点是 .copy() 方法——QImage不会复制numpy数组的内存,如果原数组在QImage使用期间被释放或修改,会导致程序崩溃。.copy() 确保QImage拥有独立的数据副本,彻底避免悬垂指针问题。
7.7 训练结果可视化标签页
TrainingResultsTab 负责展示YOLO训练产生的各类图表和数据。页面采用 顶部信息卡 + 底部多子标签页 的布局:
信息卡设计:
info_group = QGroupBox("最佳模型信息")
info_layout = QHBoxLayout(info_group)
for metric in ['mAP@0.5', 'mAP@0.5:0.95', 'Precision', 'Recall', 'Epochs']:
box = QGroupBox(metric)
v = QVBoxLayout(box)
label = QLabel("--")
label.setAlignment(Qt.AlignCenter)
label.setStyleSheet("font-size: 22px; font-weight: bold; color: #4CAF50;")
v.addWidget(label)
self.info_labels[metric] = label
info_layout.addWidget(box)
顶部信息卡以醒目的绿色大字体(22px)显示关键指标,用户一打开标签页就能立即了解模型性能。每个指标放入独立的 QGroupBox 中,带有标题边框,视觉上形成卡片效果。
图表展示:
self.PLOT_FILES = [
('训练曲线', 'results.png'),
('混淆矩阵', 'confusion_matrix.png'),
('混淆矩阵(归一化)', 'confusion_matrix_normalized.png'),
('PR 曲线', 'PR_curve.png'),
('F1 曲线', 'F1_curve.png'),
('类别分布', 'labels.jpg'),
('训练批次 0', 'train_batch0.jpg'),
('训练批次 1', 'train_batch1.jpg'),
('训练批次 2', 'train_batch2.jpg'),
('验证批次预测', 'val_batch0_pred.jpg'),
]
系统预定义了10种图表文件,通过嵌套的 QTabWidget 为每种图表创建一个子标签页。图表加载逻辑如下:
def _load_plots(self):
for fname, lbl in self.plot_labels.items():
fpath = self.TRAIN_DIR / fname
if not fpath.exists():
lbl.setText(f"未找到 {fname}")
continue
pixmap = QPixmap(str(fpath))
scaled = pixmap.scaled(
lbl.size(),
Qt.KeepAspectRatio,
Qt.SmoothTransformation
)
lbl.setPixmap(scaled)
图片缩放使用 Qt.KeepAspectRatio 保持宽高比,Qt.SmoothTransformation 启用高质量双线性插值缩放,避免图像出现锯齿或模糊。
数据表格:
除了图表,标签页还提供了 results.csv 的表格视图。由于YOLOv8的results.csv在训练第一轮可能会记录重复行(train+val后立即记一次,epoch结束又记一次),代码实现了按epoch去重:
seen_epoch = set()
unique_data = []
for r in data:
if not r:
continue
ep = r[0]
if ep in seen_epoch:
continue
seen_epoch.add(ep)
unique_data.append(r)
表格中数值列格式化为4位小数,epoch列显示整数,通过 setTextAlignment(Qt.AlignCenter) 实现居中对齐,使数据呈现整齐美观。
7.8 数据集浏览标签页
DatasetTab 提供了对原始数据集的直观浏览功能。左侧是类别列表和统计信息,右侧是图片缩略图和大图预览。
缩略图浏览:
self.image_list = QListWidget()
self.image_list.setViewMode(QListWidget.IconMode)
self.image_list.setIconSize(self.THUMB_SIZE)
self.image_list.setResizeMode(QListWidget.Adjust)
self.image_list.setMovement(QListWidget.Static)
self.image_list.setSpacing(8)
通过设置 IconMode,QListWidget 以图标网格形式展示图片,类似于Windows资源管理器的"大图标"视图。setResizeMode(QListWidget.Adjust) 使列表在窗口大小变化时自动调整列数,setSpacing(8) 设置图标之间的间距,避免图片过于拥挤。
统计信息导出:
def _export_report(self):
report = {
"total_images": sum(len(v) for v in self.class_to_files.values()),
"train_total": sum(self.class_train_count.values()),
"val_total": sum(self.class_val_count.values()),
"num_classes": len(DEFECT_CLASSES),
"class_names": DEFECT_CLASSES,
"per_class": {}
}
for cls in DEFECT_CLASSES:
report["per_class"][cls] = {
"raw": len(self.class_to_files.get(cls, [])),
"train": self.class_train_count.get(cls, 0),
"val": self.class_val_count.get(cls, 0)
}
with open(path, 'w', encoding='utf-8') as f:
json.dump(report, f, ensure_ascii=False, indent=2)
数据集统计报告导出为JSON格式,包含总体统计和每个类别的详细数据。JSON格式便于后续用Python、Excel或其他工具进一步分析处理。
7.9 训练控制标签页
TrainingControlTab 允许用户通过图形界面启动新的YOLO训练,无需编写命令行。这是系统的高级功能,适合模型调优和增量训练场景。
参数表单设计:
页面提供了完整的训练参数配置表单,包括:
- 训练轮数(Epochs):1~1000,整数调节
- 批次大小(Batch):1~128
- 图像尺寸(imgsz):64~2048,步长32
- 学习率(lr0):0.0001~1.0,4位小数精度
- 优化器:Adam / SGD / AdamW / auto
- 预训练权重:文件浏览选择
- 数据集配置:YAML文件浏览选择
- 训练设备:自动检测可用GPU并列出
所有参数输入控件都设置了合理的取值范围和步长,防止用户输入无效参数。
QProcess训练启动:
def _on_start(self):
# 验证参数
weights = self.weights_edit.text().strip()
data = self.data_edit.text().strip()
if not weights or not Path(weights).exists():
QMessageBox.warning(self, "错误", f"预训练权重不存在: {weights}")
return
if not data or not Path(data).exists():
QMessageBox.warning(self, "错误", f"数据集配置不存在: {data}")
return
# 启动 QProcess
self.process = QProcess(self)
self.process.setProcessChannelMode(QProcess.MergedChannels)
self.process.readyReadStandardOutput.connect(self._on_process_output)
self.process.finished.connect(self._on_process_finished)
# 设置环境变量
env = self.process.processEnvironment()
env.insert('YOLO_CONFIG_DIR', str(ultralytics_dir))
env.insert('YOLO_VERBOSE', 'False')
env.insert('PYTHONIOENCODING', 'utf-8')
self.process.setProcessEnvironment(env)
# 构建内联训练脚本
train_script = f"""
import os
os.environ['YOLO_VERBOSE'] = 'False'
from ultralytics import YOLO
model = YOLO(r'{weights}')
results = model.train(
data=r'{data}',
epochs={epochs},
batch={batch},
imgsz={imgsz},
lr0={lr},
optimizer='{optimizer}',
device='{device}',
project='output',
name='yolov8_train',
exist_ok=True,
pretrained=True,
seed=42,
workers=0,
verbose=False
)
print('TRAIN_DONE')
"""
cmd = [sys.executable, '-c', train_script]
self.process.start(cmd[0], cmd[1:])
训练控制的核心是 QProcess,它在独立的子进程中执行Python训练脚本,不阻塞GUI主线程。我们使用 -c 参数传递内联Python代码,这样无需额外的临时文件。
日志输出通过 readyReadStandardOutput 信号实时读取并显示在 QTextEdit 中。日志文本框设置了暗色主题样式(黑色背景、浅灰文字、等宽字体),模仿专业IDE和控制台的视觉效果,长时间查看不易疲劳。
self.log_text.setStyleSheet(
"QTextEdit { background-color: #1e1e1e; color: #d4d4d4; "
"font-family: 'Consolas', 'Courier New', monospace; font-size: 11px; }"
)
同时通过正则表达式过滤ANSI颜色码,确保日志显示干净:
def _on_process_output(self):
data = self.process.readAllStandardOutput()
text = bytes(data).decode('utf-8', errors='replace')
import re
text = re.sub(r'\x1b\[[0-9;]*m', '', text)
self.log_text.append(text.rstrip())
停止训练: 提供"停止训练"按钮,通过 QProcess.kill() 强制终止子进程。虽然强制终止会导致当前epoch的训练结果丢失,但在需要紧急停止或发现参数设置错误时非常有用。
7.10 中文字体与界面美化
src/gui/utils.py 中的 set_chinese_font 函数为整个Qt应用设置中文字体:
def set_chinese_font(app=None):
from PyQt5.QtGui import QFont
from PyQt5.QtWidgets import QApplication
font_candidates = [
("Microsoft YaHei UI", "微软雅黑"),
("Microsoft YaHei", "微软雅黑"),
("SimHei", "黑体"),
("SimSun", "宋体"),
("Arial Unicode MS", "Arial Unicode"),
]
selected_family = None
if app is None:
app = QApplication.instance()
if app is not None:
font = QFont(font_candidates[0][0], 10)
app.setFont(font)
return selected_family
Qt应用默认使用系统字体,在中文Windows上通常是"微软雅黑"或"宋体"。显式设置字体可以确保在不同系统上界面文字的一致性,避免某些控件因继承父字体而显示异常。
八、模型评估与实验结果深入分析
8.1 评估指标体系
目标检测模型的性能评估比图像分类更为复杂,需要同时考虑定位精度和分类准确性。本项目采用以下指标:
IoU(Intersection over Union,交并比):
IoU是衡量两个边界框重叠程度的指标,计算公式为:
I o U = A r e a I n t e r s e c t i o n A r e a U n i o n = A ∩ B A ∪ B IoU = \frac{Area_{Intersection}}{Area_{Union}} = \frac{A \cap B}{A \cup B} IoU=AreaUnionAreaIntersection=A∪BA∩B
IoU取值范围为0到1。通常,当IoU >= 0.5时,认为预测框与真实框"匹配"。本项目的主要评估指标mAP@0.5就是基于IoU阈值为0.5计算的。
Precision(精确率):
P r e c i s i o n = T P T P + F P Precision = \frac{TP}{TP + FP} Precision=TP+FPTP
精确率衡量模型预测为正的样本中有多少确实是正样本。高精确率意味着误检(将正常区域误判为瑕疵)较少。
Recall(召回率):
R e c a l l = T P T P + F N Recall = \frac{TP}{TP + FN} Recall=TP+FNTP
召回率衡量所有真实正样本中有多少被模型成功检出。高召回率意味着漏检较少。在实际工业场景中,漏检的代价通常高于误检,因此召回率尤为重要。
mAP(mean Average Precision,平均精度均值):
AP是PR曲线(Precision-Recall Curve)下的面积,综合反映了模型在不同置信度阈值下的精确率和召回率表现。mAP则是所有类别AP的算术平均,是目标检测领域最权威的综合性评估指标。
mAP@0.5表示在IoU阈值为0.5时计算的mAP,衡量的是"大致定位正确"的能力;mAP@0.5:0.95表示在IoU从0.5到0.95以0.05为步长取10个阈值,分别计算mAP后取平均,对定位精度要求更严格。
8.2 实验结果汇总
| 指标 | 数值 | 说明 |
|---|---|---|
| mAP@0.5 | 0.9255 | 主要评估指标,工业应用标准 |
| mAP@0.5:0.95 | 0.6910 | COCO标准严格指标 |
| Precision | 0.9090 | 精确率 |
| Recall | 0.8933 | 召回率 |
| 分类准确率 | 94.39% | ResNet18验证集准确率 |
8.3 混淆矩阵分析
混淆矩阵直观展示了模型在各类别上的分类表现。从实验结果的混淆矩阵可以观察到:
- 对角线元素(正确分类)占主导: 大多数瑕疵被正确分类到对应类别,说明模型的整体判别能力较强。
- "擦花"与"碰伤"存在一定混淆: 这两种瑕疵在视觉上具有一定相似性(都表现为表面的线性损伤),模型偶尔会将其相互误判。可以通过增加这两类的训练样本或设计更精细的特征提取策略来区分。
- "涂层开裂"的检测率相对较低: 这主要是由于该类别的训练样本数量最少(仅约65张),模型学习不够充分。数据增强和过采样技术是改善这一问题的有效手段。
- "其他"类别的误判较多: 作为兜底类别,"其他"包含了多种形态各异的非标准瑕疵,类内差异极大,检测难度本身就高于标准类别。
8.4 PR曲线与F1曲线解读
PR曲线展示了模型在不同置信度阈值下精确率和召回率的权衡关系。理想的PR曲线应该尽可能靠近右上角(高精确率、高召回率)。从本项目的PR曲线可以看出:
- 对于"漏底"、“擦花”、"脏点"等样本充足的类别,PR曲线接近右上角,AP值很高。
- 对于"涂层开裂"、"其他"等样本稀缺的类别,PR曲线相对靠近左下角,AP值较低。
F1曲线展示了F1分数随置信度阈值的变化。F1分数是精确率和召回率的调和平均:
F 1 = 2 × P r e c i s i o n × R e c a l l P r e c i s i o n + R e c a l l F1 = 2 \times \frac{Precision \times Recall}{Precision + Recall} F1=2×Precision+RecallPrecision×Recall
F1曲线的峰值对应的置信度阈值,是精确率和召回率达到最佳平衡的点。在实际部署时,可以根据业务需求调整置信度阈值:如果需要更高的精确率(减少误报),可以提高阈值;如果需要更高的召回率(减少漏检),可以降低阈值。
8.5 工业部署可行性评估
从实验结果和系统实现来看,本项目具备实际工业部署的基本条件:
精度方面: mAP@0.5达到92.55%,在工业瑕疵检测领域属于优秀水平,能够满足大多数生产线的质检要求。
速度方面: YOLOv8n在NVIDIA RTX 3060上处理单张640x640图像的推理时间约为5-10毫秒,完全满足实时检测需求(产线速度通常为每秒1-5张图像)。
易用性方面: PyQt5图形界面直观易用,产线操作人员无需编程知识即可使用。检测结果的可视化展示和JSON数据导出功能,便于与MES(制造执行系统)等上层系统集成。
可维护性方面: 模块化代码结构清晰,配置文件集中管理参数,便于后续的功能扩展和参数调优。
当然,从实验室原型到真正的工业部署,还需要考虑更多工程化问题,如多相机并行接入、与PLC的通信对接、异常报警机制、模型在线更新等。这些问题超出了本项目的范围,但为后续的系统升级指明了方向。
九、常见问题与解决方案
在实际开发和部署过程中,我们遇到了许多典型问题。本节总结了最常见的问题及其解决方案,希望能帮助读者少走弯路。
9.1 Windows控制台中文乱码
现象: 运行训练脚本时,控制台输出大量乱码方块或 UnicodeEncodeError 异常。
原因: Windows默认控制台编码为GBK(代码页936),而Ultralytics库使用Rich库输出包含大量Unicode特殊字符(如进度条边框符号、勾选符号等),GBK编码无法表示这些字符。
解决方案: 本项目采用了多重防护措施:
- 在脚本开头设置UTF-8代码页:
os.system('chcp 65001 > nul 2>&1') - 重定向标准输出流编码:
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8') - 禁用Ultralytics的Rich输出:
os.environ['YOLO_VERBOSE'] = 'False' - 将中文日志写入UTF-8文件:使用
logging.FileHandler(..., encoding='utf-8')
推荐做法: 对于新建的Python项目,在 run_gui.py 和 train_detector.py 等入口脚本的最顶部就加入编码处理代码,防患于未然。
9.2 GPU显存不足(CUDA Out of Memory)
现象: 训练过程中报错 RuntimeError: CUDA out of memory。
原因: 模型、数据、梯度、优化器状态等都占用显存。YOLOv8训练时的显存占用约为模型参数量 × 批次大小 × 4~6倍。如果batch size设置过大,或者同时运行了其他占用显存的程序,就容易OOM。
解决方案:
- 降低
batch_size:如从16降至8甚至4。 - 降低
imgsz:如从640降至416。注意图像尺寸降低会影响小瑕疵的检测精度。 - 关闭其他占用显存的程序(如游戏、其他深度学习训练进程)。
- 使用混合精度训练(YOLOv8默认已启用AMP,但如果手动关闭可以重新开启)。
- 如果以上方法都无效,可以改用CPU训练:
device='cpu',只是速度会慢很多。
9.3 PyQt5界面卡顿或无响应
现象: 点击"开始检测"后,界面卡死,标题栏显示"未响应"。
原因: 在主线程中执行了耗时的模型推理操作,阻塞了Qt的事件循环,导致界面无法刷新和响应用户输入。
解决方案: 必须将耗时操作(模型加载、图像推理、训练)移至工作线程(QThread)或子进程(QProcess)。主线程只负责UI更新,通过信号-槽机制与工作线程通信。DetectorWorker 和 TrainingControlTab 中的实现就是正确的范例。
9.4 模型首次加载缓慢
现象: 第一次运行检测时,模型加载需要几十秒甚至几分钟。
原因: YOLOv8首次使用某个模型变体时,会自动从Ultralytics服务器下载预训练权重文件(约6MB)。如果网络连接不畅或访问国外服务器速度较慢,下载时间会很长。
解决方案: 确保网络畅通,或者手动下载权重文件后放入项目目录。下载后的权重会缓存在 ~/.cache/ultralytics/ 目录,后续加载会很快。
9.5 中文标签显示为方块或乱码
现象: 检测框上的中文类别标签显示为"□"或其他乱码符号。
原因: OpenCV的 cv2.putText 不支持中文。我们的绘制方案使用PIL绘制中文,但如果系统中没有安装微软雅黑、黑体、宋体等中文字体,PIL会回退到默认字体,而默认字体通常不支持中文。
解决方案: 确保Windows系统安装了中文字体(Windows系统默认已安装)。如果是精简版系统或Linux系统,需要手动安装中文字体文件(如从Windows复制 msyh.ttc 到Linux的 /usr/share/fonts/ 目录)。
9.6 数据集划分后某些类别丢失
现象: 运行 split_dataset.py 后,训练集或验证集中缺少某些类别的样本。
原因: train_test_split 的分层采样(stratify)要求每个类别至少有2个样本(训练集和验证集各至少1个)。如果某个类别只有1张图片,分层采样会报错或排除该类别。
解决方案: split_dataset.py 中已经实现了自动过滤逻辑:检测并排除样本数不足2个的类别。对于样本极少的类别,建议先通过数据增强(旋转、翻转、裁剪)扩充样本数量,再进行数据集划分。
十、总结与展望
10.1 项目亮点回顾
本文详细介绍了一套完整的铝合金表面瑕疵识别系统,从数据集处理、弱监督标注、模型训练到图形界面开发,涵盖了深度学习工业应用的完整流程。项目的核心亮点包括:
第一,弱监督标注方案的创新性。 针对工业场景下边界框标注成本高昂的难题,我们提出了"分类模型 + Grad-CAM热力图 + 后处理"的自动化标注方案。该方案无需任何人工画框,即可生成质量足以支撑YOLOv8训练的伪标注数据,将标注成本降低了一个数量级。
第二,端到端的完整系统。 项目不仅实现了核心的检测算法,还开发了功能完善的PyQt5图形界面,包括瑕疵检测、训练可视化、数据集浏览和模型训练控制四大模块。系统开箱即用,既适合学术研究,也具备工业落地的雏形。
第三,优异的检测性能。 在11类铝合金瑕疵检测任务上,YOLOv8n模型取得了mAP@0.5=92.55%的优异成绩,证明了弱监督标注方案的有效性和YOLOv8在工业检测场景中的强大能力。
第四,工程化细节的完善。 项目充分考虑了Windows平台下中文乱码、GPU显存管理、多线程GUI交互等实际工程问题,并给出了可靠的解决方案。这些细节对于从"实验室Demo"走向"生产环境应用"至关重要。
10.2 可改进方向
尽管本项目已经实现了预期的功能目标,但在以下几个方面仍有提升空间:
数据层面: 当前数据集规模相对有限(约1500张图片),且存在类别不平衡。可以通过收集更多真实工业数据、使用GAN生成合成数据、对稀缺类别进行重采样等方式扩充和均衡数据集。
算法层面:
- 尝试使用更大规模的YOLOv8变体(如YOLOv8s、YOLOv8m),在计算资源允许的情况下进一步提升检测精度。
- 引入注意力机制(如CBAM)或Transformer结构,增强模型对微小瑕疵的感知能力。
- 探索半监督学习(Semi-Supervised Learning)方案,在少量有标注数据和大量无标注数据上联合训练。
系统层面:
- 增加摄像头实时视频流检测功能,支持从工业相机直接采集图像进行在线检测。
- 开发Web服务端接口(如基于FastAPI),使系统能够以HTTP API的形式被其他业务系统调用。
- 增加模型量化(INT8)和TensorRT加速,进一步提升推理速度,降低部署硬件成本。
- 引入异常检测机制,对于模型置信度极低的样本自动标记为"可疑",转交人工复核。
标注层面: 虽然Grad-CAM伪标注已经取得了很好的效果,但仍有改进空间。例如,可以结合显著性检测(Saliency Detection)和多尺度热力图融合,生成更精确的边界框;或者采用迭代精化策略,用初始YOLO模型的预测结果反馈优化伪标注。




更多推荐

所有评论(0)