cv2.dnn.NMSBoxesRotated源码实现

OpenCV中的cv2.dnn.NMSBoxesRotated函数是用于对旋转框进行非极大值抑制（NMS）的函数。它的源码实现如下：

def NMSBoxesRotated(boxes, scores, score_threshold, nms_threshold, eta=1.0):
    """
    boxes: 旋转框的坐标，形状为(n, 5)，其中n为旋转框的数量，每个旋转框包含5个元素，
          分别表示旋转框的中心点坐标(x,y)、宽度(w)、高度(h)和旋转角度(theta)
    scores: 旋转框的得分，形状为(n,)
    score_threshold: 得分阈值，低于该阈值的旋转框将被过滤掉
    nms_threshold: NMS阈值，用于控制旋转框的重叠度
    eta: 用于调整旋转框在NMS过程中的得分权重

    返回值：
    indices: 一个二维数组，形状为(k, 1)，其中k是保留下来的旋转框的数量，每个元素表示
             保留下来的旋转框的索引号。
    """
    # 保留得分大于等于score_threshold的旋转框
    indices = np.where(scores >= score_threshold)[0]
    if len(indices) == 0:
        return np.empty((0, 1), dtype=np.int32)

    # 提取旋转框的坐标和得分
    boxes = boxes[indices]
    scores = scores[indices]

    # 计算旋转框的面积
    areas = boxes[:, 2] * boxes[:, 3]

    # 对旋转框按得分从大到小排序
    order = np.argsort(-scores)

    # 保留下来的旋转框的索引号
    keep = []
    while order.size > 0:
        idx = order[0]  # 得分最高的旋转框的索引号
        keep.append(idx)  # 保留该旋转框

        # 计算该旋转框与其他旋转框的IoU
        xx1 = np.maximum(boxes[idx, 0] - boxes[idx, 2] / 2, boxes[:, 0] - boxes[:, 2] / 2)
        yy1 = np.maximum(boxes[idx, 1] - boxes[idx, 3] / 2, boxes[:, 1] - boxes[:, 3] / 2)
        xx2 = np.minimum(boxes[idx, 0] + boxes[idx, 2] / 2, boxes[:, 0] + boxes[:, 2] / 2)
        yy2 = np.minimum(boxes[idx, 1] + boxes[idx, 3] / 2, boxes[:, 1] + boxes[:, 3] / 2)
        theta = boxes[idx, 4] - boxes[:, 4]
        idxs = np.where((xx2 > xx1) & (yy2 > yy1) & (theta <= nms_threshold))[0]

        # 计算权重
        weight = np.ones(idxs.shape[0])
        if eta > 0:
            dx = boxes[idxs, 0] - boxes[idx, 0]
            dy = boxes[idxs, 1] - boxes[idx, 1]
            delta_theta = boxes[idxs, 4] - boxes[idx, 4]
            rho = np.sqrt(dx ** 2 + dy ** 2)
            weight = np.exp(-(rho ** 2) / (2 * eta * eta))
            weight *= np.exp(-(delta_theta ** 2) / (2 * (np.pi / 180) ** 2))

        # 应用权重
        scores[idxs] *= weight

        # 移除IoU大于等于NMS阈值的旋转框
        idxs = np.delete(idxs, np.where(scores[idxs] < score_threshold)[0])
        order = np.delete(order, np.where(order == idxs)[0])

    return np.array(keep, dtype=np.int32).reshape(-1, 1)

这个函数的实现过程与传统的NMS算法基本一致，主要是在计算IoU时需要考虑旋转框的旋转角度，并且在应用权重时还需要考虑旋转角度的差异。