YOLOv8不同模型对比和onnx部署详解

wisdom_zhe

于 2024-07-17 20:37:24 发布

阅读量176

点赞数 6

分类专栏：目标检测文章标签： YOLO 目标检测人工智能深度学习 python

本文链接：https://blog.csdn.net/qq_44231797/article/details/140504803

版权

目标检测专栏收录该内容

22 篇文章 15 订阅

订阅专栏

1、前言

本文主要介绍两点：

（1）如何将YOLOv8模型转为其他不同的部署文件格式，并且比较了YOLOv8n.pt的5种不同部署方式:包括原生yolov8n.pt的Pytorch格式、ONNX、OpenVINO-FP32、OpenVINO-int8、TensorRT在CPU和GPU下的推理速度对比；

（2）如何将YOLOv8模型.pt转为onnx格式。然后使用onnxruntime进行目标检测模型的图像推理。

2、YOLOv8模型转为不同的部署文件格式

2.1、模型导出代码

yolov8提供了很简洁的模型转换方式，代码如下所示：

from ultralytics import YOLO

# Load a model
model = YOLO("yolov8n.pt")  # load an official model
# Export the model
model.export(format="onnx")

在上面代码中可以使用 format 参数导出为任何格式，即 format='onnx' 或 format='engine' .我们也可以直接在导出的模型上进行预测或验证 yolo predict model=yolov8n.onnx ，即导出完成后，将显示模型的使用示例。

2.2、可导出的模型格式

可用的 YOLOv8 导出格式如下表所示：

Format	`format` Argument	Model	Metadata	Arguments
PyTorch	-	`yolov8n.pt`	✅	-
TorchScript	`torchscript`	`yolov8n.torchscript`	✅	`imgsz`, `optimize`, `batch`
ONNX	`onnx`	`yolov8n.onnx`	✅	`imgsz`, `half`, `dynamic`, `simplify`, `opset`, `batch`
OpenVINO	`openvino`	`yolov8n_openvino_model/`	✅	`imgsz`, `half`, `int8`, `batch`
TensorRT	`engine`	`yolov8n.engine`	✅	`imgsz`, `half`, `dynamic`, `simplify`, `workspace`, `int8`, `batch`
CoreML	`coreml`	`yolov8n.mlpackage`	✅	`imgsz`, `half`, `int8`, `nms`, `batch`
TF SavedModel	`saved_model`	`yolov8n_saved_model/`	✅	`imgsz`, `keras`, `int8`, `batch`
TF GraphDef	`pb`	`yolov8n.pb`	❌	`imgsz`, `batch`
TF Lite	`tflite`	`yolov8n.tflite`	✅	`imgsz`, `half`, `int8`, `batch`
TF Edge TPU	`edgetpu`	`yolov8n_edgetpu.tflite`	✅	`imgsz`
TF.js	`tfjs`	`yolov8n_web_model/`	✅	`imgsz`, `half`, `int8`, `batch`
PaddlePaddle	`paddle`	`yolov8n_paddle_model/`	✅	`imgsz`, `batch`
NCNN	`ncnn`	`yolov8n_ncnn_model/`	✅	`imgsz`, `half`, `batch`

其中format列表示，导出时format设置的参数名称。Arguments表示导出对应格式时可以额外设置的参数。比如，导出int8格式的openvino模型，代码如下：

from ultralytics import YOLO

# Load a model
model = YOLO("yolov8n.pt")  # load an official model
# Export the model
model.export(format="openvino",int8=True)

2.3、导出模型参数说明

下表详细介绍了可用于将 YOLO 模型导出为不同格式的配置和选项。这些设置对于优化导出模型的性能、大小以及跨各种平台和环境的兼容性至关重要。正确的配置可确保模型已准备好以最佳效率部署在预期应用程序中。

参数	类型	默认值	描述
`format`	`str`	`'torchscript'`	导出模型的目标格式，如 `'onnx'` 、 `'torchscript'` 、 `'tensorflow'` 或其他格式，用于定义与各种部署环境的兼容性。
`imgsz`	`int` 或 `tuple`	`640`	模型输入所需的图像大小。可以是方形图像的整数，也可以是特定尺寸的元组 `(height, width)` 。
`keras`	`bool`	`False`	支持将 TensorFlow SavedModel 导出为 Keras 格式，从而提供与 TensorFlow 服务和 API 的兼容性。
`optimize`	`bool`	`False`	在导出到 TorchScript 时对移动设备应用优化，从而可能减小模型大小并提高性能。
`half`	`bool`	`False`	启用 FP16（半精度）量化，减小模型大小，并可能加快在支持的硬件上的推理速度。
`int8`	`bool`	`False`	激活 INT8 量化，进一步压缩模型并加快推理速度，同时将精度损失降至最低，主要针对边缘设备。
`dynamic`	`bool`	`False`	允许 ONNX 和 TensorRT 导出的动态输入大小，从而增强处理不同图像尺寸的灵活性。
`simplify`	`bool`	`False`	使用 `onnxslim` 简化 ONNX 导出的模型图，从而可能提高性能和兼容性。
`opset`	`int`	`None`	指定 ONNX 操作集版本，以便与不同的 ONNX 分析程序和运行时兼容。如果未设置，则使用支持的最新版本。
`workspace`	`float`	`4.0`	设置最大工作空间大小（以 GiB 为单位），用于 TensorRT 优化，平衡内存使用量和性能。
`nms`	`bool`	`False`	将非最大抑制（NMS）添加到 CoreML 导出中，这对于准确高效的检测后处理至关重要。
`batch`	`int`	`1`	指定导出模型批量推理大小或导出的模型将在模式下 `predict` 并发处理的最大图像数。

调整这些参数可以自定义导出过程以满足特定要求，例如部署环境、硬件约束和性能目标。选择适当的格式和设置对于在模型大小、速度和精度之间实现最佳平衡至关重要。

2.4、模型推理速度对比

本文将yolov8n.pt模型分别导出ONNX、OpenVINO-FP32、OpenVINO-int8、TensorRT这4种格式，加上原生pytorch格式的yolov8n.pt模型，共5种格式模型。分别在CPU与GPU上进行了推理测试，测试结果如下表：

model_name	device	FPS
yolov8n.pt	GPU	77
yolov8n.onnx	GPU	81
yolov8n_openvino_model	GPU	38
yolov8n_int8_openvino_model	GPU	60
yolov8n.engine	GPU	104
yolov8n.pt	cpu	9
yolov8n.onnx	cpu	22
yolov8n_openvino_model	cpu	34
yolov8n_int8_openvino_model	cpu	51
yolov8n.engine	cpu	0

为了更直观的进行推理结果展示，我们直接将表格结果显示为图标形式，绘图代码如下：

import matplotlib.pyplot as plt
import numpy as np
import matplotlib
matplotlib.use('TkAgg')

# 示例数据
categories = ['Pytorch', 'ONNX', 'OpenVINO-FP32','OpenVINO-int8', 'TensorRT']
data_1 = [9, 22, 34, 51, 0]
data_2 = [77, 81, 38, 60,104]
# data_3 = [14, 30, 22, 36]

# 设置柱子宽度和间距
bar_width = 0.25
index = np.arange(len(categories))

# 绘制第一个数据集的条形图
bars1 = plt.bar(index, data_1, bar_width, label='CPU', color='b')

# 绘制第二个数据集的条形图，注意x坐标要偏移以避免重叠
bars2 = plt.bar(index + bar_width, data_2, bar_width, label='GPU', color='r')

# 绘制第三个数据集的条形图，继续偏移
# bars3 = plt.bar(index + 2*bar_width, data_3, bar_width, label='Dataset 3', color='g')

# 在每个柱子上方显示数值
def add_value_labels(ax, bars):
    for bar in bars:
        height = bar.get_height()
        ax.annotate('{}'.format(height),
                    xy=(bar.get_x() + bar.get_width() / 2, height),
                    xytext=(0, 3),  # 3 points vertical offset
                    textcoords="offset points",
                    ha='center', va='bottom')

add_value_labels(plt.gca(), bars1)
add_value_labels(plt.gca(), bars2)


# 设置图表标题和轴标签
plt.title('Comparison of model inference speed')
plt.xlabel('Model Name', fontsize=14)
plt.ylabel('FPS', fontsize=14)
plt.xticks(index + bar_width, categories)

# 创建图例
plt.legend()

# 显示网格
plt.grid(axis='y', linestyle='--', linewidth=0.7, alpha=0.7)

# 显示图表
plt.tight_layout()  # 自动调整子图参数，使之填充整个图像区域
# plt.show()
plt.savefig('chart.jpg')

最终绘制结果如下所示，可以更好的对比不同模型的检测速度。

在这里插入图片描述

从上述结果可以看出：在CPU设备上：

yolov8n.pt模型的性能最低，每秒处理9帧。
yolov8n.onnx模型稍微优于yolov8n.pt，每秒处理22帧。
yolov8n_openvino_model模型的性能最佳，每秒处理34帧。
yolov8n_int8_openvino_model模型略高于yolov8n_openvino_model，每秒处理51帧
yolov8n.engine模型只能在GPU运行，无法测试。

在GPU设备上：

yolov8n.pt模型的性能比CPU处理快很多，每秒处理77帧。
yolov8n.onnx模型稍微优于yolov8n.pt，每秒处理81帧。
yolov8n_openvino_model模型的性能最低，每秒处理38帧。
yolov8n_int8_openvino_model模型略高于yolov8n_openvino_model，每秒处理60帧。
yolov8n.engine模型的性能最佳，每秒处理104帧。

总体上，对于相同的模型和设备，使用GPU比使用CPU获得更高的处理帧数。此外，yolov8n.engine模型在GPU设备上表现最出色，达到了100帧/s；yolov8n.pt与yolov8n.onnx其次,约为80帧/s。在CPU上OpenVINO_int8表现出的性能最佳, 可以达到60帧/s，基本可以满足实际的检测帧率需求。

3、onnx部署详解

3.1、导入需要的库

import argparse
import cv2
import numpy as np
import onnxruntime as ort
import torch
from ultralytics import YOLO

这些库提供了处理命令行参数、图像处理、数组操作、ONNX模型推理等功能

3.2、导出ONNX模型

model = YOLO("yolov8n.pt")
model.export(format="onnx")

首先加载best.pt的YOLO模型，然后将其导出为ONNX格式。运行上述代码后，会在.pt文件的同级目录下生成一个同名的.onnx文件。

ONNX（Open Neural Network Exchange）格式是一种开源格式，用于表示深度学习模型，便于在不同的深度学习框架之间转换和部署模型。

3.3、定义YOLOv8类

class YOLOv8:
    # ... 类的定义 ...

这里定义了一个名为YOLOv8的类，该类封装了使用YOLOv8模型进行目标检测所需的所有功能，包括初始化、预处理、后处理和可视化。

3.3.1、类的初始化

def __init__(self, onnx_model, input_image, confidence_thres, iou_thres):
    """
    初始化YOLOv8类的实例。
    参数:
    onnx_model: ONNX模型的路径。
    input_image: 输入图像的路径。
    confidence_thres: 过滤检测的置信度阈值。
    iou_thres: 非极大抑制的IoU（交并比）阈值。
    """
    self.onnx_model = onnx_model
    self.input_image = input_image
    self.confidence_thres = confidence_thres
    self.iou_thres = iou_thres

    # 从COCO数据集的配置文件加载类别名称
    self.classes = yaml_load(check_yaml("coco8.yaml"))["names"]
    # 字典存储类别名称
    print(self.classes)
    # {0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane'...}

    # 为类别生成颜色调色板
    self.color_palette = np.random.uniform(0, 255, size=(len(self.classes), 3))

    # 初始化ONNX会话
    self.initialize_session(self.onnx_model)

在YOLOv8类的初始化方法中，首先定义了几个关键参数：

onnx_model: ONNX模型的路径。
input_image: 输入图像的路径。
confidence_thres: 过滤检测的置信度阈值。
iou_thres: 非极大抑制的IoU（交并比）阈值。

接着，从COCO数据集的配置文件中加载类别名称，并为这些类别生成随机的颜色调色板，用于后续在图像上绘制检测框时使用。

3.3.2、初始化ONNX会话

def initialize_session(self, onnx_model):
        """
        初始化ONNX模型会话。
        :return:
        """
        if torch.cuda.is_available():
            print("Using CUDA")
            providers = ["CUDAExecutionProvider"]
        else:
            print("Using CPU")
            providers = ["CPUExecutionProvider"]
        session_options = ort.SessionOptions()
        session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        # 使用ONNX模型创建推理会话，并指定执行提供者
        self.session = ort.InferenceSession(onnx_model,
                                            session_options=session_options,
                                            providers=providers)
        return self.session

这个方法用于初始化ONNX模型的推理会话。首先检查CUDA（GPU）是否可用，如果可用，则使用CUDAExecutionProvider，否则使用CPUExecutionProvider。然后创建一个ONNX运行时session，用于后续的模型推理。

3.3.3、预处理图像

def preprocess(self):
     """
        在进行推理之前，对输入图像进行预处理。
        返回:
            image_data: 预处理后的图像数据，准备好进行推理。
        """
        # 使用OpenCV读取输入图像(h,w,c)
        self.img = cv2.imread(self.input_image)

        # 获取输入图像的高度和宽度
        self.img_height, self.img_width = self.img.shape[:2]

        # 将图像颜色空间从BGR转换为RGB
        img = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)

        # 将图像调整为匹配输入形状(640,640,3)
        img = cv2.resize(img, (self.input_width, self.input_height))

        # 将图像数据除以255.0进行归一化
        image_data = np.array(img) / 255.0

        # 转置图像，使通道维度成为第一个维度(3,640,640)
        image_data = np.transpose(image_data, (2, 0, 1))  # 通道优先

        # 扩展图像数据的维度以匹配期望的输入形状(1,3,640,640)
        image_data = np.expand_dims(image_data, axis=0).astype(np.float32)

        # 返回预处理后的图像数据
        return image_data

在preprocess方法中，首先使用OpenCV读取输入图像，并获取其尺寸。

然后将图像从BGR颜色空间转换为RGB，调整其大小以匹配模型的输入尺寸（640x640），并进行归一化处理。

最后，对图像数据进行转置和扩展维度，以匹配模型输入的形状。

3.3.4、后处理输出

def postprocess(self, input_image, output):
    """
    对模型的输出进行后处理，以提取边界框、分数和类别ID。
    参数:
        input_image (numpy.ndarray): 输入图像。
        output (numpy.ndarray): 模型的输出。
    返回:
        numpy.ndarray: 输入图像，上面绘制了检测结果。
    """
    # 转置并压缩输出以匹配期望的形状：(8400, 84)
    outputs = np.transpose(np.squeeze(output[0]))
    # 获取输出数组的行数
    rows = outputs.shape[0]
    # 存储检测到的边界框、分数和类别ID的列表
    boxes = []
    scores = []
    class_ids = []
    # 计算边界框坐标的比例因子
    x_factor = self.img_width / self.input_width
    y_factor = self.img_height / self.input_height

    # 遍历输出数组的每一行
    for i in range(rows):
        # 从当前行提取类别的得分
        classes_scores = outputs[i][4:]
        # 找到类别得分中的最大值
        max_score = np.amax(classes_scores)

        # 如果最大得分大于或等于置信度阈值
        if max_score >= self.confidence_thres:
            # 获取得分最高的类别ID
            class_id = np.argmax(classes_scores)

            # 从当前行提取边界框坐标
            x, y, w, h = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]

            # 计算边界框的缩放坐标
            left = int((x - w / 2) * x_factor)
            top = int((y - h / 2) * y_factor)
            width = int(w * x_factor)
            height = int(h * y_factor)

            # 将类别ID、得分和边界框坐标添加到相应的列表中
            class_ids.append(class_id)
            scores.append(max_score)
            boxes.append([left, top, width, height])

    # 应用非极大抑制以过滤重叠的边界框
    indices = cv2.dnn.NMSBoxes(boxes, scores, self.confidence_thres, self.iou_thres)

    # 遍历非极大抑制后选择的索引
    for i in indices:
        # 获取与索引对应的边界框、得分和类别ID
        box = boxes[i]
        score = scores[i]
        class_id = class_ids[i]
        # 在输入图像上绘制检测结果
        self.draw_detections(input_image, box, score, class_id)
    # 返回修改后的输入图像
    return input_image

postprocess方法对模型的输出进行后处理。核心步骤如下：

首先，转置并压缩输出以匹配期望的形状。

然后，遍历输出数组，过滤掉置信度低于设定阈值的检测，并计算每个检测的边界框坐标。

接着，应用非极大抑制（NMS）来过滤掉重叠的边界框。

最后，对每个剩余的检测，调用draw_detections方法在输入图像上绘制边界框和标签。

3.3.5、绘制检测结果

def draw_detections(self, img, box, score, class_id):
    """
    根据检测到的对象在输入图像上绘制边界框和标签。
    参数:
        img: 要绘制检测的输入图像。
        box: 检测到的边界框。
        score: 对应的检测得分。
        class_id: 检测到的对象的类别ID。
    返回:
        None
    """

    # 提取边界框的坐标
    x1, y1, w, h = box

    # 获取类别ID对应的颜色
    color = self.color_palette[class_id]

    # 在图像上绘制边界框
    cv2.rectangle(img, (int(x1), int(y1)), (int(x1 + w), int(y1 + h)), color, 2)

    # 创建包含类名和得分的标签文本
    label = f"{self.classes[class_id]}: {score:.2f}"

    # 计算标签文本的尺寸
    (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)

    # 计算标签文本的位置
    label_x = x1
    label_y = y1 - 10 if y1 - 10 > label_height else y1 + 10

    # 绘制填充的矩形作为标签文本的背景
    cv2.rectangle(
        img, (label_x, label_y - label_height), (label_x + label_width, label_y + label_height), color, cv2.FILLED
    )

    # 在图像上绘制标签文本
    cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)

draw_detections方法根据检测到的对象在输入图像上绘制边界框和标签。

首先提取边界框的坐标，然后根据类别ID获取颜色，并在图像上绘制边界框。

接着，创建包含类名和得分的标签文本，并计算其在图像上的位置。

最后，cv2.putText、cv2.rectangle绘制文本背景和文本。

3.3.6、主程序

def main(self):
        """
        使用ONNX模型进行推理，并返回带有检测结果的输出图像。
        返回:
            output_img: 带有检测结果的输出图像。
        """
        # 获取模型的输入
        model_inputs = self.session.get_inputs()
        # 保存输入的形状，稍后使用
        # input_shape：(1,3,640,640)
        # self.input_width:640,self.input_height:640
        input_shape = model_inputs[0].shape
        self.input_width = input_shape[2]
        self.input_height = input_shape[3]

        # 对图像数据进行预处理
        img_data = self.preprocess()
        # 使用预处理后的图像数据运行推理,outputs:(1,84,8400)  8400 = 80*80 + 40*40 + 20*20
        outputs = self.session.run(None, {model_inputs[0].name: img_data})
        # 对输出进行后处理以获取输出图像
        return self.postprocess(self.img, outputs)  # 输出图像

main方法是进行模型推理的核心。

首先获取模型的输入形状，然后调用preprocess方法对输入图像进行预处理。

接着，使用预处理后的图像数据运行推理，并调用postprocess方法对输出进行后处理以获取输出图像。

最后，返回这个带有检测结果的图像。

3.3.7、初始化YOLOv8类实例

onnx_model_name = "yolov8n.onnx"
img_path = "test1.jpg"
# 创建用于处理命令行参数的解析器
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, default=onnx_model_name, help="请输入您的ONNX模型路径.")
parser.add_argument("--img", type=str, default=img_path, help="输入图像的路径.")
parser.add_argument("--conf-thres", type=float, default=0.3, help="置信度阈值.")
parser.add_argument("--iou-thres", type=float, default=0.5, help="IoU（交并比）阈值.")
args = parser.parse_args()
detection = YOLOv8(args.model, args.img, args.conf_thres, args.iou_thres)

这一步创建了一个YOLOv8类的实例。实例化时需要传入ONNX模型的路径、输入图像的路径、置信度阈值和IoU阈值。这些参数用于配置模型推理过程。

3.3.8、模型推理

output_image = detection.main()

调用YOLOv8实例的main方法进行模型推理。这个方法首先获取模型的输入形状，然后对输入图像进行预处理，接着运行推理，最后对输出进行后处理以获取带有检测结果的输出图像。

3.3.9、可视化检测结果

cv2.imshow("Output", output_image)
cv2.imwrite('Output.jpg', output_image)
cv2.waitKey(0)
cv2.destroyAllWindows()

最后，使用OpenCV库显示输出图像，并将检测结果保存到文件中。cv2.imshow用于在窗口中显示图像，cv2.imwrite用于将图像保存到文件，cv2.waitKey和cv2.destroyAllWindows用于处理键盘输入和关闭窗口。

4、完整代码

#coding:utf-8
import argparse
import cv2
import numpy as np
import onnxruntime as ort
import torch
from ultralytics.utils import ASSETS, yaml_load
from ultralytics.utils.checks import check_requirements, check_yaml
from ultralytics import YOLO

# 导出onnx模型
model = YOLO("MyModels/best.pt")
model.export(format="onnx")

class YOLOv8:
    """YOLOv8目标检测模型类，用于处理推理和可视化操作。"""
    def __init__(self, onnx_model, input_image, confidence_thres, iou_thres):
        """
        初始化YOLOv8类的实例。
        参数:
            onnx_model: ONNX模型的路径。
            input_image: 输入图像的路径。
            confidence_thres: 过滤检测的置信度阈值。
            iou_thres: 非极大抑制的IoU（交并比）阈值。
        """
        self.onnx_model = onnx_model
        self.input_image = input_image
        self.confidence_thres = confidence_thres
        self.iou_thres = iou_thres

        # 从COCO数据集的配置文件加载类别名称
        self.classes = yaml_load(check_yaml("coco8.yaml"))["names"]
        # 字典存储类别名称
        print(self.classes)
        # {0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane'...}

        # 为类别生成颜色调色板
        self.color_palette = np.random.uniform(0, 255, size=(len(self.classes), 3))

        # 初始化ONNX会话
        self.initialize_session(self.onnx_model)

    def draw_detections(self, img, box, score, class_id):
        """
        根据检测到的对象在输入图像上绘制边界框和标签。
        参数:
            img: 要绘制检测的输入图像。
            box: 检测到的边界框。
            score: 对应的检测得分。
            class_id: 检测到的对象的类别ID。
        返回:
            None
        """

        # 提取边界框的坐标
        x1, y1, w, h = box

        # 获取类别ID对应的颜色
        color = self.color_palette[class_id]

        # 在图像上绘制边界框
        cv2.rectangle(img, (int(x1), int(y1)), (int(x1 + w), int(y1 + h)), color, 2)

        # 创建包含类名和得分的标签文本
        label = f"{self.classes[class_id]}: {score:.2f}"

        # 计算标签文本的尺寸
        (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)

        # 计算标签文本的位置
        label_x = x1
        label_y = y1 - 10 if y1 - 10 > label_height else y1 + 10

        # 绘制填充的矩形作为标签文本的背景
        cv2.rectangle(
            img, (label_x, label_y - label_height), (label_x + label_width, label_y + label_height), color, cv2.FILLED
        )

        # 在图像上绘制标签文本
        cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)

    def preprocess(self):
        """
        在进行推理之前，对输入图像进行预处理。
        返回:
            image_data: 预处理后的图像数据，准备好进行推理。
        """
        # 使用OpenCV读取输入图像(h,w,c)
        self.img = cv2.imread(self.input_image)

        # 获取输入图像的高度和宽度
        self.img_height, self.img_width = self.img.shape[:2]

        # 将图像颜色空间从BGR转换为RGB
        img = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)

        # 将图像调整为匹配输入形状(640,640,3)
        img = cv2.resize(img, (self.input_width, self.input_height))

        # 将图像数据除以255.0进行归一化
        image_data = np.array(img) / 255.0

        # 转置图像，使通道维度成为第一个维度(3,640,640)
        image_data = np.transpose(image_data, (2, 0, 1))  # 通道优先

        # 扩展图像数据的维度以匹配期望的输入形状(1,3,640,640)
        image_data = np.expand_dims(image_data, axis=0).astype(np.float32)

        # 返回预处理后的图像数据
        return image_data

    def postprocess(self, input_image, output):
        """
        对模型的输出进行后处理，以提取边界框、分数和类别ID。
        参数:
            input_image (numpy.ndarray): 输入图像。
            output (numpy.ndarray): 模型的输出。
        返回:
            numpy.ndarray: 输入图像，上面绘制了检测结果。
        """
        # 转置并压缩输出以匹配期望的形状：(8400, 84)
        outputs = np.transpose(np.squeeze(output[0]))
        # 获取输出数组的行数
        rows = outputs.shape[0]
        # 存储检测到的边界框、分数和类别ID的列表
        boxes = []
        scores = []
        class_ids = []
        # 计算边界框坐标的比例因子
        x_factor = self.img_width / self.input_width
        y_factor = self.img_height / self.input_height

        # 遍历输出数组的每一行
        for i in range(rows):
            # 从当前行提取类别的得分
            classes_scores = outputs[i][4:]
            # 找到类别得分中的最大值
            max_score = np.amax(classes_scores)

            # 如果最大得分大于或等于置信度阈值
            if max_score >= self.confidence_thres:
                # 获取得分最高的类别ID
                class_id = np.argmax(classes_scores)

                # 从当前行提取边界框坐标
                x, y, w, h = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]

                # 计算边界框的缩放坐标
                left = int((x - w / 2) * x_factor)
                top = int((y - h / 2) * y_factor)
                width = int(w * x_factor)
                height = int(h * y_factor)

                # 将类别ID、得分和边界框坐标添加到相应的列表中
                class_ids.append(class_id)
                scores.append(max_score)
                boxes.append([left, top, width, height])

        # 应用非极大抑制以过滤重叠的边界框
        indices = cv2.dnn.NMSBoxes(boxes, scores, self.confidence_thres, self.iou_thres)

        # 遍历非极大抑制后选择的索引
        for i in indices:
            # 获取与索引对应的边界框、得分和类别ID
            box = boxes[i]
            score = scores[i]
            class_id = class_ids[i]
            # 在输入图像上绘制检测结果
            self.draw_detections(input_image, box, score, class_id)
        # 返回修改后的输入图像
        return input_image

    def initialize_session(self, onnx_model):
        """
        初始化ONNX模型会话。
        :return:
        """
        if torch.cuda.is_available():
            print("Using CUDA")
            providers = ["CUDAExecutionProvider"]
        else:
            print("Using CPU")
            providers = ["CPUExecutionProvider"]
        session_options = ort.SessionOptions()
        session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        # 使用ONNX模型创建推理会话，并指定执行提供者
        self.session = ort.InferenceSession(onnx_model,
                                            session_options=session_options,
                                            providers=providers)
        return self.session

    def main(self):
        """
        使用ONNX模型进行推理，并返回带有检测结果的输出图像。
        返回:
            output_img: 带有检测结果的输出图像。
        """
        # 获取模型的输入
        model_inputs = self.session.get_inputs()
        # 保存输入的形状，稍后使用
        # input_shape：(1,3,640,640)
        # self.input_width:640,self.input_height:640
        input_shape = model_inputs[0].shape
        self.input_width = input_shape[2]
        self.input_height = input_shape[3]

        # 对图像数据进行预处理
        img_data = self.preprocess()
        # 使用预处理后的图像数据运行推理,outputs:(1,84,8400)  8400 = 80*80 + 40*40 + 20*20
        outputs = self.session.run(None, {model_inputs[0].name: img_data})
        # 对输出进行后处理以获取输出图像
        return self.postprocess(self.img, outputs)  # 输出图像

if __name__ == "__main__":
    onnx_model_name = "yolov8n.onnx"
    img_path = "test1.jpg"
    # 创建用于处理命令行参数的解析器
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, default=onnx_model_name, help="请输入您的ONNX模型路径.")
    parser.add_argument("--img", type=str, default=img_path, help="输入图像的路径.")
    parser.add_argument("--conf-thres", type=float, default=0.3, help="置信度阈值.")
    parser.add_argument("--iou-thres", type=float, default=0.5, help="IoU（交并比）阈值.")
    args = parser.parse_args()

    # 创建YOLOv8实例
    detection = YOLOv8(args.model, args.img, args.conf_thres, args.iou_thres)
    # 模型推理
    output_image = detection.main()

    cv2.namedWindow("Output", cv2.WINDOW_NORMAL)
    cv2.imshow("Output", output_image)
    cv2.imwrite('Output.jpg', output_image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

5、目标检测系列文章

wisdom_zhe

关注

6
点赞
踩
6

收藏

觉得还不错? 一键收藏
0
评论
YOLOv8不同模型对比和onnx部署详解

（1）`如何将YOLOv8模型转为其他不同的部署文件格式`，并且比较了`YOLOv8n.pt`的5种不同部署方式:包括`原生yolov8n.pt的Pytorch格式、ONNX、OpenVINO-FP32、OpenVINO-int8、TensorRT`在CPU和GPU下的推理速度对比；（2）如何将YOLOv8模型`.pt`转为`onnx格式`。然后使用`onnxruntime`进行目标检测模型的图像推理。
复制链接

扫一扫