YOLOv8目标检测——详细记录使用ONNX Runtime进行推理部署C++/Python实现

发呆小天才O.o

已于 2025-04-05 14:49:31 修改

阅读量1w

点赞数 137

CC 4.0 BY-SA版权

分类专栏：计算机视觉文章标签： YOLO 目标检测 c++ 计算机视觉图像处理 python opencv

于 2024-12-14 19:58:06 首次发布

本文链接：https://blog.csdn.net/MariLN/article/details/144330414

概述

在之前博客中有介绍YOLOv8从环境安装到训练的完整过程，本节主要介绍ONNX Runtime的原理以及使用其进行推理加速，使用Python、C++两种编程语言来实现。
https://blog.csdn.net/MariLN/article/details/143924548?spm=1001.2014.3001.5501

1. ONNX Runtime

ONNX Runtime是一个由微软推出的跨平台机器学习模型加速器，它仅支持 ONNX 模型格式。它适用于桌面、服务器以及移动设备。

多框架支持：支持多种常见的深度学习框架，如 PyTorch、TensorFlow、Keras、scikit-learn 等，使开发者能轻松将不同框架训练的模型移植到 ONNX Runtime 中进行高效推理，促进了模型在不同框架间的共享与流转。
跨平台兼容性：可在 Linux、Windows、macOS 等多种操作系统上运行，还支持在云、边缘、网页和移动等不同环境中部署，能很好地满足各种应用场景的需求。
硬件优化：针对 GPU、CPU 以及各种 AI 加速器（如 Intel MKL、cuDNN、TensorRT 等）进行了优化，能够充分利用硬件资源提升性能。例如，在 GPU 上可实现并行计算，大大加快模型的推理速度。
高效的内存管理：采用零拷贝（Zero-Copy）技术和内存池管理，减少了数据传输的开销，提升了整体运行速度，在处理大规模数据时优势明显。
动态形状支持：允许输入尺寸在运行时发生变化，模型仍能正确处理，增加了模型应用的灵活性，可更好地适应不同的输入数据情况。

2. 模型转换

2.1 .pt与.onnx模型

2.1.1 pt 模型

.pt 模型是 PyTorch 模型的一种常见存储格式。PyTorch 是一个广泛使用的深度学习框架，在训练神经网络模型时，模型的参数（包括权重和偏置等）会被保存下来，这些参数可以以.pt 文件的形式存储在磁盘中。例如，当你使用 PyTorch 训练一个图像分类模型（如 ResNet）后，通过torch.save()函数就可以将训练好的模型保存为.pt 文件。

本质上它是一个二进制文件，它包含了模型的结构定义和参数。模型的结构定义包括网络的层数、每层的类型（如线性层、卷积层、池化层等）、激活函数的类型等信息。参数则是在训练过程中学习到的具体数值，这些数值决定了模型对输入数据的处理方式。

2.1.2 onnx 模型

ONNX（Open Neural Network Exchange）是一种开放的神经网络交换格式，.onnx 文件就是以这种格式存储的模型文件。它的出现是为了解决不同深度学习框架之间模型转换和互用的问题。许多深度学习框架（如 PyTorch、TensorFlow 等）都可以将自己的模型转换为 ONNX 格式。以 PyTorch 为例，通过torch.onnx.export()函数可以将.pt 模型转换为.onnx 模型。

.onnx 文件同样是一种结构化的文件，它以一种中间表示的形式存储了模型的计算图。这个计算图包含了模型中的各种操作（如加法、乘法、卷积等）以及操作之间的连接关系，同时也包含了模型的输入和输出信息。这种中间表示形式使得不同框架训练的模型能够在一个统一的格式下进行转换和推理。

.onnx 模型主要用于模型的跨框架部署和推理。由于它可以被多种推理引擎（如 ONNX Runtime、TensorRT 等）所支持，所以可以将在一个框架下训练好的模型转换为.onnx 格式，然后在其他环境中进行高效的推理。例如，在工业生产环境中，模型可能是在 PyTorch 中训练的，但在实际的产品线上，需要将其部署到一个对性能和效率要求更高的推理引擎上，此时将模型转换为.onnx 文件并使用 ONNX Runtime 等推理引擎进行部署就非常方便。同时，它也方便了不同团队之间的协作，即使不同团队使用不同的深度学习框架，也可以通过.onnx 文件进行模型的共享和集成。

2.2 .pt转换.onnx

将训练好的 YOLOv8 的.pt模型转换为.onnx模型。可以使用ultralytics库来进行转换。

yolo task=detect mode=export model=./runs/detect/train/weights/best.pt format=onnx

或

from ultralytics import YOLO

# Load a model
model = YOLO('./runs/detect/train/weights/best.pt')  # load a custom trained

# Export the model
success = model.export(format='onnx')

3. 模型推理

3.1 Python实现

3.1.1 环境部署

需要安装onnxruntime、numpy、cv2等库。如果使用 GPU 进行推理，还需安装onnxruntime-gpu。

pip install onnxruntime
pip install onnxruntime-gpu
pip install opencv-python
pip install numpy
pip install gradio

3.1.2 推理步骤

（1）图像预处理

读取图像并将图像的颜色空间从 BGR 格式转换为 RGB 格式。OpenCV 默认使用 BGR 格式，而许多深度学习框架和模型（如 ONNX 模型）则期望输入是 RGB 格式。
调整图像大小，通常将图像 resize 到模型要求的输入尺寸，如 640x640。
对图像进行归一化处理，将像素值归一化到 [0, 1] 区间。
调整图像通道顺序，一般从 HWC（Height, Width, Channel）转换为 CHW 格式，并增加一个批次维度，使其变为 NCHW 格式，N 为批次大小，通常设为 1。

import cv2
import numpy as np

def prepare_input(image, input_width, input_height):
	# 转换为 RGB 格式
	input_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)# cv2.imread 读取到的图像默认是 BGR 格式的

	# 调整图像尺寸
	input_img = cv2.resize(input_img, (input_width, input_height))# input_width、input_height是模型期望的输入宽度和高度

	# 归一化到 0-1
	input_img = input_img / 255.0

	# 变换通道顺序，并增加 batch 维度，HWC->NCHW
	input_img = input_img.transpose(2, 0, 1)
	input_tensor = input_img[np.newaxis, :, :, :].astype(np.float32)# np.newaxis 用于增加一个新的维度

    return input_tensor

image_path = "test.jpg"
image = cv2.imread(image_path)
input_tensor= prepare_input(image, 640, 640)

（2）模型推理

创建onnxruntime.InferenceSession对象，加载转换后的.onnx模型。
将预处理后的图像数据作为输入，传递给模型进行推理，并获取输出结果。

def inference(model_path, input_tensor):

    start = time.perf_counter() # 获取一个高精度的时间戳，主要用于代码性能测试和计算时间间隔，精确度通常远高于 time.time()
    
    # 加载 ONNX 模型
    session = onnxruntime.InferenceSession(model_path, providers=onnxruntime.get_available_providers())

	# 获取输入和输出的名字
    input_names = [model_inputs.name for model_inputs in session.get_inputs()]
    output_names = [model_outputs.name for model_outputs in session.get_outputs()]
    
    # 运行模型推理
    outputs = session.run(output_names, {
   input_names[0]: input_tensor})
   
    print(f"Inference time: {
     (time.perf_counter() - start)*1000:.2f} ms")
    return outputs

（3）后处理

对模型的输出结果去除批量维度。
获取每个检测框的置信度最高的类别，并根据置信度阈值进行筛选，过滤掉低置信度的目标检测框。
坐标转换，将预测框还原到原始图像尺寸，并将边界框的表示从中心点坐标 (x_center, y_center) 和宽高 (w, h) 格式转换为左上角和右下角坐标 (x1, y1, x2, y2) 格式。
进行非极大值抑制（NMS），去除重叠度过高的检测框，得到最终的目标检测结果。

def xywh2xyxy(x):
    # 将边界框从 (x_center, y_center, w, h) 格式转换为 (x1, y1, x2, y2)
    y = np.copy(x)

    # 计算左上角坐标 x1 和 y1
    y[..., 0] = x[..., 0] - x[..., 2] / 2  # x1 = x_center - w / 2
    y[..., 1] = x[..., 1] - x[..., 3] / 2  # y1 = y_center - h / 2

    # 计算右下角坐标 x2 和 y2
    y[..., 2] = x[..., 0] + x[..., 2] / 2  # x2 = x_center + w / 2
    y[..., 3] = x[..., 1] + x[..., 3] / 2  # y2 = y_center + h / 2

    return y

def multiclass_nms(boxes, scores, class_ids, iou_threshold):

    # 获取所有唯一的类别索引
    unique_class_ids = np.unique(class_ids)

    keep_boxes = []  # 存储最终保留的边界框索引
    
    for class_id in unique_class_ids:
        # 筛选出属于当前类别的边界框索引
        class_indices = np.where(class_ids == class_id)[0] # np.where返回元组
        
        # 提取属于当前类别的边界框和分数
        class_boxes = boxes[class_indices, :]   # 当前类别的边界框
        class_scores = scores[class_indices]   # 当前类别的分数

        # 执行 NMS 并获取保留下来的索引
        class_keep_boxes = nms(class_boxes, class_scores, iou_threshold)

        # 将保留的索引（对应原始的索引）添加到结果中
        keep_boxes.extend(class_indices[class_keep_boxes])

    return keep_boxes

def nms(boxes, scores, iou_threshold):
    # 根据 scores 对检测框从高到低进行排序，得到排序后的索引
    sorted_indices = np.argsort(scores)[::-1] # [::-1] 反转排序顺序

    keep_boxes = []
    while sorted_indices.size > 0:
        # 保留最高分数的边界框
        box_id = sorted_indices[0]
        keep_boxes.append(box_id)

        # 计算当前最高分数的边界框与剩余边界框的 IoU
        ious = compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :])

        # 找出 IoU 小于阈值的边界框索引，保留这些框，过滤重叠框
        keep_indices = np.where(ious < iou_threshold)[0]

        # 注意：由于 keep_indices 是相对于 sorted_indices[1:] 的索引，
        # 需要将其整体偏移 +1 来匹配到原始 sorted_indices
        sorted_indices = sorted_indices[keep_indices + 1]

    return keep_boxes

def compute_iou(box, boxes):

    # 计算交集区域的坐标，xmin 和 ymin: 交集左上角的坐标，xmax 和 ymax: 交集右下角的坐标
    xmin = np.maximum(box[0], boxes[:, 0]) 
    ymin = np.maximum(box[1], boxes[:, 1]) 
    xmax = np.minimum(box[2], boxes[:, 2]) 
    ymax = np.minimum(box[3], boxes[:, 3])  

    # 计算交集区域面积，如果两个框没有重叠，交集宽度和高度会为负，使用 np.maximum 保证面积非负
    intersection_area = np.maximum(0, xmax - xmin) * np.maximum(0, ymax - ymin)

    # 计算每个边界框的面积
    box_area = (box[2] - box[0]) * (box[3] - box[1])  
    boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])  

    # 计算并集区域面积
    union_area = box_area + boxes_area - intersection_area

    # 计算 IoU（交并比）
    iou = intersection_area / union_area  # 交集区域面积 / 并集区域面积

    return iou


def process_output(outputs, conf_threshold, iou_threshold, input_width, input_height, img_width, img_height):

    predictions = np.squeeze(outputs[0]).T # 去除数组中形状为1的维度，批量维度，(1, N, M)->(M, N)

    # 获取每个检测框的置信度最高的类别
    scores = np.max(predictions[:, 4:], axis=1) # 在行方向上取最大值
    # 根据置信度阈值过滤掉低置信度的检测框
    predictions = predictions[scores > conf_threshold, :]
    scores = scores[scores > conf_threshold]

    if len(scores) == 0:
        return [], [], []

    # 获取检测框的类别置信度最高的索引
    class_ids = np.argmax(predictions[:, 4:], axis=1) # 返回数组中最大值的索引

    # 提取边界框
    boxes = predictions[:, :4]
    
    # 将边界框坐标从归一化坐标还原到原图尺寸
    input_shape = np.array([input_width, input_height, input_width, input_height])
    boxes = np.divide(boxes, input_shape, dtype=np.float32) # 边界框坐标是相对于输入图像尺寸的，归一化到 [0, 1] 之间
    boxes *= np.array([img_width, img_height, img_width, img_height]) # 将归一化的坐标还原到原图尺寸

    # 转换为 xyxy 格式
    boxes = xywh2xyxy(boxes)

    # 执行非极大值抑制（NMS）
    indices = multiclass_nms(boxes, scores, class_ids, iou_threshold)

    return boxes[indices], scores[indices], class_ids[indices]

3.1.3 代码部署

utils.py

import numpy as np
import cv2

class_names = ['person','head','helmet']

# Create a list of colors for each class where each color is a tuple of 3 integer values
rng = np.random.default_rng(3)
colors = rng.uniform(0, 255, size=(len(class_names), 3))

def nms(boxes, scores, iou_threshold):
    # Sort by score
    sorted_indices = np.argsort(scores)[::-1]

    keep_boxes = []
    while sorted_indices.size > 0:
        # Pick the last box
        box_id = sorted_indices[0]
        keep_boxes.append(box_id)

        # Compute IoU of the picked box with the rest
        ious = compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :])

        # Remove boxes with IoU over the threshold
        keep_indices = np.where(ious < iou_threshold)[0]

        # print(keep_indices.shape, sorted_indices.shape)
        sorted_indices = sorted_indices[keep_indices + 1]

    return keep_boxes

def multiclass_nms(boxes, scores, class_ids, iou_threshold):

    unique_class_ids = np.unique(class_ids)

    keep_boxes = []
    for class_id in unique_class_ids:
        class_indices = np.where(class_ids == class_id)[0]
        class_boxes = boxes[class_indices,:]
        class_scores = scores[class_indices]

        class_keep_boxes = nms(class_boxes, class_scores, iou_threshold)
        keep_boxes.extend(class_indices[class_keep_boxes])

    return keep_boxes

def compute_iou(box, boxes):
    # Compute xmin, ymin, xmax, ymax for both boxes
    xmin = np.maximum(box[0], boxes[:, 0])
    ymin = np.maximum(box[1], boxes[:, 1])
    xmax = np.minimum(box[2], boxes[:, 2])
    ymax = np.minimum(box[3], boxes[:, 3])

    # Compute intersection area
    intersection_area = np.maximum(0, xmax - xmin) * np.maximum(0, ymax - ymin)

    # Compute union area
    box_area = (box[2] - box[0]) * (box[3] - box[1])
    boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    union_area = box_area + boxes_area - intersection_area

    # Compute IoU
    iou = intersection_area / union_area

    return iou

def xywh2xyxy(x):
    # Convert bounding box (x, y, w, h) to bounding box (x1, y1, x2, y2)
    y = np.copy(x)
    y[..., 0] = x[..., 0] - x[..., 2] / 2
    y[..., 1] = x[...

最低0.47元/天解锁文章