### 使用 ONNX 模型进行 YOLOv11 OBB 推理及结果后处理 YOLOv11 OBB（Oriented Bounding Box）模型是一种改进的目标检测模型，能够预测具有方向信息的边界框。使用 ONNX 模型进行推理时，通常包括以下几个步骤：模型加载、图像预处理、推理执行以及后处理。 #### 模型加载 首先，需要使用 `onnxruntime` 加载 ONNX 模型。可以使用以下代码加载模型并获取输入输出信息： ```python import onnxruntime as ort import numpy as np import cv2 def load_onnx_model(model_path): # 指定执行提供者，优先使用 GPU providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] session = ort.InferenceSession(model_path, providers=providers) # 获取模型输入信息 model_inputs = session.get_inputs() input_shape = model_inputs[0].shape input_width = input_shape[2] input_height = input_shape[3] return session, model_inputs, input_width, input_height ``` #### 图像预处理 YOLO 模型通常要求输入图像具有固定的尺寸，并且需要对图像进行归一化处理。以下是一个典型的图像预处理步骤： ```python def preprocess_image(image_path, input_width, input_height): # 读取图像 image = cv2.imread(image_path) # 获取原始图像尺寸 original_height, original_width = image.shape[:2] # 将图像调整为模型输入尺寸 resized_image = cv2.resize(image, (input_width, input_height)) # 归一化处理 normalized_image = resized_image.astype(np.float32) / 255.0 # 调整维度以匹配模型输入要求 (1, C, H, W) input_image = np.transpose(normalized_image, (2, 0, 1)) input_image = np.expand_dims(input_image, axis=0) return image, input_image.astype(np.float32) ``` #### 推理执行 使用 `onnxruntime` 进行推理时，需要将预处理后的图像输入到模型中，并获取输出结果： ```python def run_inference(session, input_image): # 获取输入名称 input_name = session.get_inputs()[0].name # 执行推理 outputs = session.run(None, {input_name: input_image}) return outputs ``` #### 后处理 YOLOv11 OBB 模型的输出通常包括边界框坐标、类别概率和方向信息。后处理步骤包括解码边界框、应用非极大值抑制（NMS）以及绘制结果。 ```python def postprocess(outputs, image, conf_threshold=0.5, iou_threshold=0.4): # 假设输出为 (1, num_anchors, 6)，其中 6 包括 [x, y, w, h, angle, class_prob] predictions = outputs[0] # 过滤低置信度的预测 confidence_mask = predictions[:, :, 4] > conf_threshold boxes = predictions[confidence_mask] # 解码边界框 x_center = boxes[:, 0] y_center = boxes[:, 1] width = boxes[:, 2] height = boxes[:, 3] angle = boxes[:, 4] class_probs = boxes[:, 5:] # 应用非极大值抑制 indices = cv2.dnn.NMSBoxesRotated( bboxes=[(x, y, w, h, a) for x, y, w, h, a in zip(x_center, y_center, width, height, angle)], scores=np.max(class_probs, axis=1), score_threshold=conf_threshold, nms_threshold=iou_threshold ) # 绘制检测结果 for i in indices: box = cv2.boxPoints(((x_center[i], y_center[i]), (width[i], height[i]), angle[i])) box = np.int0(box) cv2.drawContours(image, [box], 0, (0, 255, 0), 2) return image ``` #### 完整流程 将上述步骤整合为一个完整的推理流程： ```python def detect_objects(model_path, image_path): # 加载模型 session, model_inputs, input_width, input_height = load_onnx_model(model_path) # 预处理图像 original_image, input_image = preprocess_image(image_path, input_width, input_height) # 推理 outputs = run_inference(session, input_image) # 后处理 result_image = postprocess(outputs, original_image) # 显示结果 cv2.imshow("Detection Result", result_image) cv2.waitKey(0) cv2.destroyAllWindows() ``` ### 环境配置 为了确保 ONNX 模型能够顺利运行，建议安装以下依赖库： ```bash pip install onnxruntime-gpu==1.13.1 opencv-python==4.7.0.68 numpy==1.24.1 ``` 如果使用 GPU 进行推理，确保安装了支持 CUDA 的版本，并且系统中已安装相应的 CUDA 驱动和 cuDNN 库。 ### 结果展示 推理完成后，可以使用 OpenCV 显示检测结果。对于 OBB 检测，结果将以旋转矩形框的形式绘制在图像上，能够准确表示目标的方向和位置。 #### 示例代码 ```python detect_objects("yolov11_obb.onnx", "test_image.jpg") ``` ---