📚 课程目录第3章使用OpenVINO加速工业零件缺陷检测模型

3. 使用OpenVINO加速工业零件缺陷检测模型

任务说明

本案例将训练好的工业零件缺陷检测模型（基于PyTorch的Faster R-CNN）转换为OpenVINO IR格式，并使用OpenVINO Runtime在CPU上完成推理加速。目标是在保持检测精度的前提下，获得比原始PyTorch模型更快的推理速度。

环境依赖

OpenVINO 2023.3+（含openvino与openvino.tools）
PyTorch 2.0+（仅用于导出，推理阶段不需要）
torchvision（提供Faster R-CNN模型定义）
Python 3.9+

实现步骤

1. 导出ONNX中间格式

首先将PyTorch模型导出为ONNX格式。这里使用torchvision预训练的Faster R-CNN作为示例，实际使用时替换为你自己的训练权重。

import torch
import torchvision
from torchvision.models.detection import FasterRCNN_ResNet50_FPN_Weights

# 加载模型（示例使用预训练权重，实际替换为你的训练模型）
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
    weights=FasterRCNN_ResNet50_FPN_Weights.DEFAULT
)
model.eval()

# 构造示例输入（Faster R-CNN需要图像列表）
dummy_input = [torch.randn(3, 480, 640)]

# 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    "defect_detector.onnx",
    opset_version=17,
    input_names=["images"],
    output_names=["boxes", "labels", "scores"],
    dynamic_axes={
        "images": {0: "batch"},
        "boxes": {0: "num_detections"},
        "labels": {0: "num_detections"},
        "scores": {0: "num_detections"}
    }
)
print("ONNX导出完成")

2. 转换为OpenVINO IR格式

from openvino.tools import mo
from openvino.runtime import serialize

# 使用Model Optimizer转换
ov_model = mo.convert_model("defect_detector.onnx", compress_to_fp16=True)

# 保存IR格式（.xml + .bin）
serialize(ov_model, "defect_detector.xml", "defect_detector.bin")
print("OpenVINO IR转换完成")

3. OpenVINO推理代码

import cv2
import numpy as np
import openvino.runtime as ov

# 创建Core并加载模型
core = ov.Core()
model = core.read_model("defect_detector.xml")
compiled_model = core.compile_model(model, "CPU")

# 获取输入输出节点
input_layer = compiled_model.input(0)
output_boxes = compiled_model.output(0)
output_labels = compiled_model.output(1)
output_scores = compiled_model.output(2)

def preprocess(image_path):
    """读取图像并转换为模型输入格式"""
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 保持原始尺寸，Faster R-CNN内部会处理缩放
    input_tensor = np.expand_dims(image.transpose(2, 0, 1), axis=0).astype(np.float32)
    input_tensor /= 255.0  # 归一化到[0,1]
    return image, input_tensor

def postprocess(boxes, labels, scores, conf_threshold=0.5):
    """过滤低置信度检测结果"""
    keep = scores > conf_threshold
    return boxes[keep], labels[keep], scores[keep]

def detect(image_path):
    """执行推理并返回检测结果"""
    orig_image, input_tensor = preprocess(image_path)
    
    # 推理
    result = compiled_model([input_tensor])
    boxes = result[output_boxes][0]
    labels = result[output_labels][0]
    scores = result[output_scores][0]
    
    # 后处理
    boxes, labels, scores = postprocess(boxes, labels, scores)
    
    # 绘制结果
    for box, label, score in zip(boxes, labels, scores):
        x1, y1, x2, y2 = box.astype(int)
        cv2.rectangle(orig_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(orig_image, f"defect:{score:.2f}", (x1, y1-5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
    
    return orig_image

# 使用示例
if __name__ == "__main__":
    result_img = detect("sample_defect.jpg")
    cv2.imwrite("result_ov.jpg", result_img)
    print("推理完成，结果保存至 result_ov.jpg")

验证方法

精度验证：使用测试集图像，对比PyTorch原模型与OpenVINO模型的检测框（IoU > 0.5视为一致），确保转换后精度损失在可接受范围（通常mAP下降<1%）。
速度对比：对同一批图像分别用PyTorch和OpenVINO推理，统计平均推理时间。OpenVINO在CPU上通常可获得2-5倍加速。

# 速度对比脚本（简化版）
import time
import torch

# PyTorch推理时间
model_cpu = model.to("cpu")
start = time.time()
with torch.no_grad():
    for _ in range(100):
        _ = model([torch.randn(3, 480, 640)])
pytorch_time = (time.time() - start) / 100

# OpenVINO推理时间
start = time.time()
for _ in range(100):
    _ = compiled_model([np.random.rand(1, 3, 480, 640).astype(np.float32)])
ov_time = (time.time() - start) / 100

print(f"PyTorch CPU平均耗时: {pytorch_time*1000:.2f} ms")
print(f"OpenVINO CPU平均耗时: {ov_time*1000:.2f} ms")
print(f"加速比: {pytorch_time/ov_time:.2f}x")

常见问题

动态shape支持：Faster R-CNN输出维度不固定，导出ONNX时必须配置dynamic_axes，否则转换会失败。
FP16精度：compress_to_fp16=True可减小模型体积并提升速度，若精度下降明显可改为False。
自定义模型：若使用自己的训练模型，只需替换第1步中的模型加载部分，后续转换流程完全一致。

通过标准

成功生成defect_detector.xml与defect_detector.bin文件
OpenVINO推理结果与PyTorch原模型在测试集上检测框IoU>0.5的比例不低于95%
CPU上OpenVINO推理速度不低于PyTorch的2倍