《------往期经典推荐------》
二、机器学习实战专栏【链接】，已更新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()
好了，这篇文章就介绍到这里，喜欢的小伙伴感谢给点个赞和关注，更多精彩内容持续更新~~
关于本篇文章大家有任何建议或意见，欢迎在评论区留言交流！
扫一扫
2164
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
