首页 最新 热门 推荐

  • 首页
  • 最新
  • 热门
  • 推荐

YOLO11 目标检测 | 导出ONNX模型 | ONNX模型推理

  • 25-02-18 14:01
  • 3933
  • 13357
blog.csdn.net

本文分享YOLO11中,从xxx.pt权重文件转为.onnx文件,然后使用.onnx文件,进行目标检测任务的模型推理。

用ONNX模型推理,便于算法到开发板或芯片的部署。

备注:本文是使用Python,编写ONNX模型推理代码的

目录

1、导出ONNX模型

2、目标检测——ONNX模型推理

 2.1、ONNX模型推理-整体流程

2.2、预处理函数

2.3、后处理函数

2.4、检测效果可视化函数

2.5、YOLO11目标检测——ONNX模型推理完整代码 


1、导出ONNX模型

首先我们训练好的模型,生成xxx.pt权重文件;

然后用下面代码,导出ONNX模型(简洁版)

  1. from ultralytics import YOLO
  2. # 加载一个模型,路径为 YOLO 模型的 .pt 文件
  3. model = YOLO("runs/detect/train/weights/best.pt")
  4. # 导出模型,格式为 ONNX
  5. model.export(format="onnx")

运行代码后,会在上面路径中生成best.onnx文件的

  • 比如,填写的路径是:"runs/detect/train/weights/best.pt"
  • 那么在runs/detect/train/weights/目录中,会生成与best.pt同名的onnx文件,即best.onnx

上面代码示例是简单版,如果需要更专业设置ONNX,用下面版本的

YOLO11导出ONNX模型(专业版)

  1. from ultralytics import YOLO
  2. # 加载一个模型,路径为 YOLO 模型的 .pt 文件
  3. model = YOLO(r"runs/detect/train3/weights/best.pt")
  4. # 导出模型,设置多种参数
  5. model.export(
  6. format="onnx", # 导出格式为 ONNX
  7. imgsz=(640, 640), # 设置输入图像的尺寸
  8. keras=False, # 不导出为 Keras 格式
  9. optimize=False, # 不进行优化 False, 移动设备优化的参数,用于在导出为TorchScript 格式时进行模型优化
  10. half=False, # 不启用 FP16 量化
  11. int8=False, # 不启用 INT8 量化
  12. dynamic=False, # 不启用动态输入尺寸
  13. simplify=True, # 简化 ONNX 模型
  14. opset=None, # 使用最新的 opset 版本
  15. workspace=4.0, # 为 TensorRT 优化设置最大工作区大小(GiB)
  16. nms=False, # 不添加 NMS(非极大值抑制)
  17. batch=1, # 指定批处理大小
  18. device="cpu" # 指定导出设备为CPU或GPU,对应参数为"cpu" , "0"
  19. )

对于model.export( )函数中,各种参数说明:

  1. format="onnx":指定导出模型的格式为 onnx。
  2. imgsz=(640, 640):输入图像的尺寸设为 640x640。如果需要其他尺寸可以修改这个值。
  3. keras=False:不导出为 Keras 格式的模型。
  4. optimize=False:不应用 TorchScript 移动设备优化。
  5. half=False:不启用 FP16(半精度)量化。
  6. int8=False:不启用 INT8 量化。
  7. dynamic=False:不启用动态输入尺寸。
  8. simplify=True:简化模型以提升 ONNX 模型的性能。
  9. opset=None:使用默认的 ONNX opset 版本,如果需要可以手动指定。
  10. workspace=4.0:为 TensorRT 优化指定最大工作空间大小为 4 GiB。
  11. nms=False:不为 CoreML 导出添加非极大值抑制(NMS)。
  12. batch=1:设置批处理大小为 1。
  13. device="cpu", 指定导出设备为CPU或GPU,对应参数为"cpu" , "0"

参考官网文档:https://docs.ultralytics.com/modes/export/#arguments

当然了,YOLO11中不仅支持ONNX模型,还支持下面表格中格式

支持的导出格式format参数值生成的模型示例model.export( )函数的参数
PyTorch-yolo11n.pt-
TorchScripttorchscriptyolo11n.torchscriptimgsz, optimize, batch
ONNXonnxyolo11n.onnximgsz, half, dynamic, simplify, opset, batch
OpenVINOopenvinoyolo11n_openvino_model/imgsz, half, int8, batch
TensorRTengineyolo11n.engineimgsz, half, dynamic, simplify, workspace, int8, batch
CoreMLcoremlyolo11n.mlpackageimgsz, half, int8, nms, batch
TF SavedModelsaved_modelyolo11n_saved_model/imgsz, keras, int8, batch
TF GraphDefpbyolo11n.pbimgsz, batch
TF Litetfliteyolo11n.tfliteimgsz, half, int8, batch
TF Edge TPUedgetpuyolo11n_edgetpu.tfliteimgsz
TF.jstfjsyolo11n_web_model/imgsz, half, int8, batch
PaddlePaddlepaddleyolo11n_paddle_model/imgsz, batch
NCNNncnnyolo11n_ncnn_model/imgsz, half, batch

2、目标检测——ONNX模型推理

我们需要编写代码实现了一个使用 ONNXRuntime 执行 YOLOv11 检测模型推理的完整流程,包含图像预处理、推理、后处理和可视化 。

需要编写的代码功能包括:

  • 预处理: 读取输入图像,通过 letterbox 填充图像,自动调整图像尺寸,使其符合模型的输入尺寸要求,并归一化并转换为模型输入的格式 。

  • 模型推理: 使用 ONNXRuntime 加载YOLO11 模型,并根据系统环境自动选择 CPU 或 GPU 执行推理。预处理后的图像输入到模型中,模型返回预测的边界框、类别和分数等输出结果。

  • 后处理: 对模型输出的边界框、类别和分数进行处理,首先根据置信度阈值过滤低置信度的检测结果。然后根据缩放比例将边界框映射回原图的尺寸,应用非极大值抑制(NMS)去除重叠的边界框,最终提取有效的检测结果。

  • 可视化: 在输入图像上绘制检测到的边界框、类别名称和置信度,使用不同颜色区分不同类别的对象。可选择将带有检测结果的图像保存到文件中,最终输出检测结果保存的路径。

 2.1、ONNX模型推理-整体流程

首先编写一个用于运行YOLO11检测模型的推理类:YOLO11

ONNX推理流程:预处理 -> 推理 -> 后处理

  1. class YOLO11:
  2. """YOLO11 目标检测模型类,用于处理推理和可视化。"""
  3. def __init__(self, onnx_model, input_image, confidence_thres, iou_thres):
  4. """
  5. 初始化 YOLO11 类的实例。
  6. 参数:
  7. onnx_model: ONNX 模型的路径。
  8. input_image: 输入图像的路径。
  9. confidence_thres: 用于过滤检测结果的置信度阈值。
  10. iou_thres: 非极大值抑制(NMS)的 IoU(交并比)阈值。
  11. """
  12. self.onnx_model = onnx_model
  13. self.input_image = input_image
  14. self.confidence_thres = confidence_thres
  15. self.iou_thres = iou_thres
  16. # 加载类别名称
  17. self.classes = CLASS_NAMES
  18. # 为每个类别生成一个颜色调色板
  19. self.color_palette = np.random.uniform(0, 255, size=(len(self.classes), 3))
  20. def main(self):
  21. # 使用 ONNX 模型创建推理会话,自动选择CPU或GPU
  22. session = ort.InferenceSession(
  23. self.onnx_model,
  24. providers=["CUDAExecutionProvider", "CPUExecutionProvider"] if ort.get_device() == "GPU" else ["CPUExecutionProvider"],
  25. )
  26. # 打印模型的输入尺寸
  27. print("YOLO11 ? 目标检测 ONNXRuntime")
  28. print("模型名称:", self.onnx_model)
  29. # 获取模型的输入形状
  30. model_inputs = session.get_inputs()
  31. input_shape = model_inputs[0].shape
  32. self.input_width = input_shape[2]
  33. self.input_height = input_shape[3]
  34. print(f"模型输入尺寸:宽度 = {self.input_width}, 高度 = {self.input_height}")
  35. # 预处理图像数据,确保使用模型要求的尺寸 (640x640)
  36. img_data = self.preprocess()
  37. # 使用预处理后的图像数据运行推理
  38. outputs = session.run(None, {model_inputs[0].name: img_data})
  39. # 对输出进行后处理以获取输出图像
  40. return self.postprocess(self.img, outputs) # 输出图像

2.2、预处理函数

然后编写输入图像预处理函数:preprocess

主要功能:对输入的图像进行读取、颜色空间转换、尺寸调整、归一化等预处理操作,并返回适合模型输入的图像数据。

  • 使用 OpenCV 读取图像。
  • 将图像从 BGR 格式转换为 RGB 格式。
  • 使用 letterbox 函数保持图像的宽高比并填充图像,使其符合模型要求的输入尺寸。
  • 归一化图像数据(像素值除以 255.0)。
  • 转换图像通道的维度,符合 PyTorch 模型的输入格式(通道优先)。
  • 扩展图像维度以匹配模型输入的 batch 大小。
  1. def preprocess(self):
  2. """
  3. 对输入图像进行预处理,以便进行推理。
  4. 返回:
  5. image_data: 经过预处理的图像数据,准备进行推理。
  6. """
  7. # 使用 OpenCV 读取输入图像
  8. self.img = cv2.imread(self.input_image)
  9. # 获取输入图像的高度和宽度
  10. self.img_height, self.img_width = self.img.shape[:2]
  11. # 将图像颜色空间从 BGR 转换为 RGB
  12. img = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)
  13. # 保持宽高比,进行 letterbox 填充, 使用模型要求的输入尺寸
  14. img, self.ratio, (self.dw, self.dh) = self.letterbox(img, new_shape=(self.input_width, self.input_height))
  15. # 通过除以 255.0 来归一化图像数据
  16. image_data = np.array(img) / 255.0
  17. # 将图像的通道维度移到第一维
  18. image_data = np.transpose(image_data, (2, 0, 1)) # 通道优先
  19. # 扩展图像数据的维度,以匹配模型输入的形状
  20. image_data = np.expand_dims(image_data, axis=0).astype(np.float32)
  21. # 返回预处理后的图像数据
  22. return image_data

预处理函数preprocess,会依赖letterbox函数。

letterbox函数功能:将图像缩放并进行填充,以保持宽高比,最终将图像调整到指定的输入尺寸。

  • 计算图像的缩放比例,保证宽高比不变。
  • 根据缩放后的图像尺寸,计算需要的填充量(上下左右)。
  • 对图像进行缩放和添加边框(填充),确保最终图像尺寸符合目标输入尺寸。
  • 返回调整后的图像、缩放比例及填充尺寸。
  1. def letterbox(self, img, new_shape=(640, 640), color=(114, 114, 114), auto=False, scaleFill=False, scaleup=True):
  2. """
  3. 将图像进行 letterbox 填充,保持纵横比不变,并缩放到指定尺寸。
  4. """
  5. shape = img.shape[:2] # 当前图像的宽高
  6. print(f"Original image shape: {shape}")
  7. if isinstance(new_shape, int):
  8. new_shape = (new_shape, new_shape)
  9. # 计算缩放比例
  10. r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) # 选择宽高中最小的缩放比
  11. if not scaleup: # 仅缩小,不放大
  12. r = min(r, 1.0)
  13. # 缩放后的未填充尺寸
  14. new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r)))
  15. # 计算需要的填充
  16. dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # 计算填充的尺寸
  17. dw /= 2 # padding 均分
  18. dh /= 2
  19. # 缩放图像
  20. if shape[::-1] != new_unpad: # 如果当前图像尺寸不等于 new_unpad,则缩放
  21. img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
  22. # 为图像添加边框以达到目标尺寸
  23. top, bottom = int(round(dh)), int(round(dh))
  24. left, right = int(round(dw)), int(round(dw))
  25. img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
  26. print(f"Final letterboxed image shape: {img.shape}")
  27. return img, (r, r), (dw, dh)

2.3、后处理函数

再编写后处理函数:postprocess

该代码的功能是对模型的输出结果进行后处理,主要用于从目标检测模型的输出中提取边界框(bounding boxes)、置信度分数和类别ID,并将这些检测结果绘制在输入图像上。 

A、模型输出后处理

  • 将模型的输出数据进行转置和压缩,使其符合处理的预期格式。
  • 遍历模型的每一个检测结果,提取检测框的坐标(x, y, w, h)、类别得分,并通过阈值过滤掉低置信度的检测结果。
  • 计算检测框在原始图像上的实际位置,考虑缩放和填充的影响。

B、缩放和填充调整 

  • 根据图像的缩放比例和填充量,将检测框的坐标从网络输入尺寸(如640x640)调整回原始图像尺寸。

C、非极大值抑制(NMS)

  • 使用非极大值抑制(NMS)算法对检测结果进行过滤,删除重叠的检测框,仅保留最高置信度的框。

D、绘制检测结果

  • 使用 draw_detections 方法在输入图像上绘制边界框,并标注检测类别和置信度。
  • 最终返回包含检测结果的图像,图像上绘制了所有有效的检测框和对应的标签。
  1. def postprocess(self, input_image, output):
  2. """
  3. 对模型输出进行后处理,以提取边界框、分数和类别 ID。
  4. 参数:
  5. input_image (numpy.ndarray): 输入图像。
  6. output (numpy.ndarray): 模型的输出。
  7. 返回:
  8. numpy.ndarray: 包含检测结果的输入图像。
  9. """
  10. # 转置并压缩输出,以匹配预期形状
  11. outputs = np.transpose(np.squeeze(output[0]))
  12. rows = outputs.shape[0]
  13. boxes, scores, class_ids = [], [], []
  14. # 计算缩放比例和填充
  15. ratio = self.img_width / self.input_width, self.img_height / self.input_height
  16. for i in range(rows):
  17. classes_scores = outputs[i][4:]
  18. max_score = np.amax(classes_scores)
  19. if max_score >= self.confidence_thres:
  20. class_id = np.argmax(classes_scores)
  21. x, y, w, h = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]
  22. # 将框调整到原始图像尺寸,考虑缩放和填充
  23. x -= self.dw # 移除填充
  24. y -= self.dh
  25. x /= self.ratio[0] # 缩放回原图
  26. y /= self.ratio[1]
  27. w /= self.ratio[0]
  28. h /= self.ratio[1]
  29. left = int(x - w / 2)
  30. top = int(y - h / 2)
  31. width = int(w)
  32. height = int(h)
  33. boxes.append([left, top, width, height])
  34. scores.append(max_score)
  35. class_ids.append(class_id)
  36. indices = cv2.dnn.NMSBoxes(boxes, scores, self.confidence_thres, self.iou_thres)
  37. for i in indices:
  38. box = boxes[i]
  39. score = scores[i]
  40. class_id = class_ids[i]
  41. self.draw_detections(input_image, box, score, class_id)
  42. return input_image

2.4、检测效果可视化函数

该代码的功能是根据检测结果在输入图像上绘制边界框和标签,以可视化目标检测的结果。

  1. def draw_detections(self, img, box, score, class_id):
  2. """
  3. 在输入图像上绘制检测到的边界框和标签。
  4. 参数:
  5. img: 用于绘制检测结果的输入图像。
  6. box: 检测到的边界框。
  7. score: 对应的检测分数。
  8. class_id: 检测到的目标类别 ID。
  9. 返回:
  10. None
  11. """
  12. # 提取边界框的坐标
  13. x1, y1, w, h = box
  14. # 获取类别对应的颜色
  15. color = self.color_palette[class_id]
  16. # 在图像上绘制边界框
  17. cv2.rectangle(img, (int(x1), int(y1)), (int(x1 + w), int(y1 + h)), color, 2)
  18. # 创建包含类别名和分数的标签文本
  19. label = f"{self.classes[class_id]}: {score:.2f}"
  20. # 计算标签文本的尺寸
  21. (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
  22. # 计算标签文本的位置
  23. label_x = x1
  24. label_y = y1 - 10 if y1 - 10 > label_height else y1 + 10
  25. # 绘制填充的矩形作为标签文本的背景
  26. cv2.rectangle(img, (label_x, label_y - label_height), (label_x + label_width, label_y + label_height), color, cv2.FILLED)
  27. # 在图像上绘制标签文本
  28. cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)

2.5、YOLO11目标检测——ONNX模型推理完整代码 

完整代码,如下所示:

  1. # Ultralytics YOLO ?, AGPL-3.0 license
  2. import argparse
  3. import cv2
  4. import numpy as np
  5. import onnxruntime as ort
  6. # 类外定义类别映射关系,使用字典格式
  7. CLASS_NAMES = {
  8. 0: 'class_name1', # 类别 0 名称
  9. 1: 'class_name2', # 类别 1 名称
  10. 2: 'class_name3' # 类别 1 名称
  11. # 可以添加更多类别...
  12. }
  13. class YOLO11:
  14. """YOLO11 目标检测模型类,用于处理推理和可视化。"""
  15. def __init__(self, onnx_model, input_image, confidence_thres, iou_thres):
  16. """
  17. 初始化 YOLO11 类的实例。
  18. 参数:
  19. onnx_model: ONNX 模型的路径。
  20. input_image: 输入图像的路径。
  21. confidence_thres: 用于过滤检测结果的置信度阈值。
  22. iou_thres: 非极大值抑制(NMS)的 IoU(交并比)阈值。
  23. """
  24. self.onnx_model = onnx_model
  25. self.input_image = input_image
  26. self.confidence_thres = confidence_thres
  27. self.iou_thres = iou_thres
  28. # 加载类别名称
  29. self.classes = CLASS_NAMES
  30. # 为每个类别生成一个颜色调色板
  31. self.color_palette = np.random.uniform(0, 255, size=(len(self.classes), 3))
  32. def preprocess(self):
  33. """
  34. 对输入图像进行预处理,以便进行推理。
  35. 返回:
  36. image_data: 经过预处理的图像数据,准备进行推理。
  37. """
  38. # 使用 OpenCV 读取输入图像
  39. self.img = cv2.imread(self.input_image)
  40. # 获取输入图像的高度和宽度
  41. self.img_height, self.img_width = self.img.shape[:2]
  42. # 将图像颜色空间从 BGR 转换为 RGB
  43. img = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)
  44. # 保持宽高比,进行 letterbox 填充, 使用模型要求的输入尺寸
  45. img, self.ratio, (self.dw, self.dh) = self.letterbox(img, new_shape=(self.input_width, self.input_height))
  46. # 通过除以 255.0 来归一化图像数据
  47. image_data = np.array(img) / 255.0
  48. # 将图像的通道维度移到第一维
  49. image_data = np.transpose(image_data, (2, 0, 1)) # 通道优先
  50. # 扩展图像数据的维度,以匹配模型输入的形状
  51. image_data = np.expand_dims(image_data, axis=0).astype(np.float32)
  52. # 返回预处理后的图像数据
  53. return image_data
  54. def letterbox(self, img, new_shape=(640, 640), color=(114, 114, 114), auto=False, scaleFill=False, scaleup=True):
  55. """
  56. 将图像进行 letterbox 填充,保持纵横比不变,并缩放到指定尺寸。
  57. """
  58. shape = img.shape[:2] # 当前图像的宽高
  59. print(f"Original image shape: {shape}")
  60. if isinstance(new_shape, int):
  61. new_shape = (new_shape, new_shape)
  62. # 计算缩放比例
  63. r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) # 选择宽高中最小的缩放比
  64. if not scaleup: # 仅缩小,不放大
  65. r = min(r, 1.0)
  66. # 缩放后的未填充尺寸
  67. new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r)))
  68. # 计算需要的填充
  69. dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # 计算填充的尺寸
  70. dw /= 2 # padding 均分
  71. dh /= 2
  72. # 缩放图像
  73. if shape[::-1] != new_unpad: # 如果当前图像尺寸不等于 new_unpad,则缩放
  74. img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
  75. # 为图像添加边框以达到目标尺寸
  76. top, bottom = int(round(dh)), int(round(dh))
  77. left, right = int(round(dw)), int(round(dw))
  78. img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
  79. print(f"Final letterboxed image shape: {img.shape}")
  80. return img, (r, r), (dw, dh)
  81. def postprocess(self, input_image, output):
  82. """
  83. 对模型输出进行后处理,以提取边界框、分数和类别 ID。
  84. 参数:
  85. input_image (numpy.ndarray): 输入图像。
  86. output (numpy.ndarray): 模型的输出。
  87. 返回:
  88. numpy.ndarray: 包含检测结果的输入图像。
  89. """
  90. # 转置并压缩输出,以匹配预期形状
  91. outputs = np.transpose(np.squeeze(output[0]))
  92. rows = outputs.shape[0]
  93. boxes, scores, class_ids = [], [], []
  94. # 计算缩放比例和填充
  95. ratio = self.img_width / self.input_width, self.img_height / self.input_height
  96. for i in range(rows):
  97. classes_scores = outputs[i][4:]
  98. max_score = np.amax(classes_scores)
  99. if max_score >= self.confidence_thres:
  100. class_id = np.argmax(classes_scores)
  101. x, y, w, h = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]
  102. # 将框调整到原始图像尺寸,考虑缩放和填充
  103. x -= self.dw # 移除填充
  104. y -= self.dh
  105. x /= self.ratio[0] # 缩放回原图
  106. y /= self.ratio[1]
  107. w /= self.ratio[0]
  108. h /= self.ratio[1]
  109. left = int(x - w / 2)
  110. top = int(y - h / 2)
  111. width = int(w)
  112. height = int(h)
  113. boxes.append([left, top, width, height])
  114. scores.append(max_score)
  115. class_ids.append(class_id)
  116. indices = cv2.dnn.NMSBoxes(boxes, scores, self.confidence_thres, self.iou_thres)
  117. for i in indices:
  118. box = boxes[i]
  119. score = scores[i]
  120. class_id = class_ids[i]
  121. self.draw_detections(input_image, box, score, class_id)
  122. return input_image
  123. def draw_detections(self, img, box, score, class_id):
  124. """
  125. 在输入图像上绘制检测到的边界框和标签。
  126. 参数:
  127. img: 用于绘制检测结果的输入图像。
  128. box: 检测到的边界框。
  129. score: 对应的检测分数。
  130. class_id: 检测到的目标类别 ID。
  131. 返回:
  132. None
  133. """
  134. # 提取边界框的坐标
  135. x1, y1, w, h = box
  136. # 获取类别对应的颜色
  137. color = self.color_palette[class_id]
  138. # 在图像上绘制边界框
  139. cv2.rectangle(img, (int(x1), int(y1)), (int(x1 + w), int(y1 + h)), color, 2)
  140. # 创建包含类别名和分数的标签文本
  141. label = f"{self.classes[class_id]}: {score:.2f}"
  142. # 计算标签文本的尺寸
  143. (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
  144. # 计算标签文本的位置
  145. label_x = x1
  146. label_y = y1 - 10 if y1 - 10 > label_height else y1 + 10
  147. # 绘制填充的矩形作为标签文本的背景
  148. cv2.rectangle(img, (label_x, label_y - label_height), (label_x + label_width, label_y + label_height), color, cv2.FILLED)
  149. # 在图像上绘制标签文本
  150. cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
  151. def main(self):
  152. # 使用 ONNX 模型创建推理会话,自动选择CPU或GPU
  153. session = ort.InferenceSession(
  154. self.onnx_model,
  155. providers=["CUDAExecutionProvider", "CPUExecutionProvider"] if ort.get_device() == "GPU" else ["CPUExecutionProvider"],
  156. )
  157. # 打印模型的输入尺寸
  158. print("YOLO11 ? 目标检测 ONNXRuntime")
  159. print("模型名称:", self.onnx_model)
  160. # 获取模型的输入形状
  161. model_inputs = session.get_inputs()
  162. input_shape = model_inputs[0].shape
  163. self.input_width = input_shape[2]
  164. self.input_height = input_shape[3]
  165. print(f"模型输入尺寸:宽度 = {self.input_width}, 高度 = {self.input_height}")
  166. # 预处理图像数据,确保使用模型要求的尺寸 (640x640)
  167. img_data = self.preprocess()
  168. # 使用预处理后的图像数据运行推理
  169. outputs = session.run(None, {model_inputs[0].name: img_data})
  170. # 对输出进行后处理以获取输出图像
  171. return self.postprocess(self.img, outputs) # 输出图像
  172. if __name__ == "__main__":
  173. # 创建参数解析器以处理命令行参数
  174. parser = argparse.ArgumentParser()
  175. parser.add_argument("--model", type=str, default="runs/detect/train/weights/best.onnx", help="输入你的 ONNX 模型路径。")
  176. parser.add_argument("--img", type=str, default=r"datasets/test.jpg", help="输入图像的路径。")
  177. parser.add_argument("--conf-thres", type=float, default=0.5, help="置信度阈值")
  178. parser.add_argument("--iou-thres", type=float, default=0.45, help="NMS IoU 阈值")
  179. args = parser.parse_args()
  180. # 使用指定的参数创建 YOLO11 类的实例
  181. detection = YOLO11(args.model, args.img, args.conf_thres, args.iou_thres)
  182. # 执行目标检测并获取输出图像
  183. output_image = detection.main()
  184. # 保存输出图像到文件
  185. cv2.imwrite("det_result_picture.jpg", output_image)
  186. print("图像已保存为 det_result_picture.jpg")

需要修改类别映射关系,以及类别对应的颜色

比如,定义两个类别(car、person), 示例代码:

  1. # 类外定义类别映射关系,使用字典格式
  2. CLASS_NAMES = {
  3. 0: 'car', # 类别 0 名称
  4. 1: 'person' # 类别 1 名称
  5. # 可以添加更多类别...
  6. }

运行代码,打印信息:

YOLO11 ? 目标检测 ONNXRuntime
模型名称: runs/detect/train/weights/best.onnx
模型输入尺寸:宽度 = 640, 高度 = 640
Original image shape: (398, 700)
Final letterboxed image shape: (640, 640, 3)
图像已保存为 det_result_picture.jpg

可视化看一下分割效果,保存名称是:det_result_picture.jpg

  YOLO11相关文章推荐:

一篇文章快速认识YOLO11 | 关键改进点 | 安装使用 | 模型训练和推理-CSDN博客

一篇文章快速认识 YOLO11 | 实例分割 | 模型训练 | 自定义数据集-CSDN博客

YOLO11模型推理 | 目标检测与跟踪 | 实例分割 | 关键点估计 | OBB旋转目标检测-CSDN博客

YOLO11模型训练 | 目标检测与跟踪 | 实例分割 | 关键点姿态估计-CSDN博客

YOLO11 实例分割 | 自动标注 | 预标注 | 标签格式转换 | 手动校正标签-CSDN博客

YOLO11 实例分割 | 导出ONNX模型 | ONNX模型推理-CSDN博客

分享完成,欢迎大家多多点赞和收藏,谢谢~

注:本文转载自blog.csdn.net的一颗小树x的文章"https://blog.csdn.net/qq_41204464/article/details/142942825"。版权归原作者所有,此博客不拥有其著作权,亦不承担相应法律责任。如有侵权,请联系我们删除。
复制链接
复制链接
相关推荐
发表评论
登录后才能发表评论和回复 注册

/ 登录

评论记录:

未查询到任何数据!
回复评论:

分类栏目

后端 (14832) 前端 (14280) 移动开发 (3760) 编程语言 (3851) Java (3904) Python (3298) 人工智能 (10119) AIGC (2810) 大数据 (3499) 数据库 (3945) 数据结构与算法 (3757) 音视频 (2669) 云原生 (3145) 云平台 (2965) 前沿技术 (2993) 开源 (2160) 小程序 (2860) 运维 (2533) 服务器 (2698) 操作系统 (2325) 硬件开发 (2492) 嵌入式 (2955) 微软技术 (2769) 软件工程 (2056) 测试 (2865) 网络空间安全 (2948) 网络与通信 (2797) 用户体验设计 (2592) 学习和成长 (2593) 搜索 (2744) 开发工具 (7108) 游戏 (2829) HarmonyOS (2935) 区块链 (2782) 数学 (3112) 3C硬件 (2759) 资讯 (2909) Android (4709) iOS (1850) 代码人生 (3043) 阅读 (2841)

热门文章

101
推荐
关于我们 隐私政策 免责声明 联系我们
Copyright © 2020-2025 蚁人论坛 (iYenn.com) All Rights Reserved.
Scroll to Top