当前位置：首页 > news >正文

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

news 2024/11/23 19:22:55

本文分享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模型（简洁版）

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

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

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

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

YOLO11导出ONNX模型（专业版）


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

对于model.export( )函数中，各种参数说明：

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

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

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

支持的导出格式	`format`参数值	生成的模型示例	model.export( )函数的参数
PyTorch	-	`yolo11n.pt`	-
TorchScript	`torchscript`	`yolo11n.torchscript`	`imgsz`, `optimize`, `batch`
ONNX	`onnx`	`yolo11n.onnx`	`imgsz`, `half`, `dynamic`, `simplify`, `opset`, `batch`
OpenVINO	`openvino`	`yolo11n_openvino_model/`	`imgsz`, `half`, `int8`, `batch`
TensorRT	`engine`	`yolo11n.engine`	`imgsz`, `half`, `dynamic`, `simplify`, `workspace`, `int8`, `batch`
CoreML	`coreml`	`yolo11n.mlpackage`	`imgsz`, `half`, `int8`, `nms`, `batch`
TF SavedModel	`saved_model`	`yolo11n_saved_model/`	`imgsz`, `keras`, `int8`, `batch`
TF GraphDef	`pb`	`yolo11n.pb`	`imgsz`, `batch`
TF Lite	`tflite`	`yolo11n.tflite`	`imgsz`, `half`, `int8`, `batch`
TF Edge TPU	`edgetpu`	`yolo11n_edgetpu.tflite`	`imgsz`
TF.js	`tfjs`	`yolo11n_web_model/`	`imgsz`, `half`, `int8`, `batch`
PaddlePaddle	`paddle`	`yolo11n_paddle_model/`	`imgsz`, `batch`
NCNN	`ncnn`	`yolo11n_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 提供了不同的选项，以适应不同的应用场景：比如是否需要按指定步幅填充，是否强制拉伸以填满，或者是否允许图像放大。

def letterbox(img, new_shape=(640, 640), color=(0, 0, 0), auto=False, scale_fill=False, scale_up=False, stride=32):"""将图像调整为指定尺寸，同时保持长宽比，添加填充以适应目标输入形状。"""# 获取图像的当前高度和宽度shape = img.shape[:2]# 如果 new_shape 是整数，则将其转换为 (new_shape, new_shape) 的格式if isinstance(new_shape, int):new_shape = (new_shape, new_shape)# 计算缩放比例 r，以便图像适配到 new_shape，保持长宽比r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])# 若不允许放大(scale_up=False)，限制 r 最大值为 1.0if not scale_up:r = min(r, 1.0)# 保存缩放比例 ratio，计算未填充的新尺寸 new_unpadratio = r, rnew_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))# 计算目标尺寸与缩放后尺寸的宽度、高度差值 dw, dhdw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]# 若 auto=True，则将 dw 和 dh 调整为 stride 的倍数if auto:dw, dh = np.mod(dw, stride), np.mod(dh, stride)# 若 scale_fill=True，则忽略比例，强制缩放到目标尺寸elif scale_fill:dw, dh = 0.0, 0.0new_unpad = (new_shape[1], new_shape[0])ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]# 将填充值平分到图像四周dw /= 2  dh /= 2# 如果当前图像尺寸与目标尺寸不一致，则将图像缩放到 new_unpadif shape[::-1] != new_unpad:img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)# 应用上下左右的填充以使图像符合目标尺寸top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))left, right = int(round(dw - 0.1)), int(round(dw + 0.1))img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)# 返回填充后的图像、缩放比例和填充量return img, ratio, (dw, dh)

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

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

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

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

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

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

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

总结：

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

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

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

NMS 的目标是去除冗余的边界框，只保留最有代表性的那个，以减少重叠检测。
在这个实现中，通过 ProbIoU 来计算旋转边界框之间的相似度，用于确定哪些框需要保留。
NMS 的核心在于选取当前得分最高的边界框，将它加入保留列表中，
然后与剩余边界框进行相似性计算，通过ProbIoU计算旋转边界框之间的相似度，最终剔除那些相似度高于阈值的边界框。

def rotated_nms_with_probiou(boxes, scores, iou_threshold=0.5):"""使用 ProbIoU 执行旋转边界框的非极大值抑制（NMS）。:param boxes: 旋转边界框的集合:param scores: 每个边界框的置信度得分:param iou_threshold: IoU 阈值，用于确定是否抑制框:return: 保留的边界框索引列表"""order = scores.argsort()[::-1]  # 根据置信度得分降序排序keep = []  # 用于存储保留的边界框索引while len(order) > 0:i = order[0]  # 选择当前得分最高的边界框keep.append(i)  # 将该边界框的索引加入到保留列表中if len(order) == 1:  # 如果只剩下一个边界框，跳出循环breakremaining_boxes = boxes[order[1:]]  # 获取剩余的边界框iou_values = batch_probiou(boxes[i:i + 1], remaining_boxes).squeeze(0)  # 计算当前框与剩余框之间的 ProbIoUmask = iou_values < iou_threshold  # 找出与当前框 IoU 小于阈值的边界框order = order[1:][mask]  # 更新剩余的边界框索引，只保留 IoU 小于阈值的部分return keep  # 返回保留的边界框索引列表

7、解析ONNX模型输出函数

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

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

这里我们需要知道：

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

def parse_onnx_output(output, ratio, dwdh, conf_threshold=0.5, iou_threshold=0.5):"""解析ONNX模型的输出，提取旋转边界框坐标、置信度和类别信息，并应用旋转NMS。:param output: ONNX模型的输出，包含预测的边界框信息:param ratio: 缩放比例，用于将坐标还原到原始尺度:param dwdh: 填充的宽高，用于调整边界框的中心点坐标:param conf_threshold: 置信度阈值，过滤低于该阈值的检测框:param iou_threshold: IoU 阈值，用于旋转边界框的非极大值抑制（NMS）:return: 符合条件的旋转边界框的检测结果"""boxes, scores, classes, detections = [], [], [], []  # 初始化存储检测结果的列表num_detections = output.shape[2]  # 获取检测的边界框数量num_classes = output.shape[1] - 6  # 计算类别数量# 逐个解析每个检测结果for i in range(num_detections):detection = output[0, :, i]  # 获取第 i 个检测结果x_center, y_center, width, height = detection[0], detection[1], detection[2], detection[3]  # 提取边界框的中心坐标和宽高angle = detection[-1]  # 提取旋转角度# 处理类别信息和置信度if num_classes > 0:class_confidences = detection[4:4 + num_classes]  # 获取类别置信度if class_confidences.size == 0:continueclass_id = np.argmax(class_confidences)  # 获取置信度最高的类别索引confidence = class_confidences[class_id]  # 获取对应的置信度else:confidence = detection[4]  # 如果没有类别信息，直接使用置信度值class_id = 0  # 默认类别为 0# 过滤低置信度的检测结果if confidence > conf_threshold:x_center = (x_center - dwdh[0]) / ratio[0]  # 还原中心点 x 坐标y_center = (y_center - dwdh[1]) / ratio[1]  # 还原中心点 y 坐标width /= ratio[0]  # 还原宽度height /= ratio[1]  # 还原高度boxes.append([x_center, y_center, width, height, angle])  # 将边界框信息加入列表scores.append(confidence)  # 将置信度加入列表classes.append(class_id)  # 将类别加入列表if not boxes:  # 如果没有符合条件的边界框，返回空列表return []# 转换为 NumPy 数组boxes = np.array(boxes)scores = np.array(scores)classes = np.array(classes)# 应用旋转 NMSkeep_indices = rotated_nms_with_probiou(boxes, scores, iou_threshold=iou_threshold)# 构建最终检测结果for idx in keep_indices:x_center, y_center, width, height, angle = boxes[idx]  # 获取保留的边界框信息confidence = scores[idx]  # 获取对应的置信度class_id = classes[idx]  # 获取类别obb_corners = calculate_obb_corners(x_center, y_center, width, height, angle)  # 计算旋转边界框的四个角点# 将检测结果添加到列表中detections.append({"position": obb_corners,  # 旋转边界框的角点坐标"confidence": float(confidence),  # 置信度"class_id": int(class_id),  # 类别 ID"angle": float(angle)  # 旋转角度})return detections  # 返回检测结果

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

def calculate_obb_corners(x_center, y_center, width, height, angle):"""根据旋转角度计算旋转边界框的四个角点。:param x_center: 边界框中心的 x 坐标:param y_center: 边界框中心的 y 坐标:param width: 边界框的宽度:param height: 边界框的高度:param angle: 旋转角度:return: 旋转边界框的四个角点坐标"""cos_angle = np.cos(angle)  # 计算旋转角度的余弦值sin_angle = np.sin(angle)  # 计算旋转角度的正弦值dx = width / 2  # 计算宽度的一半dy = height / 2  # 计算高度的一半# 计算旋转边界框的四个角点坐标corners = [(int(x_center + cos_angle * dx - sin_angle * dy), int(y_center + sin_angle * dx + cos_angle * dy)),(int(x_center - cos_angle * dx - sin_angle * dy), int(y_center - sin_angle * dx + cos_angle * dy)),(int(x_center - cos_angle * dx + sin_angle * dy), int(y_center - sin_angle * dx - cos_angle * dy)),(int(x_center + cos_angle * dx + sin_angle * dy), int(y_center + sin_angle * dx - cos_angle * dy)),]return corners  # 返回角点坐标

8、完整代码

完整代码，如下所示：

import os
import cv2
import numpy as np
import onnxruntime as ort
import logging"""
YOLO11 旋转目标检测OBB
1、ONNX模型推理、可视化
2、ONNX输出格式: x_center, y_center, width, height, class1_confidence, ..., classN_confidence, angle
3、支持不同尺寸图片输入、支持旋转NMS过滤重复框、支持ProbIoU旋转IOU计算
"""def letterbox(img, new_shape=(640, 640), color=(0, 0, 0), auto=False, scale_fill=False, scale_up=False, stride=32):"""将图像调整为指定尺寸，同时保持长宽比，添加填充以适应目标输入形状。:param img: 输入图像:param new_shape: 目标尺寸:param color: 填充颜色:param auto: 是否自动调整填充为步幅的整数倍:param scale_fill: 是否强制缩放以完全填充目标尺寸:param scale_up: 是否允许放大图像:param stride: 步幅，用于自动调整填充:return: 调整后的图像、缩放比例、填充尺寸(dw, dh)"""shape = img.shape[:2]if isinstance(new_shape, int):new_shape = (new_shape, new_shape)r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])  # 计算缩放比例if not scale_up:r = min(r, 1.0)ratio = r, rnew_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]if auto:dw, dh = np.mod(dw, stride), np.mod(dh, stride)elif scale_fill:dw, dh = 0.0, 0.0new_unpad = (new_shape[1], new_shape[0])ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]dw /= 2  # 填充均分dh /= 2if shape[::-1] != new_unpad:img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))left, right = int(round(dw - 0.1)), int(round(dw + 0.1))img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)return img, ratio, (dw, dh)def load_model(weights):"""加载ONNX模型并返回会话对象。:param weights: 模型权重文件路径:return: ONNX运行会话对象"""session = ort.InferenceSession(weights, providers=['CPUExecutionProvider'])logging.info(f"模型加载成功: {weights}")return sessiondef _get_covariance_matrix(obb):"""计算旋转边界框的协方差矩阵。:param obb: 旋转边界框 (Oriented Bounding Box)，包含中心坐标、宽、高和旋转角度:return: 协方差矩阵的三个元素 a, b, c"""widths = obb[..., 2] / 2heights = obb[..., 3] / 2angles = obb[..., 4]cos_angle = np.cos(angles)sin_angle = np.sin(angles)a = (widths * cos_angle)**2 + (heights * sin_angle)**2b = (widths * sin_angle)**2 + (heights * cos_angle)**2c = widths * cos_angle * heights * sin_anglereturn a, b, cdef batch_probiou(obb1, obb2, eps=1e-7):"""计算旋转边界框之间的 ProbIoU。:param obb1: 第一个旋转边界框集合:param obb2: 第二个旋转边界框集合:param eps: 防止除零的极小值:return: 两个旋转边界框之间的 ProbIoU"""x1, y1 = obb1[..., 0], obb1[..., 1]x2, y2 = obb2[..., 0], obb2[..., 1]a1, b1, c1 = _get_covariance_matrix(obb1)a2, b2, c2 = _get_covariance_matrix(obb2)t1 = ((a1[:, None] + a2) * (y1[:, None] - y2)**2 + (b1[:, None] + b2) * (x1[:, None] - x2)**2) / ((a1[:, None] + a2) * (b1[:, None] + b2) - (c1[:, None] + c2)**2 + eps) * 0.25t2 = ((c1[:, None] + c2) * (x2 - x1[:, None]) * (y1[:, None] - y2)) / ((a1[:, None] + a2) * (b1[:, None] + b2) - (c1[:, None] + c2)**2 + eps) * 0.5t3 = np.log(((a1[:, None] + a2) * (b1[:, None] + b2) - (c1[:, None] + c2)**2) /(4 * np.sqrt((a1 * b1 - c1**2)[:, None] * (a2 * b2 - c2**2)) + eps) + eps) * 0.5bd = np.clip(t1 + t2 + t3, eps, 100.0)hd = np.sqrt(1.0 - np.exp(-bd) + eps)return 1 - hddef rotated_nms_with_probiou(boxes, scores, iou_threshold=0.5):"""使用 ProbIoU 执行旋转边界框的非极大值抑制（NMS）。:param boxes: 旋转边界框的集合:param scores: 每个边界框的置信度得分:param iou_threshold: IoU 阈值，用于确定是否抑制框:return: 保留的边界框索引列表"""order = scores.argsort()[::-1]  # 根据置信度得分降序排序keep = []while len(order) > 0:i = order[0]keep.append(i)if len(order) == 1:breakremaining_boxes = boxes[order[1:]]iou_values = batch_probiou(boxes[i:i+1], remaining_boxes).squeeze(0)mask = iou_values < iou_threshold  # 保留 IoU 小于阈值的框order = order[1:][mask]return keepdef run_inference(session, image_bytes, imgsz=(640, 640)):"""对输入图像进行预处理，然后使用ONNX模型执行推理。:param session: ONNX运行会话对象:param image_bytes: 输入图像的字节数据:param imgsz: 模型输入的尺寸:return: 推理结果、缩放比例、填充尺寸"""im0 = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_COLOR)  # 解码图像字节数据if im0 is None:raise ValueError("无法从image_bytes解码图像")img, ratio, (dw, dh) = letterbox(im0, new_shape=imgsz)  # 调整图像尺寸img = img.transpose((2, 0, 1))[::-1]  # 调整通道顺序img = np.ascontiguousarray(img)img = img[np.newaxis, ...].astype(np.float32) / 255.0  # 归一化处理input_name = session.get_inputs()[0].nameresult = session.run(None, {input_name: img})  # 执行模型推理return result[0], ratio, (dw, dh)def parse_onnx_output(output, ratio, dwdh, conf_threshold=0.5, iou_threshold=0.5):"""解析ONNX模型的输出，提取旋转边界框坐标、置信度和类别信息，并应用旋转NMS。:param output: ONNX模型的输出，包含预测的边界框信息:param ratio: 缩放比例，用于将坐标还原到原始尺度:param dwdh: 填充的宽高，用于调整边界框的中心点坐标:param conf_threshold: 置信度阈值，过滤低于该阈值的检测框:param iou_threshold: IoU 阈值，用于旋转边界框的非极大值抑制（NMS）:return: 符合条件的旋转边界框的检测结果"""boxes, scores, classes, detections = [], [], [], []num_detections = output.shape[2]  # 获取检测的边界框数量num_classes = output.shape[1] - 6  # 计算类别数量# 逐个解析每个检测结果for i in range(num_detections):detection = output[0, :, i]x_center, y_center, width, height = detection[0], detection[1], detection[2], detection[3]  # 提取边界框的中心坐标和宽高angle = detection[-1]  # 提取旋转角度if num_classes > 0:class_confidences = detection[4:4 + num_classes]  # 获取类别置信度if class_confidences.size == 0:continueclass_id = np.argmax(class_confidences)  # 获取置信度最高的类别索引confidence = class_confidences[class_id]  # 获取对应的置信度else:confidence = detection[4]  # 如果没有类别信息，直接使用置信度值class_id = 0  # 默认类别为 0if confidence > conf_threshold:  # 过滤掉低置信度的检测结果x_center = (x_center - dwdh[0]) / ratio[0]  # 还原中心点 x 坐标y_center = (y_center - dwdh[1]) / ratio[1]  # 还原中心点 y 坐标width /= ratio[0]  # 还原宽度height /= ratio[1]  # 还原高度boxes.append([x_center, y_center, width, height, angle])  # 将边界框信息加入列表scores.append(confidence)  # 将置信度加入列表classes.append(class_id)  # 将类别加入列表if not boxes:return []# 转换为 NumPy 数组boxes = np.array(boxes)scores = np.array(scores)classes = np.array(classes)# 应用旋转 NMSkeep_indices = rotated_nms_with_probiou(boxes, scores, iou_threshold=iou_threshold)# 构建最终检测结果for idx in keep_indices:x_center, y_center, width, height, angle = boxes[idx]  # 获取保留的边界框信息confidence = scores[idx]  # 获取对应的置信度class_id = classes[idx]  # 获取类别obb_corners = calculate_obb_corners(x_center, y_center, width, height, angle)  # 计算旋转边界框的四个角点detections.append({"position": obb_corners,  # 旋转边界框的角点坐标"confidence": float(confidence),  # 置信度"class_id": int(class_id),  # 类别 ID"angle": float(angle)  # 旋转角度})return detectionsdef calculate_obb_corners(x_center, y_center, width, height, angle):"""根据旋转角度计算旋转边界框的四个角点。:param x_center: 边界框中心的 x 坐标:param y_center: 边界框中心的 y 坐标:param width: 边界框的宽度:param height: 边界框的高度:param angle: 旋转角度:return: 旋转边界框的四个角点坐标"""cos_angle = np.cos(angle)  # 计算旋转角度的余弦值sin_angle = np.sin(angle)  # 计算旋转角度的正弦值dx = width / 2  # 计算宽度的一半dy = height / 2  # 计算高度的一半# 计算旋转边界框的四个角点坐标corners = [(int(x_center + cos_angle * dx - sin_angle * dy), int(y_center + sin_angle * dx + cos_angle * dy)),(int(x_center - cos_angle * dx - sin_angle * dy), int(y_center - sin_angle * dx + cos_angle * dy)),(int(x_center - cos_angle * dx + sin_angle * dy), int(y_center - sin_angle * dx - cos_angle * dy)),(int(x_center + cos_angle * dx + sin_angle * dy), int(y_center + sin_angle * dx - cos_angle * dy)),]return corners  # 返回角点坐标def save_detections(image, detections, output_path):"""在图像上绘制旋转边界框检测结果并保存。:param image: 原始图像:param detections: 检测结果列表:param output_path: 保存路径"""for det in detections:corners = det['position']  # 获取旋转边界框的四个角点confidence = det['confidence']  # 获取置信度class_id = det['class_id']  # 获取类别ID# 绘制边界框的四条边for j in range(4):pt1 = corners[j]pt2 = corners[(j + 1) % 4]cv2.line(image, pt1, pt2, (0, 0, 255), 2)# 在边界框上方显示类别和置信度cv2.putText(image, f'Class: {class_id}, Conf: {confidence:.2f}', (corners[0][0], corners[0][1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255), 3)cv2.imwrite(output_path, image)  # 保存绘制后的图像def process_images_in_folder(folder_path, model_weights, output_folder, conf_threshold, iou_threshold, imgsz):"""批量处理文件夹中的图像，执行推理、解析和可视化，保存结果。:param folder_path: 输入图像文件夹路径:param model_weights: ONNX模型权重文件路径:param output_folder: 输出结果文件夹路径:param conf_threshold: 置信度阈值:param iou_threshold: IoU 阈值，用于旋转NMS:param imgsz: 模型输入大小"""session = load_model(weights=model_weights)  # 加载ONNX模型if not os.path.exists(output_folder):os.makedirs(output_folder)  # 如果输出文件夹不存在，则创建for filename in os.listdir(folder_path):if filename.endswith(('.jpg', '.png', '.jpeg')):  # 处理图片文件image_path = os.path.join(folder_path, filename)with open(image_path, 'rb') as f:image_bytes = f.read()print("image_path:", image_path)raw_output, ratio, dwdh = run_inference(session=session, image_bytes=image_bytes, imgsz=imgsz)  # 执行推理detections = parse_onnx_output(raw_output, ratio, dwdh, conf_threshold=conf_threshold, iou_threshold=iou_threshold)  # 解析输出im0 = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_COLOR)  # 解码图像output_path = os.path.join(output_folder, filename)save_detections(im0, detections, output_path)  # 保存检测结果# 主函数：加载参数
if __name__ == "__main__":folder_path = r"point_offer_20240930_rgb"  # 输入图像文件夹路径model_weights = r"YOLO11_obb_exp39_cpu.onnx"  # ONNX模型路径output_folder = "results"  # 输出结果文件夹conf_threshold = 0.5  # 置信度阈值iou_threshold = 0.5  # IoU阈值，用于旋转NMSimgsz = (640, 640)  # 模型输入大小process_images_in_folder(folder_path, model_weights, output_folder, conf_threshold, iou_threshold, imgsz)  # 执行批量处理

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