中能建西北城市建设有限公司网站,it公论 是建立在什么网站,建设局网站查询个人信息,营销型网站建设 多少钱本文分享YOLO11中#xff0c;从xxx.pt权重文件转为.onnx文件#xff0c;然后使用.onnx文件#xff0c;进行旋转目标检测任务的模型推理。
用ONNX模型推理#xff0c;便于算法到开发板或芯片的部署。
本文提供源代码#xff0c;支持不同尺寸图片输入、支持旋转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模型简洁版
from ultralytics import YOLO# 加载一个模型路径为 YOLO 模型的 .pt 文件
model YOLO(runs/obb/train1/weights/best.pt)# 导出模型格式为 ONNX
model.export(formatonnx)运行代码后会在上面路径中生成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(rruns/obb/train1/weights/best.pt)# 导出模型设置多种参数
model.export(formatonnx, # 导出格式为 ONNXimgsz(640, 640), # 设置输入图像的尺寸kerasFalse, # 不导出为 Keras 格式optimizeFalse, # 不进行优化 False, 移动设备优化的参数用于在导出为TorchScript 格式时进行模型优化halfFalse, # 不启用 FP16 量化int8False, # 不启用 INT8 量化dynamicFalse, # 不启用动态输入尺寸simplifyTrue, # 简化 ONNX 模型opsetNone, # 使用最新的 opset 版本workspace4.0, # 为 TensorRT 优化设置最大工作区大小GiBnmsFalse, # 不添加 NMS非极大值抑制batch1, # 指定批处理大小devicecpu # 指定导出设备为CPU或GPU对应参数为cpu , 0
)
对于model.export( )函数中各种参数说明
formatonnx指定导出模型的格式为 onnx。imgsz(640, 640)输入图像的尺寸设为 640x640。如果需要其他尺寸可以修改这个值。kerasFalse不导出为 Keras 格式的模型。optimizeFalse不应用 TorchScript 移动设备优化。halfFalse不启用 FP16半精度量化。int8False不启用 INT8 量化。dynamicFalse不启用动态输入尺寸。simplifyTrue简化模型以提升 ONNX 模型的性能。opsetNone使用默认的 ONNX opset 版本如果需要可以手动指定。workspace4.0为 TensorRT 优化指定最大工作空间大小为 4 GiB。nmsFalse不为 CoreML 导出添加非极大值抑制NMS。batch1设置批处理大小为 1。devicecpu, 指定导出设备为CPU或GPU对应参数为cpu , 0
参考官网文档https://docs.ultralytics.com/modes/export/#arguments
当然了YOLO11中不仅支持ONNX模型还支持下面表格中格式
支持的导出格式format参数值生成的模型示例model.export( )函数的参数PyTorch-yolo11n.pt-TorchScripttorchscriptyolo11n.torchscriptimgsz, optimize, batchONNXonnxyolo11n.onnximgsz, half, dynamic, simplify, opset, batchOpenVINOopenvinoyolo11n_openvino_model/imgsz, half, int8, batchTensorRTengineyolo11n.engineimgsz, half, dynamic, simplify, workspace, int8, batchCoreMLcoremlyolo11n.mlpackageimgsz, half, int8, nms, batchTF SavedModelsaved_modelyolo11n_saved_model/imgsz, keras, int8, batchTF GraphDefpbyolo11n.pbimgsz, batchTF Litetfliteyolo11n.tfliteimgsz, half, int8, batchTF Edge TPUedgetpuyolo11n_edgetpu.tfliteimgszTF.jstfjsyolo11n_web_model/imgsz, half, int8, batchPaddlePaddlepaddleyolo11n_paddle_model/imgsz, batchNCNNncnnyolo11n_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), autoFalse, scale_fillFalse, scale_upFalse, stride32):将图像调整为指定尺寸同时保持长宽比添加填充以适应目标输入形状。# 获取图像的当前高度和宽度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_upFalse)限制 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]# 若 autoTrue则将 dw 和 dh 调整为 stride 的倍数if auto:dw, dh np.mod(dw, stride), np.mod(dh, stride)# 若 scale_fillTrue则忽略比例强制缩放到目标尺寸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, interpolationcv2.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, valuecolor)# 返回填充后的图像、缩放比例和填充量return img, ratio, (dw, dh)5、旋转边界框IoU计算 函数
这里编写batch_probiou 函数用于旋转边界框IoU计算其中使用到ProbIoU方法。
这是一个基于协方差矩阵的方法用于评估旋转边界框之间的相似性。
def batch_probiou(obb1, obb2, eps1e-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 - hdhd 越小表示相似度越高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, cwidths 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_threshold0.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_threshold0.5, iou_threshold0.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_thresholdiou_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), # 类别 IDangle: 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), autoFalse, scale_fillFalse, scale_upFalse, stride32):将图像调整为指定尺寸同时保持长宽比添加填充以适应目标输入形状。: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, interpolationcv2.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, valuecolor)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, cwidths 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, eps1e-7):计算旋转边界框之间的 ProbIoU。:param obb1: 第一个旋转边界框集合:param obb2: 第二个旋转边界框集合:param eps: 防止除零的极小值:return: 两个旋转边界框之间的 ProbIoUx1, 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_threshold0.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:i1], 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_shapeimgsz) # 调整图像尺寸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_threshold0.5, iou_threshold0.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_thresholdiou_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), # 类别 IDangle: 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, fClass: {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(weightsmodel_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(sessionsession, image_bytesimage_bytes, imgszimgsz) # 执行推理detections parse_onnx_output(raw_output, ratio, dwdh, conf_thresholdconf_threshold, iou_thresholdiou_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 rpoint_offer_20240930_rgb # 输入图像文件夹路径model_weights rYOLO11_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) # 执行批量处理使用这个代码时需要修改配置参数 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博客 分享完成欢迎大家多多点赞和收藏谢谢~
