【YOLOv8模型onnx部署详解】YOLOv8模型转onnx格式并使用onnxruntime 进行推理部署

阿_旭

17175人浏览 · 2024-06-28 13:01:25

阿_旭 · 2024-06-28 13:01:25 发布

《------往期经典推荐------》

项目名称	项目名称
1.【人脸识别与管理系统开发】	2.【车牌识别与自动收费管理系统开发】
3.【手势识别系统开发】	4.【人脸面部活体检测系统开发】
5.【图片风格快速迁移软件开发】	6.【人脸表表情识别系统】
7.【YOLOv8多目标识别与自动标注软件开发】	8.【基于YOLOv8深度学习的行人跌倒检测系统】
9.【基于YOLOv8深度学习的PCB板缺陷检测系统】	10.【基于YOLOv8深度学习的生活垃圾分类目标检测系统】
11.【基于YOLOv8深度学习的安全帽目标检测系统】	12.【基于YOLOv8深度学习的120种犬类检测与识别系统】
13.【基于YOLOv8深度学习的路面坑洞检测系统】	14.【基于YOLOv8深度学习的火焰烟雾检测系统】
15.【基于YOLOv8深度学习的钢材表面缺陷检测系统】	16.【基于YOLOv8深度学习的舰船目标分类检测系统】
17.【基于YOLOv8深度学习的西红柿成熟度检测系统】	18.【基于YOLOv8深度学习的血细胞检测与计数系统】
19.【基于YOLOv8深度学习的吸烟/抽烟行为检测系统】	20.【基于YOLOv8深度学习的水稻害虫检测与识别系统】
21.【基于YOLOv8深度学习的高精度车辆行人检测与计数系统】	22.【基于YOLOv8深度学习的路面标志线检测与识别系统】
23.【基于YOLOv8深度学习的智能小麦害虫检测识别系统】	24.【基于YOLOv8深度学习的智能玉米害虫检测识别系统】
25.【基于YOLOv8深度学习的200种鸟类智能检测与识别系统】	26.【基于YOLOv8深度学习的45种交通标志智能检测与识别系统】
27.【基于YOLOv8深度学习的人脸面部表情识别系统】	28.【基于YOLOv8深度学习的苹果叶片病害智能诊断系统】
29.【基于YOLOv8深度学习的智能肺炎诊断系统】	30.【基于YOLOv8深度学习的葡萄簇目标检测系统】
31.【基于YOLOv8深度学习的100种中草药智能识别系统】	32.【基于YOLOv8深度学习的102种花卉智能识别系统】
33.【基于YOLOv8深度学习的100种蝴蝶智能识别系统】	34.【基于YOLOv8深度学习的水稻叶片病害智能诊断系统】
35.【基于YOLOv8与ByteTrack的车辆行人多目标检测与追踪系统】	36.【基于YOLOv8深度学习的智能草莓病害检测与分割系统】
37.【基于YOLOv8深度学习的复杂场景下船舶目标检测系统】	38.【基于YOLOv8深度学习的农作物幼苗与杂草检测系统】
39.【基于YOLOv8深度学习的智能道路裂缝检测与分析系统】	40.【基于YOLOv8深度学习的葡萄病害智能诊断与防治系统】
41.【基于YOLOv8深度学习的遥感地理空间物体检测系统】	42.【基于YOLOv8深度学习的无人机视角地面物体检测系统】
43.【基于YOLOv8深度学习的木薯病害智能诊断与防治系统】	44.【基于YOLOv8深度学习的野外火焰烟雾检测系统】
45.【基于YOLOv8深度学习的脑肿瘤智能检测系统】	46.【基于YOLOv8深度学习的玉米叶片病害智能诊断与防治系统】
47.【基于YOLOv8深度学习的橙子病害智能诊断与防治系统】	48.【车辆检测追踪与流量计数系统】
49.【行人检测追踪与双向流量计数系统】	50.【基于YOLOv8深度学习的反光衣检测与预警系统】
51.【危险区域人员闯入检测与报警系统】	52.【高压输电线绝缘子缺陷智能检测系统】

二、机器学习实战专栏【链接】，已更新31期，欢迎关注，持续更新中~~
三、深度学习【Pytorch】专栏【链接】
四、【Stable Diffusion绘画系列】专栏【链接】
五、YOLOv8改进专栏【链接】，持续更新中~~
六、YOLO性能对比专栏【链接】，持续更新中~

《------正文------》

引言

本文主要介绍如何将YOLOv8模型.pt转为onnx格式。然后使用onnxruntime进行目标检测模型的图像推理。包含详细的代码解释和完整代码，供小伙伴们学习交流。如果有帮助，感谢点赞关注。

在这里插入图片描述

详细步骤

导入需要的库

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

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

步骤1：导出ONNX模型

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

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

在这里插入图片描述

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

步骤2：定义YOLOv8类

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

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

2.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数据集的配置文件中加载类别名称，并为这些类别生成随机的颜色调色板，用于后续在图像上绘制检测框时使用。

2.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，用于后续的模型推理。

2.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），并进行归一化处理。

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

2.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方法在输入图像上绘制边界框和标签。

2.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绘制文本背景和文本。

2.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：初始化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阈值。这些参数用于配置模型推理过程。

步骤4：模型推理

output_image = detection.main()

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

步骤5：可视化检测结果

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用于处理键盘输入和关闭窗口。

完整代码

#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()