yolov_obb检测onnx模型部署

本文介绍了yolo系列中obb检测的onnx模型部署，在借鉴源码的基础上按照自己的需求做了相应的更改。

🏷️一、模型输出

首先介绍一下yolov_obb模型的输出。YOLOv_OBB（Oriented Bounding Box）检测模型的原始输出通常是一个三维张量，其每个维度的含义如下：

第一个维度：批量大小（Batch Size）

第一个维度表示输入图像的批量大小，即一次处理多少张图像。例如，如果批量大小为8，则张量的第一个维度为8，表示同时处理8张图像。

第二个维度：特征大小（Feature Size）

第二个维度表示特征的大小，即模型输出的特征的长度。特征的大小取决于当前任务类别的数量。例如，如果类别数量为20，则特征的大小为：4 + 20 + 1 = 25。其中4标语与检测框的中心点坐标以及宽和高；其中1代表检测框的倾斜角度。
其数据顺序为：x，y，w，h，num_classes，theta

第三维度：目标的数量

第三个维度表示检测出目标的数量，即模型输出的检测框的数量。

🏷️二、整个流程的大致为：

1、初始化：模型加载、配置参数设置

2、预处理：图像 resize、padding、归一化

3、推理：ONNX Runtime 推理

4、后处理：置信度过滤、NMS、坐标还原、检测框处理

5、可视化：绘制边界框、文字标注

🏷️三、以下为完整的代码：

import argparse
import cv2
import numpy as np
import onnxruntime as ort


class YOLOv8_Obb:
    """YOLOv8 obb model."""

    def __init__(self, onnx_model, conf_threshold=0.5, iou_threshold=0.45):
        """
        Initialization.

        Args:
            onnx_model (str): Path to the ONNX model.
            conf_threshold (float): confidence threshold for filtering predictions.
            iou_threshold (float): iou threshold for NMS.
        """
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold
        # Build Ort session
        self.session = ort.InferenceSession(
            onnx_model,
            providers=[("CUDAExecutionProvider", {
   "cudnn_conv_algo_search": "DEFAULT"}), 'CPUExecutionProvider']
            if ort.get_device() == "GPU"
            else ["CPUExecutionProvider"],
        )

        # Numpy dtype: support both FP32 and FP16 onnx model
        self.ndtype = np.half if self.session.get_inputs()[0].type == "tensor(float16)" else np.single

        # Get model width and height(YOLOv8-seg only has one input)
        self.model_height, self.model_width = 640, 640  # [x.shape for x in self.session.get_inputs()][0][-2:]

        # Load COCO class names
        self.classes = [str(i) for i in range(1, 25)]  #yaml_load(check_yaml("coco8.yaml"))["names"]

    def __call__(self, im0):
        """
        The whole pipeline: pre-process -> inference -> post-process.

        Args:
            im0 (Numpy.ndarray): original input image.
        Returns:
            boxes (List): list of obboxes (n, x_y_w_h_r_conf_cls).
        """
        # Pre-process
        im, ratio, (pad_w, pad_h) = self.preprocess(im0)

        # Ort inference
        preds = self.session.run(None, {
   self.session.get_inputs()[0].name: im})

        # Post-process
        boxes = self.postprocess(
            preds,
            im0=im0,
            ratio=ratio,
            pad_w=pad_w,
            pad_h=pad_h
        )
        return boxes

    def mask_iou(self, mask1, mask2):
        """
        计算两个二进制掩码的交并比 (IoU)。
        :param mask1: 第一个掩码，形状为 (H, W) 的布尔数组。
        :param mask2: 第二个掩码，形状为 (H, W) 的布尔数组。
        :return: 交并比 (IoU)。
        """
        # 计算交集
        intersection = np.logical_and(mask1, mask2)
        intersection_area = np.sum(intersection)

        # 计算并集
        union = np.logical_or(mask1, mask2)
        union_area = np.sum(union)

        # 计算 IoU
        iou = intersection_area / (union_area + 1e-6)  # 防止除以零

        if intersection_area / np.sum(mask1) > 0.8 or intersection_area / np.sum(mask2) > 0.8:
            return 1
        else:
            return iou

    def probiou(self, obb1: np.ndarray, obb2: np.ndarray, eps=1e-7):
        """
        Calculate the prob IoU between oriented bounding boxes, https://arxiv.org/pdf/2106.06072v1.pdf.

        Args:
            obb1 (np.ndarray): A tensor of shape (N, 5) representing ground truth obbs, with xywhr format.
            obb2 (np.ndarray): A tensor of shape (M, 5) representing predicted obbs, with xywhr format.
            eps (float, optional): A small value to avoid division by zero. Defaults to 1e-7.

        Returns:
            (torch.Tensor): A tensor of shape (N, M) representing obb similarities.
        """
        # obb1 = torch.from_numpy(obb1) if isinstance(obb1, np.ndarray) else obb1
        # obb2 = torch.from_numpy(obb2) if isinstance(obb2, np.ndarray) else obb2
        if len(obb1.shape) == 2:
            x1, y1 = np.split(obb1[..., :2], 2, axis=-1)  # [n,1]
            x2, y2 = (np.squeeze