以下是使用Ultralytics YOLOv11模型转换为ONNX格式并用Python调用的完整指南:
一、模型转换步骤
- 安装依赖库
pip install ultralytics onnx onnxruntime
- 导出ONNX模型
from ultralytics import YOLO
# 加载预训练模型
model = YOLO("yolov11n.pt") # 支持自定义训练模型路径
# 导出为ONNX格式
model.export(
format="onnx",
imgsz=(640, 640), # 输入尺寸
simplify=True, # 简化模型结构
opset=13, # ONNX算子版本
device="cpu" # 导出设备(可选GPU)
)
二、Python调用ONNX模型
1.预处理图像
import cv2
import numpy as np
def preprocess_image(image, input_size):
"""
图像预处理函数
:param image: 输入的图像(BGR 格式)
:param input_size: 模型输入的尺寸 (w, h)
:return: 预处理后的图像张量
"""
# 转换为 RGB 格式
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# 调整图像大小
resized_image = cv2.resize(image, input_size)
# 归一化
input_image = resized_image / 255.0
# 增加批次维度
input_image = np.expand_dims(input_image.transpose(2, 0, 1), axis=0).astype(np.float32)
return resized_image, input_image
2.执行推理
import onnxruntime
def main():
# 加载 ONNX 模型
onnx_path = '/mnt/d/tmp/yolo11n.onnx' # 替换为你的 ONNX 模型路径
session = ort.InferenceSession(onnx_path)
# 获取输入和输出名称
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
# 加载图像
image_path = '../images/dogs.jpg' # 替换为你的图像路径
image = cv2.imread(image_path)
# 预处理图像
input_size = (640, 640)
resized_image, input_image = preprocess_image(image, input_size)
# 运行推理
outputs = session.run([output_name], {input_name: input_image})
# 后处理
boxes, scores, class_ids = postprocess(outputs)
# 将图像转换为 RGB 格式用于 plt 显示
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# 绘制检测结果
result_image = draw_boxes(resized_image, boxes, scores, class_ids, COCO_CLASSES)
# 使用 plt 显示结果
plt.figure(figsize=(10, 10))
plt.imshow(result_image)
plt.axis('off')
plt.show()
3.后处理输出
def postprocess(outputs, conf_threshold=0.5, iou_threshold=0.4):
"""
后处理函数,包括过滤低置信度检测结果和非极大值抑制
:param outputs: 模型的输出
:param conf_threshold: 置信度阈值
:param iou_threshold: 交并比阈值
:return: 过滤后的边界框、置信度和类别索引
"""
predictions = outputs[0]
boxes = []
scores = []
class_ids = []
predictions = predictions[0].T # 输出维度为 (8400, 84)
# 遍历每个检测结果
for pred in predictions:
# 提取置信度
# obj_conf = pred[4]
class_probs = pred[4:]
class_score = np.max(class_probs)
class_id = np.argmax(class_probs)
# confidence = obj_conf * class_score # 最终的置信度是目标置信度和类别概率的乘积
confidence = class_score
if confidence > conf_threshold:
x_center, y_center, width, height = pred[:4]
x1 = int(x_center - width / 2)
y1 = int(y_center - height / 2)
x2 = int(x_center + width / 2)
y2 = int(y_center + height / 2)
boxes.append([x1, y1, x2, y2])
scores.append(confidence)
class_ids.append(class_id)
# 非极大值抑制
indices = cv2.dnn.NMSBoxes(boxes, scores, conf_threshold, iou_threshold)
final_boxes = []
final_scores = []
final_class_ids = []
for i in indices:
i = i.item()
final_boxes.append(boxes[i])
final_scores.append(scores[i])
final_class_ids.append(class_ids[i])
return final_boxes, final_scores, final_class_ids
4.可视化结果
import cv2
import numpy as np
import onnxruntime as ort
import matplotlib.pyplot as plt
# 类别名称列表,这里假设是 COCO 数据集的 80 个类别
COCO_CLASSES = [
'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]
def draw_boxes(image, boxes, scores, class_ids, classes):
"""
在图像上绘制检测结果
:param image: 输入的图像(RGB 格式)
:param boxes: 边界框列表
:param scores: 置信度列表
:param class_ids: 类别索引列表
:param classes: 类别名称列表
:return: 绘制了检测结果的图像
"""
for box, score, class_id in zip(boxes, scores, class_ids):
x1, y1, x2, y2 = box
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
label = f'{classes[class_id]}: {score:.2f}'
cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
return image
三、关键参数说明
参数 | 作用 | 推荐值 |
imgsz | 输入图像尺寸 | 640x640 |
simplify | 优化模型结构 | True |
opset | ONNX算子版本兼容性 | 13+ |
half | FP16量化(需GPU支持) | False(CPU) |
常见问题
- 动态输入尺寸
设置dynamic=True导出支持可变输入尺寸的模型 - INT8量化
添加int8=True参数,需校准数据集 - 性能优化
使用ONNX Runtime的GPU加速:
session = onnxruntime.InferenceSession("model.onnx", providers=['CUDAExecutionProvider'])