YOLO11 旋转目标检测 | OBB定向检测 | ONNX模型推理 | 旋转NMS

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

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

本文提供源代码，支持不同尺寸图片输入、支持旋转NMS过滤重复框、支持旋转IOU计算。

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

目录

1、导出ONNX模型

2、所需依赖库

3、整体框架思路

4、支持不同尺寸图像输入 函数

5、旋转边界框IoU计算 函数

6、旋转NMS过滤重复框 函数

7、解析ONNX模型输出 函数

8、完整代码

1、导出ONNX模型

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

然后用下面代码，导出ONNX模型（简洁版）

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

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

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

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

YOLO11导出ONNX模型（专业版）

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

对于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、所需依赖库

本文的代码中，主要依赖opencv、onnxruntime、numpy这三个库；

不需要安装torch、ultralytics等库的。

import os

import cv2

import numpy as np

import onnxruntime as ort

import logging

3、整体框架思路

我们编写代码，实现了一个基于 YOLO11 旋转目标检测（OBB）的推理和检测可视化系统。

以下是代码的整体思路分析：

3.1、基本功能与目标

  • YOLO11模型推理：使用ONNX格式的YOLO11模型，对图像中的旋转目标进行检测。
  • 输出解析：解析模型输出，获取检测的旋转边界框坐标、类别和置信度。
  • 旋转边界框处理：支持旋转NMS（非极大值抑制）和 ProbIoU（概率交并比）来处理重复检测框。
  • 可视化检测结果：在图像上绘制旋转边界框，标注检测结果。

3.2、图像预处理 (letterbox 函数)

  • 将输入图像调整为指定的 imgsz 大小（默认为 640x640），保持长宽比并添加填充。
  • 返回图像缩放的 ratio 和填充偏移 dw, dh，以便后续解析输出时恢复原图坐标。

3.3、加载模型 (load_model 函数)

  • 加载 ONNX 格式的 YOLO11 模型，使用 onnxruntime 进行推理。

3.4、旋转边界框的协方差矩阵计算 (_get_covariance_matrix 函数)

  • 基于边界框的宽、高和角度，计算协方差矩阵的元素 a, b, c，这是 ProbIoU 计算的前提。

3.5、旋转边界框的ProbIoU计算 (batch_probiou 函数)

  • 基于两个旋转边界框集合，使用协方差矩阵计算 ProbIoU 值。ProbIoU 计算边界框之间的相似性，考虑了旋转角度的影响。
  • 返回一个矩阵，表示每对边界框的 ProbIoU 值。

3.6、旋转NMS过滤重复框 (rotated_nms_with_probiou 函数)

  • 使用 ProbIoU 执行旋转边界框的 NMS 操作，去除重叠度过高的检测框。
  • 根据置信度分数 scores 降序排列，逐步计算当前检测框与其余框的 ProbIoU，如果 ProbIoU 小于设定的阈值 iou_threshold，则保留当前框。

3.7、推理处理 (run_inference 函数)

  • 从输入字节流中解码图像数据。
  • 将图像经过 letterbox 预处理后转换为模型输入格式。
  • 使用 ONNX 模型进行推理，返回推理结果 result 及缩放 ratio 和填充 dwdh。

3.8、解析模型输出 (parse_onnx_output 函数)

  • 提取每个检测框的中心坐标、宽高、类别置信度及旋转角度。
  • 根据设定的 conf_threshold 过滤置信度低的检测框。
  • 多类别模型中，提取所有类别的置信度，并选择置信度最高的类别。
  • 使用 rotated_nms_with_probiou 函数执行旋转NMS去除重复框。
  • 将检测框坐标按比例缩放并去除填充，返回处理后的检测结果，包括位置坐标、类别和置信度。

3.9、旋转边界框四角点计算 (calculate_obb_corners 函数)

  • 根据旋转角度和宽高，计算旋转边界框的四个角点坐标，以便在图像上进行绘制。

3.10、绘制检测结果 (save_detections 函数)

  • 在原始图像上绘制旋转边界框和类别标签，标注置信度。
  • 将绘制完成的图像保存到指定的输出路径。

3.11、批量处理文件夹图像 (process_images_in_folder 函数)

  • 从指定文件夹中逐张读取图像文件，对每张图像执行推理、解析输出、绘制结果，并将绘制后的图像保存到输出文件夹。

3.12、主函数执行参数

  • 定义输入图像文件夹、模型路径、输出文件夹、置信度阈值、IoU阈值及模型输入大小 imgsz 等参数。
  • 执行 process_images_in_folder 对文件夹内的图像批量处理并保存结果。

4、支持不同尺寸图像输入 函数

letterbox 函数的主要任务是将输入图像调整为特定尺寸，使用模型所需的输入大小（如 640x640），同时保持图像的长宽比，并为目标尺寸添加适量的填充。

  • 它通过保持原始宽高比来避免图像的变形，通过对称填充使得图像保持中心化
  • 最终得到标准尺寸的图像，适应深度学习模型的输入需求。
  • 参数 auto、scale_fill 和 scale_up 提供了不同的选项，以适应不同的应用场景：比如是否需要按指定步幅填充，是否强制拉伸以填满，或者是否允许图像放大。
  1. def letterbox(img, new_shape=(640, 640), color=(0, 0, 0), auto=False, scale_fill=False, scale_up=False, stride=32):
  2.     """
  3.     将图像调整为指定尺寸，同时保持长宽比，添加填充以适应目标输入形状。
  4.     """
  5.     # 获取图像的当前高度和宽度
  6.     shape = img.shape[:2]
  7.     
  8.     # 如果 new_shape 是整数，则将其转换为 (new_shape, new_shape) 的格式
  9.     if isinstance(new_shape, int):
  10.         new_shape = (new_shape, new_shape)
  11.  
  12.     # 计算缩放比例 r，以便图像适配到 new_shape，保持长宽比
  13.     r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
  14.     
  15.     # 若不允许放大(scale_up=False)，限制 r 最大值为 1.0
  16.     if not scale_up:
  17.         r = min(r, 1.0)
  18.  
  19.     # 保存缩放比例 ratio，计算未填充的新尺寸 new_unpad
  20.     ratio = r, r
  21.     new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
  22.     
  23.     # 计算目标尺寸与缩放后尺寸的宽度、高度差值 dw, dh
  24.     dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]
  25.  
  26.     # 若 auto=True，则将 dw 和 dh 调整为 stride 的倍数
  27.     if auto:
  28.         dw, dh = np.mod(dw, stride), np.mod(dh, stride)
  29.     # 若 scale_fill=True，则忽略比例，强制缩放到目标尺寸
  30.     elif scale_fill:
  31.         dw, dh = 0.0, 0.0
  32.         new_unpad = (new_shape[1], new_shape[0])
  33.         ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]
  34.  
  35.     # 将填充值平分到图像四周
  36.     dw /= 2  
  37.     dh /= 2
  38.  
  39.     # 如果当前图像尺寸与目标尺寸不一致，则将图像缩放到 new_unpad
  40.     if shape[::-1] != new_unpad:
  41.         img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
  42.  
  43.     # 应用上下左右的填充以使图像符合目标尺寸
  44.     top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
  45.     left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
  46.     img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
  47.     
  48.     # 返回填充后的图像、缩放比例和填充量
  49.     return img, ratio, (dw, dh)

5、旋转边界框IoU计算 函数

这里编写batch_probiou 函数，用于旋转边界框IoU计算，其中使用到ProbIoU方法。

这是一个基于协方差矩阵的方法，用于评估旋转边界框之间的相似性。

  1. def batch_probiou(obb1, obb2, eps=1e-7):
  2.     """
  3.     计算旋转边界框之间的 ProbIoU。
  4.     :param obb1: 第一个旋转边界框集合
  5.     :param obb2: 第二个旋转边界框集合
  6.     :param eps: 防止除零的极小值
  7.     :return: 两个旋转边界框之间的 ProbIoU
  8.     """
  9.     # 提取两个旋转边界框的中心坐标 (x, y)
  10.     x1, y1 = obb1[..., 0], obb1[..., 1]
  11.     x2, y2 = obb2[..., 0], obb2[..., 1]
  12.  
  13.     # 计算两个旋转边界框的协方差矩阵元素 a, b, c
  14.     a1, b1, c1 = _get_covariance_matrix(obb1)
  15.     a2, b2, c2 = _get_covariance_matrix(obb2)
  16.  
  17.     # 计算 ProbIoU 的三个部分 t1, t2, t3
  18.     # t1 表示中心点位置差异的贡献
  19.     t1 = ((a1[:, None] + a2) * (y1[:, None] - y2) ** 2 + (b1[:, None] + b2) * (x1[:, None] - x2) ** 2) / (
  20.             (a1[:, None] + a2) * (b1[:, None] + b2) - (c1[:, None] + c2) ** 2 + eps) * 0.25
  21.     # t2 表示旋转角度的耦合贡献
  22.     t2 = ((c1[:, None] + c2) * (x2 - x1[:, None]) * (y1[:, None] - y2)) / (
  23.             (a1[:, None] + a2) * (b1[:, None] + b2) - (c1[:, None] + c2) ** 2 + eps) * 0.5
  24.     # t3 表示面积和形状之间的差异贡献
  25.     t3 = np.log(((a1[:, None] + a2) * (b1[:, None] + b2) - (c1[:, None] + c2) ** 2) /
  26.                 (4 * np.sqrt((a1 * b1 - c1 ** 2)[:, None] * (a2 * b2 - c2 ** 2)) + eps) + eps) * 0.5
  27.  
  28.     # 计算 Bhattacharyya 距离 bd
  29.     bd = np.clip(t1 + t2 + t3, eps, 100.0)  # 将 bd 限制在 [eps, 100.0] 范围内
  30.  
  31.     # 计算 ProbIoU 值 hd
  32.     hd = np.sqrt(1.0 - np.exp(-bd) + eps)  # 使用 Bhattacharyya 距离计算 hd
  33.     return 1 - hd  # 返回 1 - hd，hd 越小表示相似度越高，1 - hd 即为 ProbIoU

ProbIoU的论文地址：https://arxiv.org/pdf/2106.06072v1.pdf

batch_probiou 函数需要用到协方差矩阵，这里也编写一个函数进行封装

  1. def _get_covariance_matrix(obb):
  2.     """
  3.     计算旋转边界框的协方差矩阵。
  4.     :param obb: 旋转边界框 (Oriented Bounding Box)，包含中心坐标、宽、高和旋转角度
  5.     :return: 协方差矩阵的三个元素 a, b, c
  6.     """
  7.     widths = obb[..., 2] / 2  # 获取宽度的一半
  8.     heights = obb[..., 3] / 2  # 获取高度的一半
  9.     angles = obb[..., 4]  # 获取旋转角度
  10.  
  11.     cos_angle = np.cos(angles)  # 计算旋转角度的余弦值
  12.     sin_angle = np.sin(angles)  # 计算旋转角度的正弦值
  13.  
  14.     # 计算协方差矩阵的三个元素 a, b, c
  15.     a = (widths * cos_angle) ** 2 + (heights * sin_angle) ** 2
  16.     b = (widths * sin_angle) ** 2 + (heights * cos_angle) ** 2
  17.     c = widths * cos_angle * heights * sin_angle
  18.  
  19.     return a, b, c

总结：

  • 该代码用于计算两个旋转边界框之间的 ProbIoU，以量化它们的相似程度。
  • 核心思想是通过协方差矩阵来描述每个旋转边界框的形状和旋转角度，结合 Bhattacharyya 距离来评估边界框之间的相似性。
  • ProbIoU 的计算考虑了旋转角度、边界框的中心位置以及大小差异，是比传统 IoU 更加复杂和准确的度量旋转边界框相似性的方法，特别适用于场景中有方向性的对象。

6、旋转NMS过滤重复框 函数

这里编写rotated_nms_with_probiou函数，实现了旋转边界框的非极大值抑制NMS。

  • NMS 的目标是去除冗余的边界框，只保留最有代表性的那个，以减少重叠检测。
  • 在这个实现中，通过 ProbIoU 来计算旋转边界框之间的相似度，用于确定哪些框需要保留。
  • NMS 的核心在于选取当前得分最高的边界框，将它加入保留列表中，
  • 然后与剩余边界框进行相似性计算，通过ProbIoU计算旋转边界框之间的相似度，最终剔除那些相似度高于阈值的边界框。
  1. def rotated_nms_with_probiou(boxes, scores, iou_threshold=0.5):
  2.     """
  3.     使用 ProbIoU 执行旋转边界框的非极大值抑制（NMS）。
  4.     :param boxes: 旋转边界框的集合
  5.     :param scores: 每个边界框的置信度得分
  6.     :param iou_threshold: IoU 阈值，用于确定是否抑制框
  7.     :return: 保留的边界框索引列表
  8.     """
  9.     order = scores.argsort()[::-1]  # 根据置信度得分降序排序
  10.     keep = []  # 用于存储保留的边界框索引
  11.  
  12.     while len(order) > 0:
  13.         i = order[0]  # 选择当前得分最高的边界框
  14.         keep.append(i)  # 将该边界框的索引加入到保留列表中
  15.  
  16.         if len(order) == 1:  # 如果只剩下一个边界框，跳出循环
  17.             break
  18.  
  19.         remaining_boxes = boxes[order[1:]]  # 获取剩余的边界框
  20.         iou_values = batch_probiou(boxes[i:i + 1], remaining_boxes).squeeze(0)  # 计算当前框与剩余框之间的 ProbIoU
  21.  
  22.         mask = iou_values < iou_threshold  # 找出与当前框 IoU 小于阈值的边界框
  23.         order = order[1:][mask]  # 更新剩余的边界框索引，只保留 IoU 小于阈值的部分
  24.  
  25.     return keep  # 返回保留的边界框索引列表

7、解析ONNX模型输出 函数

这里编写parse_onnx_output函数，实现了对 ONNX 模型的输出进行解析

  • 提取旋转边界框的坐标、旋转角度、置信度和类别信息，
  • 应用旋转边界框的非极大值抑制（NMS），并计算旋转边界框的四个角点

这里我们需要知道：

ONNX输出格式: x_center, y_center, width, height, class1_confidence, ..., classN_confidence, angle 

  1. def parse_onnx_output(output, ratio, dwdh, conf_threshold=0.5, iou_threshold=0.5):
  2.     """
  3.     解析ONNX模型的输出，提取旋转边界框坐标、置信度和类别信息，并应用旋转NMS。
  4.     :param output: ONNX模型的输出，包含预测的边界框信息
  5.     :param ratio: 缩放比例，用于将坐标还原到原始尺度
  6.     :param dwdh: 填充的宽高，用于调整边界框的中心点坐标
  7.     :param conf_threshold: 置信度阈值，过滤低于该阈值的检测框
  8.     :param iou_threshold: IoU 阈值，用于旋转边界框的非极大值抑制（NMS）
  9.     :return: 符合条件的旋转边界框的检测结果
  10.     """
  11.     boxes, scores, classes, detections = [], [], [], []  # 初始化存储检测结果的列表
  12.     num_detections = output.shape[2]  # 获取检测的边界框数量
  13.     num_classes = output.shape[1] - 6  # 计算类别数量
  14.  
  15.     # 逐个解析每个检测结果
  16.     for i in range(num_detections):
  17.         detection = output[0, :, i]  # 获取第 i 个检测结果
  18.         x_center, y_center, width, height = detection[0], detection[1], detection[2], detection[3]  # 提取边界框的中心坐标和宽高
  19.         angle = detection[-1]  # 提取旋转角度
  20.  
  21.         # 处理类别信息和置信度
  22.         if num_classes > 0:
  23.             class_confidences = detection[4:4 + num_classes]  # 获取类别置信度
  24.             if class_confidences.size == 0:
  25.                 continue
  26.             class_id = np.argmax(class_confidences)  # 获取置信度最高的类别索引
  27.             confidence = class_confidences[class_id]  # 获取对应的置信度
  28.         else:
  29.             confidence = detection[4]  # 如果没有类别信息，直接使用置信度值
  30.             class_id = 0  # 默认类别为 0
  31.  
  32.         # 过滤低置信度的检测结果
  33.         if confidence > conf_threshold:
  34.             x_center = (x_center - dwdh[0]) / ratio[0]  # 还原中心点 x 坐标
  35.             y_center = (y_center - dwdh[1]) / ratio[1]  # 还原中心点 y 坐标
  36.             width /= ratio[0]  # 还原宽度
  37.             height /= ratio[1]  # 还原高度
  38.  
  39.             boxes.append([x_center, y_center, width, height, angle])  # 将边界框信息加入列表
  40.             scores.append(confidence)  # 将置信度加入列表
  41.             classes.append(class_id)  # 将类别加入列表
  42.  
  43.     if not boxes:  # 如果没有符合条件的边界框，返回空列表
  44.         return []
  45.  
  46.     # 转换为 NumPy 数组
  47.     boxes = np.array(boxes)
  48.     scores = np.array(scores)
  49.     classes = np.array(classes)
  50.  
  51.     # 应用旋转 NMS
  52.     keep_indices = rotated_nms_with_probiou(boxes, scores, iou_threshold=iou_threshold)
  53.  
  54.     # 构建最终检测结果
  55.     for idx in keep_indices:
  56.         x_center, y_center, width, height, angle = boxes[idx]  # 获取保留的边界框信息
  57.         confidence = scores[idx]  # 获取对应的置信度
  58.         class_id = classes[idx]  # 获取类别
  59.         obb_corners = calculate_obb_corners(x_center, y_center, width, height, angle)  # 计算旋转边界框的四个角点
  60.  
  61.         # 将检测结果添加到列表中
  62.         detections.append({
  63.             "position": obb_corners,  # 旋转边界框的角点坐标
  64.             "confidence": float(confidence),  # 置信度
  65.             "class_id": int(class_id),  # 类别 ID
  66.             "angle": float(angle)  # 旋转角度
  67.         })
  68.  
  69.     return detections  # 返回检测结果
  70.

parse_onnx_output函数需要用到calculate_obb_corners函数，根据旋转角度计算旋转边界框的四个角点

  1. def calculate_obb_corners(x_center, y_center, width, height, angle):
  2.     """
  3.     根据旋转角度计算旋转边界框的四个角点。
  4.     :param x_center: 边界框中心的 x 坐标
  5.     :param y_center: 边界框中心的 y 坐标
  6.     :param width: 边界框的宽度
  7.     :param height: 边界框的高度
  8.     :param angle: 旋转角度
  9.     :return: 旋转边界框的四个角点坐标
  10.     """
  11.     cos_angle = np.cos(angle)  # 计算旋转角度的余弦值
  12.     sin_angle = np.sin(angle)  # 计算旋转角度的正弦值
  13.     dx = width / 2  # 计算宽度的一半
  14.     dy = height / 2  # 计算高度的一半
  15.  
  16.     # 计算旋转边界框的四个角点坐标
  17.     corners = [
  18.         (int(x_center + cos_angle * dx - sin_angle * dy), int(y_center + sin_angle * dx + cos_angle * dy)),
  19.         (int(x_center - cos_angle * dx - sin_angle * dy), int(y_center - sin_angle * dx + cos_angle * dy)),
  20.         (int(x_center - cos_angle * dx + sin_angle * dy), int(y_center - sin_angle * dx - cos_angle * dy)),
  21.         (int(x_center + cos_angle * dx + sin_angle * dy), int(y_center + sin_angle * dx - cos_angle * dy)),
  22.     ]
  23.     return corners  # 返回角点坐标

8、完整代码

完整代码，如下所示：

  1. import os
  2. import cv2
  3. import numpy as np
  4. import onnxruntime as ort
  5. import logging
  6.  
  7. """
  8. YOLO11 旋转目标检测OBB
  9. 1、ONNX模型推理、可视化
  10. 2、ONNX输出格式: x_center, y_center, width, height, class1_confidence, ..., classN_confidence, angle
  11. 3、支持不同尺寸图片输入、支持旋转NMS过滤重复框、支持ProbIoU旋转IOU计算
  12. """
  13.  
  14. def letterbox(img, new_shape=(640, 640), color=(0, 0, 0), auto=False, scale_fill=False, scale_up=False, stride=32):
  15.     """
  16.     将图像调整为指定尺寸，同时保持长宽比，添加填充以适应目标输入形状。
  17.     :param img: 输入图像
  18.     :param new_shape: 目标尺寸
  19.     :param color: 填充颜色
  20.     :param auto: 是否自动调整填充为步幅的整数倍
  21.     :param scale_fill: 是否强制缩放以完全填充目标尺寸
  22.     :param scale_up: 是否允许放大图像
  23.     :param stride: 步幅，用于自动调整填充
  24.     :return: 调整后的图像、缩放比例、填充尺寸(dw, dh)
  25.     """
  26.     shape = img.shape[:2]
  27.     if isinstance(new_shape, int):
  28.         new_shape = (new_shape, new_shape)
  29.  
  30.     r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])  # 计算缩放比例
  31.     if not scale_up:
  32.         r = min(r, 1.0)
  33.  
  34.     ratio = r, r
  35.     new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
  36.     dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]
  37.  
  38.     if auto:
  39.         dw, dh = np.mod(dw, stride), np.mod(dh, stride)
  40.     elif scale_fill:
  41.         dw, dh = 0.0, 0.0
  42.         new_unpad = (new_shape[1], new_shape[0])
  43.         ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]
  44.  
  45.     dw /= 2  # 填充均分
  46.     dh /= 2
  47.  
  48.     if shape[::-1] != new_unpad:
  49.         img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
  50.     top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
  51.     left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
  52.     img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
  53.     return img, ratio, (dw, dh)
  54.  
  55. def load_model(weights):
  56.     """
  57.     加载ONNX模型并返回会话对象。
  58.     :param weights: 模型权重文件路径
  59.     :return: ONNX运行会话对象
  60.     """
  61.     session = ort.InferenceSession(weights, providers=['CPUExecutionProvider'])
  62.     logging.info(f"模型加载成功: {weights}")
  63.     return session
  64.  
  65. def _get_covariance_matrix(obb):
  66.     """
  67.     计算旋转边界框的协方差矩阵。
  68.     :param obb: 旋转边界框 (Oriented Bounding Box)，包含中心坐标、宽、高和旋转角度
  69.     :return: 协方差矩阵的三个元素 a, b, c
  70.     """
  71.     widths = obb[..., 2] / 2
  72.     heights = obb[..., 3] / 2
  73.     angles = obb[..., 4]
  74.  
  75.     cos_angle = np.cos(angles)
  76.     sin_angle = np.sin(angles)
  77.  
  78.     a = (widths * cos_angle)**2 + (heights * sin_angle)**2
  79.     b = (widths * sin_angle)**2 + (heights * cos_angle)**2
  80.     c = widths * cos_angle * heights * sin_angle
  81.  
  82.     return a, b, c
  83.  
  84. def batch_probiou(obb1, obb2, eps=1e-7):
  85.     """
  86.     计算旋转边界框之间的 ProbIoU。
  87.     :param obb1: 第一个旋转边界框集合
  88.     :param obb2: 第二个旋转边界框集合
  89.     :param eps: 防止除零的极小值
  90.     :return: 两个旋转边界框之间的 ProbIoU
  91.     """
  92.     x1, y1 = obb1[..., 0], obb1[..., 1]
  93.     x2, y2 = obb2[..., 0], obb2[..., 1]
  94.     a1, b1, c1 = _get_covariance_matrix(obb1)
  95.     a2, b2, c2 = _get_covariance_matrix(obb2)
  96.  
  97.     t1 = ((a1[:, None] + a2) * (y1[:, None] - y2)**2 + (b1[:, None] + b2) * (x1[:, None] - x2)**2) / (
  98.             (a1[:, None] + a2) * (b1[:, None] + b2) - (c1[:, None] + c2)**2 + eps) * 0.25
  99.     t2 = ((c1[:, None] + c2) * (x2 - x1[:, None]) * (y1[:, None] - y2)) / (
  100.             (a1[:, None] + a2) * (b1[:, None] + b2) - (c1[:, None] + c2)**2 + eps) * 0.5
  101.     t3 = np.log(((a1[:, None] + a2) * (b1[:, None] + b2) - (c1[:, None] + c2)**2) /
  102.                 (4 * np.sqrt((a1 * b1 - c1**2)[:, None] * (a2 * b2 - c2**2)) + eps) + eps) * 0.5
  103.  
  104.     bd = np.clip(t1 + t2 + t3, eps, 100.0)
  105.     hd = np.sqrt(1.0 - np.exp(-bd) + eps)
  106.     return 1 - hd
  107.  
  108. def rotated_nms_with_probiou(boxes, scores, iou_threshold=0.5):
  109.     """
  110.     使用 ProbIoU 执行旋转边界框的非极大值抑制（NMS）。
  111.     :param boxes: 旋转边界框的集合
  112.     :param scores: 每个边界框的置信度得分
  113.     :param iou_threshold: IoU 阈值，用于确定是否抑制框
  114.     :return: 保留的边界框索引列表
  115.     """
  116.     order = scores.argsort()[::-1]  # 根据置信度得分降序排序
  117.     keep = []
  118.  
  119.     while len(order) > 0:
  120.         i = order[0]
  121.         keep.append(i)
  122.  
  123.         if len(order) == 1:
  124.             break
  125.  
  126.         remaining_boxes = boxes[order[1:]]
  127.         iou_values = batch_probiou(boxes[i:i+1], remaining_boxes).squeeze(0)
  128.  
  129.         mask = iou_values < iou_threshold  # 保留 IoU 小于阈值的框
  130.         order = order[1:][mask]
  131.  
  132.     return keep
  133.  
  134. def run_inference(session, image_bytes, imgsz=(640, 640)):
  135.     """
  136.     对输入图像进行预处理，然后使用ONNX模型执行推理。
  137.     :param session: ONNX运行会话对象
  138.     :param image_bytes: 输入图像的字节数据
  139.     :param imgsz: 模型输入的尺寸
  140.     :return: 推理结果、缩放比例、填充尺寸
  141.     """
  142.     im0 = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_COLOR)  # 解码图像字节数据
  143.     if im0 is None:
  144.         raise ValueError("无法从image_bytes解码图像")
  145.  
  146.     img, ratio, (dw, dh) = letterbox(im0, new_shape=imgsz)  # 调整图像尺寸
  147.     img = img.transpose((2, 0, 1))[::-1]  # 调整通道顺序
  148.     img = np.ascontiguousarray(img)
  149.     img = img[np.newaxis, ...].astype(np.float32) / 255.0  # 归一化处理
  150.  
  151.     input_name = session.get_inputs()[0].name
  152.     result = session.run(None, {input_name: img})  # 执行模型推理
  153.  
  154.     return result[0], ratio, (dw, dh)
  155.  
  156. def parse_onnx_output(output, ratio, dwdh, conf_threshold=0.5, iou_threshold=0.5):
  157.     """
  158.     解析ONNX模型的输出，提取旋转边界框坐标、置信度和类别信息，并应用旋转NMS。
  159.     :param output: ONNX模型的输出，包含预测的边界框信息
  160.     :param ratio: 缩放比例，用于将坐标还原到原始尺度
  161.     :param dwdh: 填充的宽高，用于调整边界框的中心点坐标
  162.     :param conf_threshold: 置信度阈值，过滤低于该阈值的检测框
  163.     :param iou_threshold: IoU 阈值，用于旋转边界框的非极大值抑制（NMS）
  164.     :return: 符合条件的旋转边界框的检测结果
  165.     """
  166.     boxes, scores, classes, detections = [], [], [], []
  167.     num_detections = output.shape[2]  # 获取检测的边界框数量
  168.     num_classes = output.shape[1] - 6  # 计算类别数量
  169.  
  170.     # 逐个解析每个检测结果
  171.     for i in range(num_detections):
  172.         detection = output[0, :, i]
  173.         x_center, y_center, width, height = detection[0], detection[1], detection[2], detection[3]  # 提取边界框的中心坐标和宽高
  174.         angle = detection[-1]  # 提取旋转角度
  175.  
  176.         if num_classes > 0:
  177.             class_confidences = detection[4:4 + num_classes]  # 获取类别置信度
  178.             if class_confidences.size == 0:
  179.                 continue
  180.             class_id = np.argmax(class_confidences)  # 获取置信度最高的类别索引
  181.             confidence = class_confidences[class_id]  # 获取对应的置信度
  182.         else:
  183.             confidence = detection[4]  # 如果没有类别信息，直接使用置信度值
  184.             class_id = 0  # 默认类别为 0
  185.  
  186.         if confidence > conf_threshold:  # 过滤掉低置信度的检测结果
  187.             x_center = (x_center - dwdh[0]) / ratio[0]  # 还原中心点 x 坐标
  188.             y_center = (y_center - dwdh[1]) / ratio[1]  # 还原中心点 y 坐标
  189.             width /= ratio[0]  # 还原宽度
  190.             height /= ratio[1]  # 还原高度
  191.  
  192.             boxes.append([x_center, y_center, width, height, angle])  # 将边界框信息加入列表
  193.             scores.append(confidence)  # 将置信度加入列表
  194.             classes.append(class_id)  # 将类别加入列表
  195.  
  196.     if not boxes:
  197.         return []
  198.  
  199.     # 转换为 NumPy 数组
  200.     boxes = np.array(boxes)
  201.     scores = np.array(scores)
  202.     classes = np.array(classes)
  203.  
  204.     # 应用旋转 NMS
  205.     keep_indices = rotated_nms_with_probiou(boxes, scores, iou_threshold=iou_threshold)
  206.  
  207.     # 构建最终检测结果
  208.     for idx in keep_indices:
  209.         x_center, y_center, width, height, angle = boxes[idx]  # 获取保留的边界框信息
  210.         confidence = scores[idx]  # 获取对应的置信度
  211.         class_id = classes[idx]  # 获取类别
  212.         obb_corners = calculate_obb_corners(x_center, y_center, width, height, angle)  # 计算旋转边界框的四个角点
  213.  
  214.         detections.append({
  215.             "position": obb_corners,  # 旋转边界框的角点坐标
  216.             "confidence": float(confidence),  # 置信度
  217.             "class_id": int(class_id),  # 类别 ID
  218.             "angle": float(angle)  # 旋转角度
  219.         })
  220.  
  221.     return detections
  222.  
  223. def calculate_obb_corners(x_center, y_center, width, height, angle):
  224.     """
  225.     根据旋转角度计算旋转边界框的四个角点。
  226.     :param x_center: 边界框中心的 x 坐标
  227.     :param y_center: 边界框中心的 y 坐标
  228.     :param width: 边界框的宽度
  229.     :param height: 边界框的高度
  230.     :param angle: 旋转角度
  231.     :return: 旋转边界框的四个角点坐标
  232.     """
  233.     cos_angle = np.cos(angle)  # 计算旋转角度的余弦值
  234.     sin_angle = np.sin(angle)  # 计算旋转角度的正弦值
  235.     dx = width / 2  # 计算宽度的一半
  236.     dy = height / 2  # 计算高度的一半
  237.  
  238.     # 计算旋转边界框的四个角点坐标
  239.     corners = [
  240.         (int(x_center + cos_angle * dx - sin_angle * dy), int(y_center + sin_angle * dx + cos_angle * dy)),
  241.         (int(x_center - cos_angle * dx - sin_angle * dy), int(y_center - sin_angle * dx + cos_angle * dy)),
  242.         (int(x_center - cos_angle * dx + sin_angle * dy), int(y_center - sin_angle * dx - cos_angle * dy)),
  243.         (int(x_center + cos_angle * dx + sin_angle * dy), int(y_center + sin_angle * dx - cos_angle * dy)),
  244.     ]
  245.     return corners  # 返回角点坐标
  246.  
  247. def save_detections(image, detections, output_path):
  248.     """
  249.     在图像上绘制旋转边界框检测结果并保存。
  250.     :param image: 原始图像
  251.     :param detections: 检测结果列表
  252.     :param output_path: 保存路径
  253.     """
  254.     for det in detections:
  255.         corners = det['position']  # 获取旋转边界框的四个角点
  256.         confidence = det['confidence']  # 获取置信度
  257.         class_id = det['class_id']  # 获取类别ID
  258.  
  259.         # 绘制边界框的四条边
  260.         for j in range(4):
  261.             pt1 = corners[j]
  262.             pt2 = corners[(j + 1) % 4]
  263.             cv2.line(image, pt1, pt2, (0, 0, 255), 2)
  264.         
  265.         # 在边界框上方显示类别和置信度
  266.         cv2.putText(image, f'Class: {class_id}, Conf: {confidence:.2f}', 
  267.                     (corners[0][0], corners[0][1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255), 3)
  268.  
  269.     cv2.imwrite(output_path, image)  # 保存绘制后的图像
  270.  
  271. def process_images_in_folder(folder_path, model_weights, output_folder, conf_threshold, iou_threshold, imgsz):
  272.     """
  273.     批量处理文件夹中的图像，执行推理、解析和可视化，保存结果。
  274.     :param folder_path: 输入图像文件夹路径
  275.     :param model_weights: ONNX模型权重文件路径
  276.     :param output_folder: 输出结果文件夹路径
  277.     :param conf_threshold: 置信度阈值
  278.     :param iou_threshold: IoU 阈值，用于旋转NMS
  279.     :param imgsz: 模型输入大小
  280.     """
  281.     session = load_model(weights=model_weights)  # 加载ONNX模型
  282.     
  283.     if not os.path.exists(output_folder):
  284.         os.makedirs(output_folder)  # 如果输出文件夹不存在，则创建
  285.  
  286.     for filename in os.listdir(folder_path):
  287.         if filename.endswith(('.jpg', '.png', '.jpeg')):  # 处理图片文件
  288.             image_path = os.path.join(folder_path, filename)
  289.             with open(image_path, 'rb') as f:
  290.                 image_bytes = f.read()
  291.             print("image_path:", image_path)
  292.  
  293.             raw_output, ratio, dwdh = run_inference(session=session, image_bytes=image_bytes, imgsz=imgsz)  # 执行推理
  294.             detections = parse_onnx_output(raw_output, ratio, dwdh, conf_threshold=conf_threshold, iou_threshold=iou_threshold)  # 解析输出
  295.  
  296.             im0 = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_COLOR)  # 解码图像
  297.             output_path = os.path.join(output_folder, filename)
  298.             save_detections(im0, detections, output_path)  # 保存检测结果
  299.  
  300. # 主函数：加载参数
  301. if __name__ == "__main__":
  302.     folder_path = r"point_offer_20240930_rgb"  # 输入图像文件夹路径
  303.     model_weights = r"YOLO11_obb_exp39_cpu.onnx"  # ONNX模型路径
  304.     output_folder = "results"  # 输出结果文件夹
  305.     conf_threshold = 0.5  # 置信度阈值
  306.     iou_threshold = 0.5  # IoU阈值，用于旋转NMS
  307.     imgsz = (640, 640)  # 模型输入大小
  308.  
  309.     process_images_in_folder(folder_path, model_weights, output_folder, conf_threshold, iou_threshold, imgsz)  # 执行批量处理

使用这个代码时，需要修改配置参数：

  • folder_path: 输入图像文件夹路径，指向包含待检测图像的目录。

  • model_weights: ONNX 模型文件路径，指向训练好的模型文件。

  • output_folder: 检测结果保存的文件夹路径，输出检测后的图片。

  • conf_threshold: 置信度阈值，用于过滤低置信度的检测框。建议调整以平衡检测精度，默认值为 0.5。

  • iou_threshold: IoU 阈值，用于旋转边界框的非极大值抑制（NMS），默认值为 0.5。较低值减少重叠，较高值保留更多边界框。

  • imgsz: 输入图像的尺寸，例如 (640, 640)。应与模型训练时的输入尺寸一致。

 YOLO11相关文章推荐：

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

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

一篇文章快速认识YOLO11 | 旋转目标检测 | 原理分析 | 模型训练 | 模型推理_yolov11 obb-CSDN博客

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

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

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

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

YOLO11 目标检测 | 导出ONNX模型 | ONNX模型推理-CSDN博客

YOLO11 目标检测 | 自动标注 | 预标注 | 标签格式转换 | 手动校正标签_yolo11 标注平台-CSDN博客

YOLO11 图像缩放 | 图像填充 | 自适应不同尺寸的图片_yolov11 中图像预处理-CSDN博客

YOLO11 旋转目标检测 | 数据标注 | 自定义数据集 | 模型训练 | 模型推理-CSDN博客

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

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

评论记录：

未查询到任何数据！