一文掌握工业缺陷检测项目实战（Pytorch算法训练、部署、C++ DLL制作、Qt集成）

一、前言

在智能制造的浪潮中，将AI视觉检测算法成功部署到产线，是连接实验室与工厂的“最后一公里”。然而，众多教程止步于Python脚本的演示，离真正的工业级应用相去甚远。

这篇博客的任务，便是提供一个 经过重新审视和优化的项目蓝图，旨在定义一个标准化、高可靠性的开发模式。未来承接任何类似的工业瑕疵检测项目，都可以此为范本，高效、可靠地完成从数据处理到最终交付的全过程。

本文遵循一个端到端的黄金标准工作流：

• 奠定通用数据基石：从任何团队都易于上手的LabelMe标注开始，构建标准化的COCO数据集，告别繁琐的格式转换。
• 坚持精度与易用性并重：选用PyTorch官方视觉库torchvision中内置的Faster R-CNN模型。它不仅继承了该算法的高精度特性，而且无需复杂的编译环境，极大地简化了训练流程。
• 构筑现代部署架构：采用业界领先的ONNX Runtime作为部署后端。通过将模型导出为标准的ONNX格式，实现与训练框架的彻底解耦，并通过其C++ API在生产环境中实现高性能的GPU推理。
• 实现算法/应用分离：最终将经过验证的C++推理逻辑封装为动态链接库(DLL)，供Qt等上层应用调用。这种架构是技术上的最佳实践，更是团队高效协同、项目稳健迭代的保证。

跟随本文，不仅可以学会一个案例，更是掌握一套可复制的、足以应对真实工业挑战的AI落地方法论。

二、核心选型解析

2.1 模型选择：为何是Torchvision内置的Faster R-CNN？

• 精度是不可动摇的底线：工业质检中，漏检或误检都可能造成巨大损失。Faster R-CNN作为两阶段检测算法的翘楚，其“先提取候选框，后精细分类回归”的机制，在定位与识别精度上通常优于单阶段算法。
• 小目标检测的天然优势：结合特征金字塔网络(FPN)，模型能够在不同深度的特征图上进行预测，有效融合高层语义与底层细节，对工业图像中常见的微小瑕疵有出色的检测能力。
• 易用性和维护性：Torchvision是PyTorch的核心组件，其提供的模型经过了充分测试，API稳定。相较于需要独立编译安装的第三方库，使用Torchvision可以极大简化环境配置，降低项目维护成本。

2.2 部署方案：为何是ONNX Runtime？

• 跨平台与高性能：ONNX (Open Neural Network Exchange)是一种开放的模型表示格式，而ONNX Runtime是微软推出的、用于加速ONNX模型推理的高性能引擎。它支持多种硬件加速（CPU, GPU via CUDA/TensorRT, etc.），一次导出，即可在Windows, Linux等多平台部署。
• 彻底解耦：将模型转换为.onnx文件，意味着交付产物不再依赖于Python或PyTorch。C++应用只需引入ONNX Runtime库即可运行，部署环境极为纯净。
• 性能优化：ONNX Runtime内部集成了图优化、内核融合等多种技术，并能无缝对接TensorRT等后端，可以自动榨干硬件的极致性能，通常比直接使用LibTorch有更好的开箱性能表现。

2.3 Faster R-CNN算法原理简介

Faster R-CNN（Faster Region-based Convolutional Neural Network）是目标检测领域一个里程碑式的算法，以其卓越的精度而著称。作为一个经典的“两阶段”（Two-Stage）检测器，其核心思想是将检测过程分解为两个主要步骤：首先生成可能包含目标的候选区域，然后对这些区域进行精确的分类和位置回归。

其整体架构主要由以下四个核心部分组成：

1. 特征提取主干网络 (Backbone):
   这是算法的起点。输入图像首先经过一个预训练的深度卷积神经网络（如ResNet、VGG等），用于提取图像的深层特征图（Feature Map）。这些特征图既包含了丰富的语义信息，也保留了物体的空间位置信息，是后续所有步骤的基础。

2. 区域提议网络 (Region Proposal Network, RPN):
   这是Faster R-CNN相较于其前辈（R-CNN, Fast R-CNN）最关键的创新。RPN取代了耗时的“选择性搜索”等传统候选框生成算法，将区域提议功能集成到了神经网络中，实现了端到端的训练。其工作流程如下：

   • 生成锚框 (Anchors)：在主干网络输出的特征图上，以每个像素点为中心，预设多种不同尺度和长宽比的“锚框”。这些锚框覆盖了图像中所有可能的位置和形状。
   • 二分类与回归：RPN通过一个小型的卷积网络，对每个锚框执行两个任务：一是判断该锚框是前景（包含物体）还是背景，二是对前景锚框的位置进行初步的微调，使其更贴近真实物体。
   • 生成候选框 (Proposals)：最后，RPN输出一系列得分较高的、经过初步修正的矩形框，这些框被称为“候选区域”或“感兴趣区域”（Region of Interest, RoI）。

3. RoI池化层 (RoI Pooling / RoI Align):
   由于RPN生成的候选框尺寸各不相同，而后续的全连接层分类器要求固定尺寸的输入。RoI池化层的作用就是将这些不同大小的候选框，从主干网络的特征图中提取出对应的特征，并将其转换为统一大小的特征向量。现代实现中普遍使用RoI Align，它通过双线性插值等方法避免了RoI Pooling中的量化误差，进一步提升了定位精度。

4. 检测头 (Detector Head):
   这是算法的终点。经过RoI池化层输出的、固定尺寸的特征向量被送入检测头，通常由几个全连接层构成。检测头对每个候选框执行两个最终任务：

   • 精细分类：对候选框中的物体进行多类别分类（例如，判断是short、spur还是background）。
   • 边界框精修 (Bounding Box Regression)：对候选框的位置进行第二次、更精确的回归，使其与物体的真实边界框完美对齐。

综上所述，Faster R-CNN的“两阶段”流程可以清晰地概括为：第一阶段，由RPN快速高效地提出高质量的物体候选框；第二阶段，对这些候选框进行精细的分类和定位。这种“先提议、后识别”的策略，是其能够在复杂场景和微小目标检测任务中保持高精度的关键所在。

---

三、基础环境搭建

本章节将详细介绍如何配置一个功能完备的算法训练与导出环境。

3.1 核心依赖：GPU环境

为了充分利用GPU进行高效训练，必须首先搭建好底层的硬件驱动和计算库。

• 必备组件：
  1. NVIDIA显卡驱动
  2. CUDA Toolkit
  3. cuDNN (CUDA Deep Neural Network library)

这些基础组件的安装过程环环相扣，版本匹配至关重要。详细的安装步骤较为繁琐，可以参考我的技术博客，该博客基于Ubuntu搭建深度学习环境，本文推荐在Ubuntu环境下开展算法研发。部署环节再转换到Windows平台。

3.2 Python算法环境

• 环境配置参考：

  1. 操作系统：Ubuntu 20.04
  2. Python：3.8 - 3.10
  3. PyTorch：1.13.1 (或更高版本，需与CUDA版本匹配)
  4. Torchvision：0.14.1 (或与PyTorch匹配的更高版本)

• 安装核心库：
  首先，访问PyTorch官网，根据本地的CUDA版本选择并运行正确的安装命令。

  # 例如，针对CUDA 11.6
  pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu116
  

• 安装其他辅助库：

  pip install opencv-python pycocotools tqdm labelme onnx onnxruntime -i https://pypi.tuna.tsinghua.edu.cn/simple
  

  • 库功能解析：
    • opencv-python: 用于图像的读取与处理。
    • pycocotools: 用于处理COCO数据集格式及进行标准评估。
    • tqdm: 提供可视化的循环进度条。
    • labelme: 图像标注工具。
    • onnx, onnxruntime: 用于将模型导出为ONNX格式并进行验证。

四、数据集准备：从自主标注到COCO格式

数据是AI模型的基石。本章节将分为两大部分：首先介绍一个可以直接使用的高质量PCB缺陷数据集；然后再详细讲解如何使用业界通用的LabelMe工具，从零开始标注自定义工业图像数据集，以创建用于目标检测的标注文件。

4.1 数据集介绍：PCB印刷电路板瑕疵数据集

本文选用一个由北京大学发布的印刷电路板（PCB）瑕疵数据集作为案例。该数据集质量高、标注精细，非常适合用于训练工业缺陷检测模型。

• 数据集网址：https://pan.baidu.com/s/1w942I0wBIRbACoz3lyVOJw?pwd=u9f7 提取码: u9f7
• 数据内容：包含693张适用于检测任务的图像，并已提供与LabelMe兼容的JSON标注文件(共693个json文件)。
• 瑕疵类别 (6种)：
  1. missing_hole (缺失孔)
  2. mouse_bite (鼠标咬伤/板边缺口)
  3. open_circuit (开路)
  4. short (短路)
  5. spur (杂散/铜刺)
  6. spurious_copper (伪铜/多余的铜)

4.2 自行标注指南：使用LabelMe创建目标检测数据

本节讲解如何使用LabelMe工具创建与上述数据集格式完全一致的标注文件。

假设有一批待标注的PCB图像，存放于 my_pcb_images 文件夹。

1. 启动LabelMe并指定目录
在终端中，进入图片文件夹，然后启动LabelMe。这样做可以方便地在图片间切换。

cd path/to/my_pcb_images
labelme .

然后通过界面工具栏的 “打开目录” 按钮打开图像文件夹。

打开后，需要做两个设置：

• 单击菜单栏顶部 文件，将“同时保存图像数据”复选框的勾去掉。
• 单击菜单栏顶部 文件->自动保存。

2. 创建矩形框
对于目标检测任务，一般使用矩形框来标注目标。

• 在LabelMe界面的工具栏中，单击 “Create RectBox” 按钮（或按快捷键 Ctrl+R）。
• 将鼠标移动到图像中的缺陷区域，按住左键拖动，绘制一个刚好能框住整个缺陷的矩形框。

3. 分配类别标签

• 当松开鼠标左键后，会立即弹出一个对话框，要求输入标签名称。
• 在此处输入缺陷的类别，例如 short、spur 等。注意，标签名称中不得包含中文或特殊字符。
• 点击 “OK” 完成该标注。LabelMe会自动在图像同文件夹内生成一个同名的 .json 文件。例如，为 0001.jpg 创建的标注文件就是 0001.json。这个JSON文件详细记录了绘制的所有矩形框的位置和类别信息。可以继续对图中的其他缺陷重复此操作。

完成所有图片的标注后，就拥有了一套与本文提供的公开数据集结构相同的、可用于后续步骤的自定义数据集。

4.3 自动划分训练与验证集

无论是使用下载的已标注数据集，还是刚刚完成了自己的标注，下一步都是将数据划分为训练集（train）和验证集（val）。这一步骤对于评估模型的泛化能力至关重要。

手动按比例分配文件不仅效率低下，而且会影响数据划分的随机性。因此，本文将通过Python脚本自动完成数据集的随机划分。

1. 准备数据源

首先，创建一个项目根目录（例如 PCB_DATASET），并在其中建立一个名为 all_data 的子目录。将所有成对的图片（.jpg）和标注文件（.json）全部放入 all_data 文件夹中。

初始目录结构：

PCB_DATASET/
└── all_data/├── 01_missing_hole_01.jpg├── 01_missing_hole_01.json├── 01_missing_hole_02.jpg├── 01_missing_hole_02.json└── ...

2. 创建并运行划分脚本

在项目根目录（PCB_DATASET的同级目录或内部）创建一个名为 split_dataset.py 的Python脚本，并将以下代码复制进去。

split_dataset.py 脚本：

import os
import glob
import random
import shutil
from tqdm import tqdmdef split_dataset_to_train_val(base_dir, train_ratio=0.8):"""自动将源文件夹中的数据集按比例划分为训练集和验证集。Args:base_dir (str): 数据集根目录。此目录下应包含一个名为 'all_data' 的子目录。train_ratio (float): 训练集所占的比例，剩余的为验证集。"""# 1. 定义源路径和目标路径source_dir = os.path.join(base_dir, "all_data")train_dir = os.path.join(base_dir, "train")val_dir = os.path.join(base_dir, "val")# 确保源目录存在if not os.path.isdir(source_dir):print(f"错误：源目录 '{source_dir}' 不存在。")print("请确保所有图片和JSON文件都已放入 'all_data' 文件夹中。")return# 2. 创建训练集和验证集目录os.makedirs(train_dir, exist_ok=True)os.makedirs(val_dir, exist_ok=True)# 3. 获取所有图片文件并进行随机打乱# 我们基于图片文件进行划分，并假定每个图片都有一个同名的JSON文件image_files = glob.glob(os.path.join(source_dir, "*.jpg"))random.shuffle(image_files)if not image_files:print(f"在 '{source_dir}' 中未找到.jpg文件，程序终止。")return# 4. 计算切分点split_index = int(len(image_files) * train_ratio)# 5. 分配训练集和验证集文件列表train_files = image_files[:split_index]val_files = image_files[split_index:]print(f"共找到 {len(image_files)} 张图片。")print(f"划分比例: {train_ratio*100}% 训练 / {(1-train_ratio)*100}% 验证")print(f"训练集数量: {len(train_files)}, 验证集数量: {len(val_files)}")# 6. 定义文件移动函数def move_files(file_list, destination_dir):for img_path in tqdm(file_list, desc=f"移动文件到 {os.path.basename(destination_dir)}"):# 获取不带扩展名的文件名base_filename = os.path.splitext(os.path.basename(img_path))[0]json_filename = base_filename + ".json"json_path = os.path.join(source_dir, json_filename)# 移动图片文件shutil.move(img_path, os.path.join(destination_dir, os.path.basename(img_path)))# 移动对应的JSON文件if os.path.exists(json_path):shutil.move(json_path, os.path.join(destination_dir, json_filename))# 7. 执行移动move_files(train_files, train_dir)move_files(val_files, val_dir)print("\n数据集划分完成！")# 检查all_data是否为空，如果为空，可以选择删除if not os.listdir(source_dir):os.rmdir(source_dir)print(f"源文件夹 '{source_dir}' 已清空并删除。")if __name__ == '__main__':# --- 配置参数 ---# 1. 设置数据集的根目录dataset_base_dir = './PCB_DATASET' # 2. 设置训练集比例train_split_ratio = 0.9  # --- 执行划分 ---split_dataset_to_train_val(dataset_base_dir, train_split_ratio)

3. 检查结果

运行此脚本后，它会自动将 all_data 文件夹中的文件随机分配到 train 和 val 文件夹中，并保持图片和标注文件的配对关系。脚本执行完毕后，all_data 文件夹会被清空并删除。

此时，项目目录将变成后续步骤所期望的标准结构：

PCB_DATASET/
├── train/
│   ├── 01_missing_hole_01.jpg
│   ├── 01_missing_hole_01.json
│   └── ...
└── val/├── 01_missing_hole_02.jpg├── 01_missing_hole_02.json└── ...

至此，数据集已准备就绪，可以进入下一阶段的格式转换工作。

4.4 格式转换：从LabelMe到COCO

目前的数据集是以“一个图像一个JSON”的形式松散地存在。为了在处理数万甚至数百万张图片时保持极致的读写效率，并方便进行标准的评估，可以采用一种更通用的目标检测数据集格式：COCO数据集格式。

COCO格式的核心思想，是将所有图像的信息（images）、所有标注（annotations）以及所有类别（categories）整合到一个单一的JSON索引文件中。这样做的好处是：

• 高效加载：框架启动时只需加载一个JSON文件，即可掌握整个数据集的全貌，无需频繁地开关和解析成千上万个小文件。
• 标准化：COCO是业界公认的标准，拥有丰富的配套工具和评估指标（如mAP），便于复现和比较。

因此，接下来将编写一个脚本，将train和val文件夹内所有零散的LabelMe JSON文件，分别聚合成train_coco.json和val_coco.json。

在项目根目录（例如PCB_DATASET的同级目录）下创建以下Python脚本。这个脚本经过了优化，会自动扫描数据集来发现所有类别，无需手动指定，这使得它更加通用和不易出错。

labelme_to_coco.py 脚本:

import os
import json
import glob
from tqdm import tqdm
import collectionsdef convert_labelme_to_coco(dataset_base_dir, output_dir):"""将已划分为 train/val 的 LabelMe 格式数据集转换为 COCO 格式。该函数会自动发现所有类别，无需手动指定。Args:dataset_base_dir (str): 包含 train/ 和 val/ 子目录的数据集根目录。output_dir (str): 用于保存生成的 COCO JSON 文件的目录。"""# --- 1. 自动发现所有类别 ---print("Scanning for all unique categories...")all_json_files = glob.glob(os.path.join(dataset_base_dir, "*", "*.json"))if not all_json_files:print(f"Error: No .json files found in subdirectories of '{dataset_base_dir}'. Please check your directory structure.")returnall_categories = set()for json_path in tqdm(all_json_files, desc="Discovering categories"):with open(json_path, 'r', encoding='utf-8') as f:labelme_data = json.load(f)for shape in labelme_data.get("shapes", []):all_categories.add(shape["label"])# 排序以确保每次运行ID一致sorted_categories = sorted(list(all_categories))cat_name_to_id = {name: i for i, name in enumerate(sorted_categories)}# 构建COCO格式的categories字段coco_categories = [{"id": v, "name": k, "supercategory": "defect"} for k, v in cat_name_to_id.items()]print(f"Found {len(sorted_categories)} categories: {sorted_categories}")os.makedirs(output_dir, exist_ok=True)# --- 2. 分别处理 train 和 val 两个子集 ---for split in ["train", "val"]:split_dir = os.path.join(dataset_base_dir, split)if not os.path.isdir(split_dir):print(f"Directory '{split_dir}' not found. Skipping '{split}' split.")continueprint(f"\nProcessing '{split}' split...")# 初始化COCO数据结构coco_output = {"info": {"description": f"PCB Defect Dataset - {split} split"},"licenses": [],"images": [],"annotations": [],"categories": coco_categories}image_id_counter = 0annotation_id_counter = 0json_files = glob.glob(os.path.join(split_dir, "*.json"))for json_path in tqdm(json_files, desc=f"Converting {split} data"):with open(json_path, 'r', encoding='utf-8') as f:labelme_data = json.load(f)# --- a. 添加图像信息 ---image_info = {"id": image_id_counter,"file_name": os.path.basename(labelme_data["imagePath"]),"height": labelme_data["imageHeight"],"width": labelme_data["imageWidth"],"license": 0}coco_output["images"].append(image_info)# --- b. 添加标注信息 ---for shape in labelme_data.get("shapes", []):cat_name = shape["label"]if shape["shape_type"] == "rectangle":points = shape["points"]# LabelMe矩形框为[ [x1,y1], [x2,y2] ]# COCO bbox为 [x_min, y_min, width, height]x_coords = [p[0] for p in points]y_coords = [p[1] for p in points]xmin = min(x_coords)ymin = min(y_coords)xmax = max(x_coords)ymax = max(y_coords)bbox_w = xmax - xminbbox_h = ymax - yminannotation_info = {"id": annotation_id_counter,"image_id": image_id_counter,"category_id": cat_name_to_id[cat_name],"bbox": [xmin, ymin, bbox_w, bbox_h],"area": bbox_w * bbox_h,"segmentation": [], # 目标检测任务通常为空"iscrowd": 0}coco_output["annotations"].append(annotation_info)annotation_id_counter += 1image_id_counter += 1# --- 3. 写入最终的COCO JSON文件 ---output_path = os.path.join(output_dir, f"{split}_coco.json")with open(output_path, 'w', encoding='utf-8') as f:json.dump(coco_output, f, indent=4, ensure_ascii=False)print(f"Successfully created COCO annotation file at: {output_path}")if __name__ == '__main__':# 1. 指定已划分好 train/val 的PCB数据集根目录#    脚本将自动查找此目录下的 'train' 和 'val' 文件夹dataset_base_dir = './PCB_DATASET' # 2. 指定输出COCO json文件的目录coco_output_dir = os.path.join(dataset_base_dir, 'coco_annotations')# 3. 运行转换convert_labelme_to_coco(dataset_base_dir, coco_output_dir)

保存脚本后，直接在终端运行python labelme_to_coco.py。

成功运行后，数据集目录结构将更新为：

PCB_DATASET/
├── train/
│   ├── 01_missing_hole_01.jpg
│   ├── 01_missing_hole_01.json
│   └── ...
├── val/
│   ├── 01_missing_hole_02.jpg
│   ├── 01_missing_hole_02.json
│   └── ...
└── coco_annotations/├── train_coco.json  <-- 生成的训练集索引└── val_coco.json    <-- 生成的验证集索引

至此，数据集已完全准备就绪，可以进行算法模型训练。

五、模型训练：基于Torchvision

本章将使用torchvision内置的Faster R-CNN模型来训练。

5.1 编写训练脚本

在项目根目录下创建train.py文件。

train.py 脚本：

import os
import torch
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.datasets import CocoDetection
import torchvision.transforms as T
from pycocotools.cocoeval import COCOeval
from tqdm import tqdm# --- 1. 自定义数据集类 ---
class PCBDataloader(CocoDetection):def __init__(self, root, annFile, transform=None):super(PCBDataloader, self).__init__(root, annFile, transform=transform)def __getitem__(self, idx):img, target = super(PCBDataloader, self).__getitem__(idx)image_id = self.ids[idx]boxes = []labels = []for annot in target: xmin, ymin, w, h = annot['bbox']xmax, ymax = xmin + w, ymin + hboxes.append([xmin, ymin, xmax, ymax])labels.append(annot['category_id'])boxes = torch.as_tensor(boxes, dtype=torch.float32)if boxes.shape[0] == 0:boxes = torch.zeros((0, 4), dtype=torch.float32)labels = torch.as_tensor(labels, dtype=torch.int64)target_dict = {}target_dict["boxes"] = boxestarget_dict["labels"] = labelsreturn img, target_dict# --- 2. 辅助函数 ---
def get_transform():return T.Compose([T.ToTensor()])def get_model(num_classes):# 加载预训练的Faster R-CNN模型model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)# 获取分类器的输入特征数in_features = model.roi_heads.box_predictor.cls_score.in_features# 替换预训练的头部model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)return model# --- 3. 训练与评估主函数 ---
def main():device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')print(f"Using device: {device}")# --- 数据集路径配置 ---DATASET_ROOT = './PCB_DATASET'TRAIN_IMG_DIR = os.path.join(DATASET_ROOT, "train")VAL_IMG_DIR = os.path.join(DATASET_ROOT, "val")TRAIN_ANN_FILE = os.path.join(DATASET_ROOT, "coco_annotations", "train_coco.json")VAL_ANN_FILE = os.path.join(DATASET_ROOT, "coco_annotations", "val_coco.json")# 类别数 (6种瑕疵 + 1个背景)NUM_CLASSES = 7 # --- 创建数据集和数据加载器 ---dataset_train = PCBDataloader(TRAIN_IMG_DIR, TRAIN_ANN_FILE, get_transform())dataset_val = PCBDataloader(VAL_IMG_DIR, VAL_ANN_FILE, get_transform())data_loader_train = torch.utils.data.DataLoader(dataset_train, batch_size=4, shuffle=True, num_workers=4,collate_fn=lambda x: tuple(zip(*x)))data_loader_val = torch.utils.data.DataLoader(dataset_val, batch_size=1, shuffle=False, num_workers=4,collate_fn=lambda x: tuple(zip(*x)))# --- 初始化模型、优化器 ---model = get_model(NUM_CLASSES)model.to(device)params = [p for p in model.parameters() if p.requires_grad]optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)num_epochs = 100output_dir = "./output_torchvision"os.makedirs(output_dir, exist_ok=True)# --- 训练循环 ---for epoch in range(num_epochs):model.train()for images, targets in tqdm(data_loader_train, desc=f"Epoch {epoch+1}/{num_epochs}"):images = list(image.to(device) for image in images)targets = [{k: v.to(device) for k, v in t.items()} for t in targets]loss_dict = model(images, targets)losses = sum(loss for loss in loss_dict.values())optimizer.zero_grad()losses.backward()optimizer.step()lr_scheduler.step()# --- 每个epoch后进行评估 ---evaluate(model, data_loader_val, device)# --- 保存最终模型 ---final_model_path = os.path.join(output_dir, 'model_final.pth')torch.save(model.state_dict(), final_model_path)print(f"Final model saved to {final_model_path}")def evaluate(model, data_loader, device):model.eval()coco_gt = data_loader.dataset.cocococo_dt_list = []with torch.no_grad():for images, _ in tqdm(data_loader, desc="Evaluating"):images = list(img.to(device) for img in images)outputs = model(images)for i, output in enumerate(outputs):image_id = data_loader.dataset.ids[i]boxes = output['boxes'].cpu().numpy()scores = output['scores'].cpu().numpy()labels = output['labels'].cpu().numpy()for box, score, label in zip(boxes, scores, labels):if score < 0.05: continuex, y, x2, y2 = boxw, h = x2 - x, y2 - ycoco_dt_list.append({'image_id': image_id,'category_id': label,'bbox': [x, y, w, h],'score': score.item()})if not coco_dt_list:print("No detections found for evaluation.")returncoco_dt = coco_gt.loadRes(coco_dt_list)coco_eval = COCOeval(coco_gt, coco_dt, 'bbox')coco_eval.evaluate()coco_eval.accumulate()coco_eval.summarize()if __name__ == '__main__':main()

5.2 关键点解析

1. NUM_CLASSES: 这是最关键的参数。其值必须是 瑕疵类别总数 + 1 (用于背景)。对于PCB数据集，共有6种瑕疵，因此设置为7。
2. get_model: 此函数加载在COCO上预训练的fasterrcnn_resnet50_fpn模型，并将其最后的分类层替换为适应新数据集类别数的新层。
3. PCBDataloader: 自定义的Dataset类，继承自CocoDetection，并对输出格式进行微调，以匹配torchvision模型训练时所需的输入格式。
4. collate_fn: DataLoader中的整理函数，用于处理批次中图像尺寸可能不一的情况。
5. 训练与评估: 脚本包含一个标准的训练循环和一个基于pycocotools的评估函数，可以在每个epoch结束后报告mAP等标准指标。

5.3 启动训练

在终端中，确保处于train.py所在的目录下，然后执行：

python train.py

训练日志和最终的模型权重(model_final.pth)将被保存在./output_torchvision目录中。

5.4 单张图像推理

在完成模型训练后，一个重要的验证步骤是使用训练好的模型对单张图片进行推理，并直观地检查其检测效果。这有助于快速评估模型在特定样本上的表现，并为后续的模型优化或部署提供依据。

下面是一个独立的Python脚本，用于加载已保存的 .pth 权重文件，对指定的单张图片进行预测，并将检测结果（边界框、类别和置信度）可视化地绘制在图片上。

在项目根目录下创建一个名为 inference_image.py 的新文件，并将以下代码复制进去。

inference_image.py 脚本：

import torch
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
import cv2
import torchvision.transforms as T
import randomdef get_model(num_classes):"""加载与训练时结构相同的Faster R-CNN模型。"""# 加载一个在COCO上预训练的模型，作为基础结构model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=None) # weights=None,因为我们要加载自己的权重# 获取分类器的输入特征数in_features = model.roi_heads.box_predictor.cls_score.in_features# 替换为新的头部，类别数为我们自定义的数量model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)return modeldef main():# --- 1. 配置参数 ---# 类别定义 (必须与训练时一致，0通常是背景)CLASS_NAMES = ['__background__', 'missing_hole', 'mouse_bite', 'open_circuit','short', 'spur', 'spurious_copper']NUM_CLASSES = len(CLASS_NAMES)# 模型权重路径MODEL_WEIGHTS_PATH = './output_torchvision/model_final.pth'# 要进行推理的图片路径 (可以从验证集val中任选一张)TEST_IMAGE_PATH = './PCB_DATASET/val/01_missing_hole_08.jpg'# 推理结果保存路径OUTPUT_IMAGE_PATH = './inference_result.jpg'# 检测置信度阈值CONFIDENCE_THRESHOLD = 0.5# --- 2. 加载模型 ---device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')print(f"Using device: {device}")# 初始化模型结构model = get_model(NUM_CLASSES)# 加载训练好的权重model.load_state_dict(torch.load(MODEL_WEIGHTS_PATH, map_location=device))model.to(device)# 设置为评估模式model.eval()print("Model loaded successfully.")# --- 3. 图像预处理 ---# 使用OpenCV读取图片image_bgr = cv2.imread(TEST_IMAGE_PATH)if image_bgr is None:print(f"Error: Could not read image at {TEST_IMAGE_PATH}")return# 将BGR图像转换为RGBimage_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)# 定义转换操作transform = T.Compose([T.ToTensor()])image_tensor = transform(image_rgb).to(device)# 在第0维增加一个批次维度image_tensor = image_tensor.unsqueeze(0)# --- 4. 执行推理 ---with torch.no_grad():predictions = model(image_tensor)# --- 5. 结果解析与可视化 ---# predictions[0]包含了'boxes', 'labels', 'scores'boxes = predictions[0]['boxes'].cpu().numpy()labels = predictions[0]['labels'].cpu().numpy()scores = predictions[0]['scores'].cpu().numpy()# 随机生成颜色以便区分不同类别的框colors = {name: [random.randint(0, 255) for _ in range(3)] for name in CLASS_NAMES}# 绘制结果到原始图片上result_image = image_bgr.copy() # 在BGR图上绘制num_detections = 0for box, label_id, score in zip(boxes, labels, scores):if score > CONFIDENCE_THRESHOLD:num_detections += 1class_name = CLASS_NAMES[label_id]color = colors[class_name]# 坐标转换成整数x1, y1, x2, y2 = map(int, box)# 绘制矩形框cv2.rectangle(result_image, (x1, y1), (x2, y2), color, 2)# 准备要显示的文本label_text = f"{class_name}: {score:.2f}"# 计算文本大小(text_width, text_height), baseline = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)# 绘制文本背景cv2.rectangle(result_image, (x1, y1 - text_height - baseline), (x1 + text_width, y1), color, -1)# 绘制文本cv2.putText(result_image, label_text, (x1, y1 - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)print(f"Found {num_detections} defects with confidence > {CONFIDENCE_THRESHOLD}")# 保存结果图片cv2.imwrite(OUTPUT_IMAGE_PATH, result_image)print(f"Inference result saved to {OUTPUT_IMAGE_PATH}")if __name__ == '__main__':main()

运行脚本：在终端中执行命令：

python inference_image.py

脚本运行后，会在项目根目录下生成一个名为 inference_result.jpg 的文件。打开这张图片，可以看到模型检测出的瑕疵已经被带有类别和置信度分数的彩色矩形框标注了出来。

这一步骤是模型部署前最直观、最有效的“冒烟测试”，确保了算法模型的基本可用性。

六、模型导出为ONNX

模型训练完成后，得到的 model_final.pth 文件本质上是PyTorch框架专用的权重集合。为了让模型能够脱离Python和PyTorch环境，能够被C++等其他语言调用，需要将其转换成一种通用的、跨平台的格式。ONNX (Open Neural Network Exchange) 正是为此而生的行业标准。

接下来的脚本 export_onnx.py 的核心任务就是：加载训练好的PyTorch模型，并将其“翻译”成一个名为 model.onnx 的独立文件。这个过程也被称为“模型序列化”，它将模型的网络结构和权重参数固化下来，使其成为一个可移植的计算图，准备好被任何支持ONNX标准的推理引擎（如ONNX Runtime）加载和执行。

该脚本主要执行以下几个关键步骤：

1. 构建并加载模型：首先，必须创建一个与训练时完全相同的模型结构，然后将保存在 .pth 文件中的权重加载到这个模型中。注意，此时不需要加载预训练权重，因为即将加载的是自己训练好的权重。随后，必须调用 model.eval() 将模型切换到评估模式，这会关闭Dropout等只在训练时使用的层。

2. 创建虚拟输入：ONNX的导出机制是通过“追踪”一次模型的前向传播过程来确定计算图的。因此，我们需要创建一个符合模型输入尺寸要求的张量（Tensor）作为示例输入。这个张量的内容是随机的，但其形状（例如 [1, 3, 800, 800]）必须是确定的，代表着未来C++部署时实际输入的图像尺寸。

3. 执行导出命令：调用PyTorch核心的 torch.onnx.export() 函数。在此函数中，需要指定输入和输出节点的名称（如 'input', 'boxes', 'labels', 'scores'）。为节点命名是一个非常好的工程实践，它使得在后续C++中可以按名获取输入输出，而不是依赖于模糊的索引号，极大地增强了代码的可读性和健壮性。

4. 验证ONNX模型 (可选但推荐)：导出完成后，脚本会使用 onnx 和 onnxruntime 库尝试加载并运行刚刚生成的 model.onnx 文件。这一步就像是质量检查，如果能够成功运行并得到与虚拟输入相对应的输出，就证明模型已经成功转换，并且是准确、有效的，为后续的C++部署扫清了障碍。

在项目根目录下创建export_onnx.py文件，内容如下：

export_onnx.py 脚本：

import torch
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictordef get_model(num_classes):"""构建一个与训练时结构相同的模型实例。"""# 注意：这里不加载预训练权重(pretrained=False或weights=None)，因为我们将加载自己训练好的权重文件。model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=None) in_features = model.roi_heads.box_predictor.cls_score.in_featuresmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)return modeldef main():# --- 1. 配置参数 ---NUM_CLASSES = 7  # 6种瑕疵 + 1个背景MODEL_WEIGHTS_PATH = './output_torchvision/model_final.pth'OUTPUT_ONNX_PATH = './output_torchvision/model.onnx'# 导出模型的固定输入尺寸，必须与C++部署时使用的尺寸一致INPUT_HEIGHT = 800INPUT_WIDTH = 800# --- 2. 加载模型 ---# 导出过程在CPU上进行即可，无需GPUdevice = torch.device('cpu') # 初始化模型结构model = get_model(NUM_CLASSES)# 加载训练好的权重model.load_state_dict(torch.load(MODEL_WEIGHTS_PATH, map_location=device))# 切换到评估模式model.eval()print("Model loaded successfully.")# --- 3. 创建一个符合模型输入的虚拟张量 ---# 形状为 [batch_size, channels, height, width]dummy_input = torch.randn(1, 3, INPUT_HEIGHT, INPUT_WIDTH, device=device)# --- 4. 导出为ONNX ---print(f"Exporting model to {OUTPUT_ONNX_PATH}...")torch.onnx.export(model,dummy_input,OUTPUT_ONNX_PATH,verbose=False,# 为输入和输出节点命名，方便C++端调用input_names=['input'],output_names=['boxes', 'labels', 'scores'],# ONNX算子集版本，11是一个稳定且常用的版本opset_version=11)print("Model has been converted to ONNX successfully.")# --- 5. (可选) 验证导出的ONNX模型 ---try:import onnximport onnxruntime# 检查模型格式是否正确onnx_model = onnx.load(OUTPUT_ONNX_PATH)onnx.checker.check_model(onnx_model)# 使用ONNX Runtime创建一个推理会话ort_session = onnxruntime.InferenceSession(OUTPUT_ONNX_PATH, providers=['CPUExecutionProvider'])def to_numpy(tensor):return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()# 准备输入数据ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(dummy_input)}# 执行推理ort_outs = ort_session.run(None, ort_inputs)print("ONNX model has been checked and tested with ONNX Runtime. It is valid.")except Exception as e:print(f"Error during ONNX model verification: {e}")if __name__ == '__main__':main()

保存脚本后，执行：

python export_onnx.py

此脚本会加载训练好的model_final.pth，创建一个虚拟输入，然后调用torch.onnx.export将其计算图和权重序列化到model.onnx文件中。脚本最后还包含一个验证步骤，使用ONNX Runtime加载并测试导出的模型，确保其有效性。

这个model.onnx文件就是算法阶段的最终交付产物。

七、C++部署实战：ONNX Runtime与控制台验证

进入部署阶段的核心：在C++环境中通过ONNX Runtime调用model.onnx模型。

本章算法集成环境：

• 操作系统：Windows 10
• Visual Studio：2019 或更高版本
• ONNX Runtime：1.13.1 (重要 严格参考ONNX Runtime官网版本发行说明，需要与当前CUDA的版本号一致，本文在Windows上的CUDA版本为11.6)
• OpenCV：4.11.0

7.1 准备依赖库

1. ONNX Runtime: 从其GitHub Releases页面下载预编译的Windows二进制包 (例如 onnxruntime-win-x64-gpu-1.13.1.zip )。解压到一个不含中文或空格的路径，例如 D:\toolplace\onnxruntime。
2. OpenCV: 下载并安装OpenCV，例如 D:\toolplace\opencv。

7.2 创建并配置Visual Studio项目

1. 创建项目: 创建一个新的 “Windows控制台应用程序” 项目，命名为 OnnxInferenceTest。
2. 配置属性 (Release x64):
   • C/C++ -> 常规 -> 附加包含目录:
     • D:\toolplace\onnxruntime\include
     • D:\toolplace\opencv\build\include
   • 链接器 -> 常规 -> 附加库目录:
     • D:\toolplace\onnxruntime\lib
     • D:\toolplace\opencv\build\x64\vc16\lib
   • 链接器 -> 输入 -> 附加依赖项:
     • onnxruntime.lib（尽可能将所有lib文件全部包含进来，后续运行成功后再逐步剔除）
     • opencv_world4110.lib (文件名取决于OpenCV版本)
     • onnxruntime_providers_cuda.lib(若缺失该文件，就只会调用CPU执行推理)
     • onnxruntime_providers_shared.lib

7.3 编写C++推理代码

将以下代码复制到主源文件中，并修改模型和图片路径。

OnnxInferenceTest.cpp

#include <iostream>
#include <vector>
#include <string>
#include <chrono>#include <onnxruntime_cxx_api.h>
#include <opencv2/opencv.hpp>struct DetectionBox {float x1, y1, x2, y2;float score;int class_id;
};int main() {// ================== 1. 配置参数 ==================std::wstring model_path = L"./model.onnx"; // 使用宽字符路径std::string image_path = "./04_mouse_bite_04.jpg";const int INPUT_WIDTH = 800;const int INPUT_HEIGHT = 800;const float CONF_THRESHOLD = 0.5f;const std::vector<std::string> CLASS_NAMES = {"background", "missing_hole", "mouse_bite", "open_circuit","short", "spur", "spurious_copper"};// ================== 2. 初始化ONNX Runtime ==================Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "OnnxInferenceTest");Ort::SessionOptions session_options;// 使用CUDA进行GPU加速OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0);Ort::Session session(env, model_path.c_str(), session_options);// ================== 3. 加载并预处理图像 ==================cv::Mat image = cv::imread(image_path);cv::Mat original_image = image.clone();cv::resize(image, image, cv::Size(INPUT_WIDTH, INPUT_HEIGHT));cv::cvtColor(image, image, cv::COLOR_BGR2RGB); // BGR -> RGBstd::vector<float> input_tensor_values(3 * INPUT_WIDTH * INPUT_HEIGHT);for (int i = 0; i < image.rows; i++) {for (int j = 0; j < image.cols; j++) {input_tensor_values[(0 * image.rows * image.cols) + (i * image.cols) + j] = image.at<cv::Vec3b>(i, j)[0] / 255.0f;input_tensor_values[(1 * image.rows * image.cols) + (i * image.cols) + j] = image.at<cv::Vec3b>(i, j)[1] / 255.0f;input_tensor_values[(2 * image.rows * image.cols) + (i * image.cols) + j] = image.at<cv::Vec3b>(i, j)[2] / 255.0f;}}// 创建输入Tensorstd::vector<int64_t> input_shape = { 1, 3, INPUT_HEIGHT, INPUT_WIDTH };Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, input_tensor_values.data(), input_tensor_values.size(), input_shape.data(), input_shape.size());// ================== 4. 执行推理 ==================const char* input_names[] = { "input" };const char* output_names[] = { "boxes", "labels", "scores" };auto start_time = std::chrono::high_resolution_clock::now();auto output_tensors = session.Run(Ort::RunOptions{ nullptr }, input_names, &input_tensor, 1, output_names, 3);auto end_time = std::chrono::high_resolution_clock::now();auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time);std::cout << "Inference took " << duration.count() << " ms." << std::endl;// ================== 5. 解析输出 ==================// Torchvision Faster R-CNN的输出已经经过NMS处理float* boxes = output_tensors[0].GetTensorMutableData<float>();int64_t* labels = output_tensors[1].GetTensorMutableData<int64_t>();float* scores = output_tensors[2].GetTensorMutableData<float>();auto result_shape = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape();int num_detections = result_shape[0];std::vector<DetectionBox> detections;for (int i = 0; i < num_detections; ++i) {if (scores[i] > CONF_THRESHOLD) {detections.push_back({boxes[i * 4], boxes[i * 4 + 1],boxes[i * 4 + 2], boxes[i * 4 + 3],scores[i],static_cast<int>(labels[i])});}}std::cout << "Found " << detections.size() << " defects." << std::endl;// ================== 6. 结果可视化 ==================float scale_x = (float)original_image.cols / INPUT_WIDTH;float scale_y = (float)original_image.rows / INPUT_HEIGHT;for (const auto& box : detections) {int x1 = static_cast<int>(box.x1 * scale_x);int y1 = static_cast<int>(box.y1 * scale_y);int x2 = static_cast<int>(box.x2 * scale_x);int y2 = static_cast<int>(box.y2 * scale_y);cv::rectangle(original_image, cv::Point(x1, y1), cv::Point(x2, y2), cv::Scalar(0, 255, 0), 2);std::string label = CLASS_NAMES[box.class_id] + ": " + cv::format("%.2f", box.score);cv::putText(original_image, label, cv::Point(x1, y1 - 5), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 255, 0), 1);}cv::imwrite("inference_result_cpp_onnx.jpg", original_image);return 0;
}

7.4 运行与验证

1. 拷贝依赖: 首次编译前，将D:\toolplace\onnxruntime\lib目录下的onnxruntime.dll、onnxruntime_providers_cuda.dll和onnxruntime_providers_shared.dll，以及OpenCV的dll，全部拷贝到项目的x64\Release目录下。
2. 编译运行: 点击VS的 “本地Windows调试器” 按钮。程序将成功加载模型，对图片进行推理，并在项目目录下生成inference_result_cpp_onnx.jpg。

运行后可能会出现下述错误：

Could not locate zlibwapi.dll. Please make sure it is in your library path!

原因就是window系统里没有 zlibwapi.dll。
解决方法就是取zlib官网下载一个 zlibwapi.dll 放到某个path下边，比如 C:\Windows\System32

读者也可以从我的一个资源网站进行下载https://aistudio.baidu.com/datasetdetail/252154，直接下载zlibwapi.dll文件即可，然后放置到C:\Windows\System32路径下。或者放置在项目根目录下也可以。

成功运行后，效果如下：

控制台程序的成功运行，证明了model.onnx和C++ ONNX Runtime环境配置的正确性，可以进行下一步的DLL封装。

八、封装为工业级DLL并集成到Qt（高性能内存版）

将经过验证的C++推理代码封装成动态链接库（DLL），是实现算法与应用分离的关键一步，也是工业软件开发的标准实践。然而，对于性能要求严苛的工业场景，例如高频次的产线拍照检测，通过文件路径在应用程序和算法模块间传递图像，会因频繁的磁盘I/O而成为性能瓶颈。

为解决此问题，本章将采用一种更为高效的架构：通过内存直接传递图像数据。上层应用程序（Qt）负责将图像文件解码到内存中，然后将指向像素数据的指针以及图像元信息（宽高、通道数）直接传递给DLL。DLL则基于此内存块构造一个OpenCV Mat对象进行推理，从而完全规避了磁盘读写开销。

此方法的核心优势在于：

• 极致性能：消除了磁盘I/O延迟，最大化数据传输效率，为实时视频流或高速相机序列检测奠定了基础。
• 模块化与可维护性：算法工程师可以独立更新和迭代DLL（例如，替换为新训练的模型），而无需触碰或重新编译主应用程序。
• 关注点分离：UI开发人员无需配置复杂的AI环境，只需按约定格式准备好内存中的图像数据，即可调用DLL提供的简洁接口集成强大的AI功能。
• 跨语言/平台复用：一个标准的C接口DLL可以被多种语言（C++, C#, Python）和框架调用，具备极高的复用价值。

本章将引导完成从创建DLL项目、定义内存接口，到最终在Qt界面中调用并实现可视化检测的全过程。

8.1 创建DLL项目并定义接口

首先，在Visual Studio中创建一个新项目，用于构建包含推理逻辑的DLL。

1. 创建DLL项目：

   • 在Visual Studio中，选择“创建新项目”，搜索并选择“创建具有导出项的动态链接库”模板。
   • 将项目命名为 PcbDefectDetector。

2. 项目配置：

   • 确保项目平台设置为 Release x64。
   • 重复与上一章节控制台应用完全相同的配置步骤：添加ONNX Runtime和OpenCV的附加包含目录、附加库目录以及附加依赖项。

3. 定义导出接口:
   为了保证最大的兼容性，使用C语言风格的函数接口（通过extern "C"）来避免C++的名称修饰（name mangling）问题。

   创建一个头文件 PcbDefectDetector.h，用于声明将从DLL中导出的函数和数据结构。

   PcbDefectDetector.h

   #ifndef PCB_DEFECT_DETECTOR_H
   #define PCB_DEFECT_DETECTOR_H// 用于在DLL项目内部定义为导出，在外部应用程序中定义为导入
   #ifdef PCBDEFECTDETECTOR_EXPORTS
   #define API __declspec(dllexport)
   #else
   #define API __declspec(dllimport)
   #endif// 定义检测结果的数据结构
   struct DetectionResult {int class_id;char class_name[32]; // 预留足够空间float score;float x1, y1, x2, y2;
   };extern "C" {/*** @brief 初始化检测器* @param model_path ONNX模型文件的路径 (UTF-8编码)* @param use_gpu 是否使用GPU (0 for CPU, 1 for GPU)* @return 返回一个指向检测器实例的句柄，失败则返回 nullptr*/API void* __cdecl create_detector(const char* model_path, int use_gpu);/*** @brief 对内存中的图像数据执行缺陷检测* @param detector_handle create_detector返回的句柄* @param image_data 指向图像像素数据的指针 (期望格式为 BGR BGR BGR ...)* @param image_width 图像宽度* @param image_height 图像高度* @param results 指向 DetectionResult 数组的指针，用于接收结果* @param result_count 指向整数的指针，用于接收检测到的目标数量* @return 成功返回0，失败返回负值*/API int __cdecl detect_image_from_memory(void* detector_handle, const unsigned char* image_data, int image_width, int image_height, DetectionResult** results, int* result_count);/*** @brief 释放检测结果占用的内存* @param results detect_image函数分配的结果数组*/API void __cdecl free_results(DetectionResult* results);/*** @brief 销毁检测器并释放所有资源* @param detector_handle create_detector返回的句柄*/API void __cdecl destroy_detector(void* detector_handle);
   }#endif // PCB_DEFECT_DETECTOR_H
   

   • 关键设计变更：
     • detect_image 函数被替换为 detect_image_from_memory。
     • 新接口不再接收image_path字符串，而是接收一个const unsigned char* image_data指针，以及image_width和image_height整数，用于精确描述内存中的图像。
     • 其他接口（生命周期管理和内存释放）保持不变，这是良好的接口设计实践。

8.2 实现DLL内部逻辑

现在，将推理代码迁移并封装到DLL的源文件 PcbDefectDetector.cpp 中，使其适配新的内存输入接口。

PcbDefectDetector.cpp 完整代码：

#include "pch.h"
#include "PcbDefectDetector.h"#include <onnxruntime_cxx_api.h>
#include <opencv2/opencv.hpp>#include <vector>
#include <string>
#include <iostream>
#include <stdexcept>/*** @class Detector* @brief 一个内部类，封装了所有ONNX Runtime逻辑和状态。** 此类处理ONNX Runtime环境和会话的初始化、图像预处理、模型推理以及结果的后处理。* 它对DLL使用者是隐藏的，使用者仅通过一个不透明指针(void*)与其交互。*/
class Detector {
public:Detector(const char* model_path, bool use_gpu);std::vector<DetectionResult> run_detection(const unsigned char* image_data, int image_width, int image_height);private:// ONNX Runtime成员Ort::Env env;Ort::Session session{ nullptr };Ort::SessionOptions session_options;Ort::MemoryInfo memory_info{ nullptr };// 模型配置const int INPUT_WIDTH = 800;const int INPUT_HEIGHT = 800;const float CONF_THRESHOLD = 0.5f;const std::vector<std::string> CLASS_NAMES = {"background", "missing_hole", "mouse_bite", "open_circuit","short", "spur", "spurious_copper"};const char* input_names[1] = { "input" };const char* output_names[3] = { "boxes", "labels", "scores" };
};// --- Detector类实现 ---Detector::Detector(const char* model_path, bool use_gpu): env(ORT_LOGGING_LEVEL_WARNING, "PcbDefectDetector"),memory_info(Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault)) 
{session_options.SetIntraOpNumThreads(1);if (use_gpu) {OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0);}std::string str_path(model_path);std::wstring wstr_path(str_path.begin(), str_path.end());session = Ort::Session(env, wstr_path.c_str(), session_options);
}/*** @brief 对内存中的图像数据运行推理。* @param image_data 指向BGR格式图像数据的指针。* @param image_width 原始图像宽度。* @param image_height 原始图像高度。* @return 检测到的对象向量。*/
std::vector<DetectionResult> Detector::run_detection(const unsigned char* image_data, int image_width, int image_height) {std::vector<DetectionResult> detections;// 1. 从内存数据构造OpenCV Mat对象// 这是一个高效的操作，通常不会复制数据，而是共享内存。// 假设传入的数据是BGR格式，这与OpenCV的默认格式和Qt转换后的格式一致。cv::Mat image(image_height, image_width, CV_8UC3, const_cast<unsigned char*>(image_data));if (image.empty()) {std::cerr << "Error: Could not create cv::Mat from memory data." << std::endl;return detections;}// 2. 预处理图像cv::Mat resized_image, rgb_image;cv::resize(image, resized_image, cv::Size(INPUT_WIDTH, INPUT_HEIGHT));cv::cvtColor(resized_image, rgb_image, cv::COLOR_BGR2RGB);std::vector<float> input_tensor_values(3 * INPUT_WIDTH * INPUT_HEIGHT);for (int i = 0; i < rgb_image.rows; i++) {for (int j = 0; j < rgb_image.cols; j++) {for (int c = 0; c < 3; c++) {input_tensor_values[c * (INPUT_WIDTH * INPUT_HEIGHT) + i * INPUT_WIDTH + j] =rgb_image.at<cv::Vec3b>(i, j)[c] / 255.0f;}}}// 3. 创建输入Tensorstd::vector<int64_t> input_shape = { 1, 3, INPUT_HEIGHT, INPUT_WIDTH };Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, input_tensor_values.data(), input_tensor_values.size(),input_shape.data(), input_shape.size());// 4. 运行推理auto output_tensors = session.Run(Ort::RunOptions{ nullptr }, input_names, &input_tensor, 1, output_names, 3);// 5. 后处理结果float* boxes = output_tensors[0].GetTensorMutableData<float>();int64_t* labels = output_tensors[1].GetTensorMutableData<int64_t>();float* scores = output_tensors[2].GetTensorMutableData<float>();auto result_shape = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape();int num_detections = static_cast<int>(result_shape[0]);float scale_x = static_cast<float>(image_width) / INPUT_WIDTH;float scale_y = static_cast<float>(image_height) / INPUT_HEIGHT;for (int i = 0; i < num_detections; ++i) {if (scores[i] > CONF_THRESHOLD) {DetectionResult res;res.class_id = static_cast<int>(labels[i]);res.score = scores[i];res.x1 = boxes[i * 4] * scale_x;res.y1 = boxes[i * 4 + 1] * scale_y;res.x2 = boxes[i * 4 + 2] * scale_x;res.y2 = boxes[i * 4 + 3] * scale_y;if (res.class_id >= 0 && res.class_id < CLASS_NAMES.size()) {strncpy_s(res.class_name, sizeof(res.class_name), CLASS_NAMES[res.class_id].c_str(), _TRUNCATE);} else {strncpy_s(res.class_name, sizeof(res.class_name), "unknown", _TRUNCATE);}detections.push_back(res);}}return detections;
}// --- 导出的C风格API实现 ---API void* __cdecl create_detector(const char* model_path, int use_gpu) {try {return new Detector(model_path, use_gpu != 0);} catch (const std::exception& e) {std::cerr << "ERROR creating detector: " << e.what() << std::endl;return nullptr;}
}API int __cdecl detect_image_from_memory(void* detector_handle, const unsigned char* image_data, int image_width, int image_height, DetectionResult** results, int* result_count) {if (!detector_handle || !image_data || !results || !result_count) {return -1; // 无效参数}Detector* detector = static_cast<Detector*>(detector_handle);try {std::vector<DetectionResult> detections = detector->run_detection(image_data, image_width, image_height);*result_count = static_cast<int>(detections.size());if (*result_count == 0) {*results = nullptr;return 0; // 成功，无检测结果}// 为将传递给调用者的C风格数组分配内存。// 此内存必须由调用者使用 free_results() 函数释放。*results = new DetectionResult[*result_count];std::copy(detections.begin(), detections.end(), *results);return 0; // 成功} catch (const std::exception& e) {std::cerr << "ERROR during detection: " << e.what() << std::endl;*results = nullptr;*result_count = 0;return -2; // 检测失败}
}API void __cdecl free_results(DetectionResult* results) {if (results) {delete[] results;}
}API void __cdecl destroy_detector(void* detector_handle) {if (detector_handle) {delete static_cast<Detector*>(detector_handle);}
}

编译此项目，将生成PcbDefectDetector.dll (动态库) 和 PcbDefectDetector.lib (导入库)。

8.3 在Qt中集成并调用DLL

现在，创建一个Qt应用程序来调用这个高性能的AI视觉库。

1. 创建Qt项目：

   • 在Qt Creator中，创建一个新的“Qt Widgets应用程序”。
   • 重要: 确保Qt项目使用的编译器和架构与DLL项目完全一致（例如，MSVC 2019, 64-bit）。

2. 准备文件:

   • 在Qt项目根目录下创建一个名为 3rdparty 的文件夹。
   • 将PcbDefectDetector.h, PcbDefectDetector.lib 拷贝到 3rdparty 文件夹中。
   • 将PcbDefectDetector.dll, onnxruntime.dll, onnxruntime_providers_cuda.dll 以及OpenCV的DLL文件，全部拷贝到Qt的构建输出目录（与生成的可执行文件.exe同级）。

3. 配置.pro文件:
   编辑Qt项目的.pro文件，告诉Qt如何找到头文件和链接库。

   QT       += core gui
   greaterThan(QT_MAJOR_VERSION, 4): QT += widgetsCONFIG += c++17# ... 其他默认配置 ...# 链接 PcbDefectDetector 库
   INCLUDEPATH += $$PWD/3rdparty
   LIBS += -L$$PWD/3rdparty -lPcbDefectDetector
   

4. 编写Qt界面与调用逻辑:
   设计一个简单的界面，包含一个用于显示图片的QLabel和一个“加载图片并检测”的QPushButton。在调用逻辑中，使用QImage加载图片，并将其像素数据直接传递给DLL。

   mainwindow.h

   #ifndef MAINWINDOW_H
   #define MAINWINDOW_H#include <QMainWindow>
   #include "3rdparty/PcbDefectDetector.h" // 包含DLL头文件QT_BEGIN_NAMESPACE
   namespace Ui { class MainWindow; }
   QT_END_NAMESPACEclass MainWindow : public QMainWindow
   {Q_OBJECTpublic:MainWindow(QWidget *parent = nullptr);~MainWindow();private slots:void on_detectButton_clicked();private:Ui::MainWindow *ui;void* m_detector = nullptr; // 用于保存检测器句柄
   };
   #endif // MAINWINDOW_H
   

   mainwindow.cpp

#include "mainwindow.h"
#include "ui_mainwindow.h"
#include <QFileDialog>
#include <QImage>
#include <QPixmap>
#include <QPainter>
#include <QMessageBox>MainWindow::MainWindow(QWidget *parent): QMainWindow(parent), ui(new Ui::MainWindow)
{ui->setupUi(this);// 初始化检测器const char* model_path = "./model.onnx";m_detector = create_detector(model_path, 1); // 1 for GPUif (!m_detector) {QMessageBox::critical(this, "Error", "Failed to initialize detector.");}
}MainWindow::~MainWindow()
{// 释放检测器资源if (m_detector) {destroy_detector(m_detector);}delete ui;
}void MainWindow::on_detectButton_clicked()
{if (!m_detector) return;QString file_path = QFileDialog::getOpenFileName(this, "Select Image", "", "Images (*.png *.jpg *.bmp)");if (file_path.isEmpty()) return;// 1. 使用QImage加载图片到内存QImage image(file_path);if (image.isNull()) {QMessageBox::warning(this, "Error", "Failed to load image.");return;}// 2. 转换图像格式以匹配DLL的输入要求（BGR）QImage rgb_image = image.convertToFormat(QImage::Format_RGB888);QImage bgr_image = rgb_image.rgbSwapped();// 3. 调用DLL进行检测，传递内存指针和图像尺寸DetectionResult* results = nullptr;int count = 0;int status = detect_image_from_memory(m_detector,bgr_image.constBits(), // 获取指向像素数据的指针bgr_image.width(),bgr_image.height(),&results,&count);// 4. 可视化结果QPixmap pixmap = QPixmap::fromImage(image); // 在原始（RGB）图像上绘制if (status == 0 && count > 0) {QPainter painter(&pixmap);for (int i = 0; i < count; ++i) {const auto& res = results[i];painter.setPen(QPen(Qt::red, 2));painter.drawRect(res.x1, res.y1, res.x2 - res.x1, res.y2 - res.y1);QString label = QString("%1: %2").arg(res.class_name).arg(res.score, 0, 'f', 2);painter.drawText(QPoint(res.x1, res.y1 - 5), label);}painter.end();// 5. 释放DLL分配的结果内存！free_results(results);}ui->imageLabel->setPixmap(pixmap.scaled(ui->imageLabel->size(), Qt::KeepAspectRatio, Qt::SmoothTransformation));ui->statusbar->showMessage(QString("Detection finished. Found %1 defects.").arg(count));
}

运行Qt应用程序，点击按钮选择一张图片，即可看到带有检测框和类别信息的可视化结果。通过采用内存传递，整个检测流程的效率得到了显著提升，为构建高性能的工业AI视觉系统打下了坚实的基础。至此，一个从模型训练到工业级高性能DLL封装，再到最终应用程序集成的完整闭环已经成功实现。

九、小结

通过本篇指南，读者不仅掌握了使用PyTorch/torchvision训练一个高精度瑕疵检测模型的技术细节，更重要的是，获得了一套基于ONNX Runtime的、高性能、高可维护性的AI项目开发蓝图。从通用的LabelMe数据处理，到解耦的C++ DLL部署架构，本文所展示的每一步，都旨在解决AI落地过程中的实际痛点。

这种将算法核心封装为独立组件，由上层应用调用的模式，是连接AI研发与工业现场的坚实桥梁，也是实现团队高效协作、保障项目长期稳定交付的关键。