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

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、目标检测系列文章

1. YOLOv5s网络模型讲解(一看就会)
2. 生活垃圾数据集（YOLO版）
3. YOLOv5如何训练自己的数据集
4. 双向控制舵机（树莓派版）
5. 树莓派部署YOLOv5目标检测（详细篇）
6. YOLO_Tracking 实践 （环境搭建 & 案例测试）
7. 目标检测：数据集划分 & XML数据集转YOLO标签
8. DeepSort行人车辆识别系统（实现目标检测+跟踪+统计）
9. YOLOv5参数大全(parse_opt篇)
10. YOLOv5改进（一）-- 轻量化YOLOv5s模型
11. YOLOv5改进（二）-- 目标检测优化点（添加小目标头检测）
12. YOLOv5改进（三）-- 引进Focaler-IoU损失函数
13. YOLOv5改进（四）–轻量化模型ShuffleNetv2
14. YOLOv5改进（五）-- 轻量化模型MobileNetv3
15. YOLOv5改进（六）–引入YOLOv8中C2F模块
16. YOLOv5改进（七）–改进损失函数EIoU、Alpha-IoU、SIoU、Focal-EIOU
17. YOLOv5改进（八）–引入Soft-NMS非极大值抑制
18. YOLOv5改进（九）–引入BiFPN模块
19. 基于YOLOv10的车辆统计跟踪与车速计算应用
20. 初探 YOLOv8（训练参数解析）

原文地址：https://blog.csdn.net/qq_44231797/article/details/140504803

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！