浅析解析 3D NMS 算法及实现

深入解析 3D NMS 算法及实现

在 3D 目标检测任务中，如基于点云的物体检测或 3D 边界框回归，深度学习模型通常会输出一组重叠的 3D 边界框。这些框可能包含冗余的预测，因此需要通过非极大值抑制（Non-Maximum Suppression，NMS）来筛选出最优的检测结果。本文将详细介绍如何实现 3D NMS 算法，并展示修正后的代码实现。

什么是 3D NMS？

NMS 是一种常用的后处理算法，目的是从一组重叠的边界框中选择出最优的框，通常基于它们的置信分数。对于 3D NMS，我们将考虑 3D 边界框之间的交并比（IoU, Intersection over Union），并利用其来筛选掉冗余的框。

基本原理

输入：一组 3D 边界框及对应的置信分数。
输出：筛选后的 3D 边界框。
步骤：
1. 按照置信分数对边界框进行降序排序。
2. 依次选择得分最高的框，并与剩余框计算 IoU。
3. 移除与当前框重叠度高于给定阈值的框。
4. 重复上述过程，直到所有框都被处理完。

3D NMS 的关键部分

1. 3D IoU 的计算

IoU 是 NMS 算法的核心，用来衡量两个 3D 边界框的重叠度。在 3D 空间中，IoU 的计算不仅仅是计算 2D 面积的交并比，更要涉及到体积的计算。具体来说，IoU 的计算包括以下步骤：

交集体积：计算两个框在宽度（w）、高度（h）和深度（d）三个方向上的重叠部分，并求得交集体积。
并集体积：并集体积是两个框的体积之和减去交集体积。

交集体积公式：
$max_inter − x min_inter ) ⋅ max ⁡ ( 0 , y max_inter − y min_inter ) ⋅ max ⁡ ( 0 , z max_inter − z min_inter ) V_{\text{intersection}} = \max(0, x_{\text{max\_inter}} - x_{\text{min\_inter}}) \cdot \max(0, y_{\text{max\_inter}} - y_{\text{min\_inter}}) \cdot \max(0, z_{\text{max\_inter}} - z_{\text{min\_inter}})$

并集体积公式：
$V_{\text{union}} = V_{\text{box1}} + V_{\text{box2}} - V_{\text{intersection}}$

2. 处理特殊情况

在实际应用中，可能会遇到以下两种情况：

交集体积为 0：如果两个框没有交集，IoU 应直接返回 0。
框为空或无效：如果其中一个框为空，IoU 也应直接返回 0。

3D NMS 的代码实现

计算 3D IoU

def compute_3d_iou(box1, box2):
    """
    计算两个3D边界框的IoU（假设无旋转）
    :param box1: [x, y, z, w, h, d] 第一个框的中心坐标和尺寸
    :param box2: [x, y, z, w, h, d] 第二个框的中心坐标和尺寸
    :return: IoU 值
    """
    # 如果 box2 为空，直接返回 IoU = 0
    if box2 is None or len(box2) == 0:
        return 0.0

    # 计算交集体积
    inter_w = max(0, min(box1[0] + box1[3]/2, box2[0] + box2[3]/2) - max(box1[0] - box1[3]/2, box2[0] - box2[3]/2))
    inter_h = max(0, min(box1[1] + box1[4]/2, box2[1] + box2[4]/2) - max(box1[1] - box1[4]/2, box2[1] - box2[4]/2))
    inter_d = max(0, min(box1[2] + box1[5]/2, box2[2] + box2[5]/2) - max(box1[2] - box1[5]/2, box2[2] - box2[5]/2))
    inter_volume = inter_w * inter_h * inter_d

    # 如果交集体积为 0，直接返回 IoU = 0
    if inter_volume == 0:
        return 0.0

    # 计算并集体积
    volume1 = box1[3] * box1[4] * box1[5]
    volume2 = box2[3] * box2[4] * box2[5]
    union_volume = volume1 + volume2 - inter_volume

    # 返回交并比
    return inter_volume / union_volume

3D NMS 的实现

def nms_3d(boxes, scores, iou_threshold):
    """
    3D NMS 算法
    :param boxes: List of [x, y, z, w, h, d], 所有边界框
    :param scores: List of float, 每个框的置信分数
    :param iou_threshold: float, IoU 阈值
    :return: 保留框的索引列表
    """
    # 按得分排序
    indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
    selected_indices = []

    while len(indices) > 0:
        # 当前最高分框
        current = indices[0]
        selected_indices.append(current)
        
        # 剩余框
        rest_indices = indices[1:]
        
        # 计算 IoU，并筛选 IoU 小于阈值的框
        ious = np.array([compute_3d_iou(boxes[current], boxes[i]) for i in rest_indices])
        indices = [rest_indices[i] for i in range(len(ious)) if ious[i] < iou_threshold]

    return selected_indices

代码详解

compute_3d_iou：
- 该函数计算两个 3D 边界框的交并比（IoU）。
- 我们首先计算交集体积，并且在交集体积为 0 时返回 IoU 为 0。
- 然后我们计算并集体积，最后返回交并比。
nms_3d：
- 该函数执行 3D NMS 操作。
- 首先按照置信分数对所有框进行排序，然后依次选择得分最高的框，并与其他框计算 IoU，筛选掉 IoU 大于给定阈值的框。