[Fix] make iou3d.boxes_iou3d actually calculate 3D IoU (#2018)

* fix iou3d

* rename variables in test

* fix comments

mmcv/ops/__init__.py

@@ -29,8 +29,8 @@ from .gather_points import gather_points
 from .group_points import GroupAll, QueryAndGroup, grouping_operation
 from .info import (get_compiler_version, get_compiling_cuda_version,
                    get_onnxruntime_op_path)
-from .iou3d import (boxes_iou3d, boxes_iou_bev, nms3d, nms3d_normal, nms_bev,
-                    nms_normal_bev)
+from .iou3d import (boxes_iou3d, boxes_iou_bev, boxes_overlap_bev, nms3d,
+                    nms3d_normal, nms_bev, nms_normal_bev)
 from .knn import knn
 from .masked_conv import MaskedConv2d, masked_conv2d
 from .min_area_polygons import min_area_polygons
@@ -91,12 +91,12 @@ __all__ = [
     'three_interpolate', 'MultiScaleDeformableAttention', 'BorderAlign',
     'border_align', 'gather_points', 'furthest_point_sample',
     'furthest_point_sample_with_dist', 'PointsSampler', 'Correlation',
-    'boxes_iou3d', 'boxes_iou_bev', 'nms_bev', 'nms_normal_bev', 'nms3d',
-    'nms3d_normal', 'Voxelization', 'voxelization', 'dynamic_scatter',
-    'DynamicScatter', 'RoIAwarePool3d', 'SparseConv2d', 'SparseConv3d',
-    'SparseConvTranspose2d', 'SparseConvTranspose3d', 'SparseInverseConv2d',
-    'SparseInverseConv3d', 'SubMConv2d', 'SubMConv3d', 'SparseModule',
-    'SparseSequential', 'SparseMaxPool2d', 'SparseMaxPool3d',
+    'boxes_iou3d', 'boxes_iou_bev', 'boxes_overlap_bev', 'nms_bev',
+    'nms_normal_bev', 'nms3d', 'nms3d_normal', 'Voxelization', 'voxelization',
+    'dynamic_scatter', 'DynamicScatter', 'RoIAwarePool3d', 'SparseConv2d',
+    'SparseConv3d', 'SparseConvTranspose2d', 'SparseConvTranspose3d',
+    'SparseInverseConv2d', 'SparseInverseConv3d', 'SubMConv2d', 'SubMConv3d',
+    'SparseModule', 'SparseSequential', 'SparseMaxPool2d', 'SparseMaxPool3d',
     'SparseConvTensor', 'scatter_nd', 'points_in_boxes_part',
     'points_in_boxes_cpu', 'points_in_boxes_all', 'points_in_polygons',
     'min_area_polygons', 'active_rotated_filter', 'convex_iou', 'convex_giou',

mmcv/ops/csrc/common/cuda/iou3d_cuda_kernel.cuh

@@ -216,11 +216,11 @@ __device__ inline float iou_bev(const float *box_a, const float *box_b) {
   return s_overlap / fmaxf(sa + sb - s_overlap, EPS);
 }
 
-__global__ void iou3d_boxes_iou3d_forward_cuda_kernel(const int num_a,
-                                                      const float *boxes_a,
-                                                      const int num_b,
-                                                      const float *boxes_b,
-                                                      float *ans_iou) {
+__global__ void iou3d_boxes_overlap_bev_forward_cuda_kernel(
+    const int num_a, const float *boxes_a, const int num_b,
+    const float *boxes_b, float *ans_overlap) {
+  // params boxes_a: (N, 7) [x, y, z, dx, dy, dz, heading]
+  // params boxes_b: (M, 7) [x, y, z, dx, dy, dz, heading]
   CUDA_2D_KERNEL_LOOP(b_idx, num_b, a_idx, num_a) {
     if (a_idx >= num_a || b_idx >= num_b) {
       return;
@@ -228,8 +228,8 @@ __global__ void iou3d_boxes_iou3d_forward_cuda_kernel(const int num_a,
 
     const float *cur_box_a = boxes_a + a_idx * 7;
     const float *cur_box_b = boxes_b + b_idx * 7;
-    float cur_iou_bev = iou_bev(cur_box_a, cur_box_b);
-    ans_iou[a_idx * num_b + b_idx] = cur_iou_bev;
+    float cur_overlap = box_overlap(cur_box_a, cur_box_b);
+    ans_overlap[a_idx * num_b + b_idx] = cur_overlap;
   }
 }
 

mmcv/ops/csrc/pytorch/cuda/cudabind.cpp

@@ -564,11 +564,11 @@ REGISTER_DEVICE_IMPL(group_points_forward_impl, CUDA,
 REGISTER_DEVICE_IMPL(group_points_backward_impl, CUDA,
                      group_points_backward_cuda);
 
-void IoU3DBoxesIoU3DForwardCUDAKernelLauncher(const int num_a,
-                                              const Tensor boxes_a,
-                                              const int num_b,
-                                              const Tensor boxes_b,
-                                              Tensor ans_iou);
+void IoU3DBoxesOverlapBevForwardCUDAKernelLauncher(const int num_a,
+                                                   const Tensor boxes_a,
+                                                   const int num_b,
+                                                   const Tensor boxes_b,
+                                                   Tensor ans_overlap);
 
 void IoU3DNMS3DForwardCUDAKernelLauncher(const Tensor boxes,
                                          unsigned long long* mask,
@@ -580,11 +580,11 @@ void IoU3DNMS3DNormalForwardCUDAKernelLauncher(const Tensor boxes,
                                                int boxes_num,
                                                float nms_overlap_thresh);
 
-void iou3d_boxes_iou3d_forward_cuda(const int num_a, const Tensor boxes_a,
-                                    const int num_b, const Tensor boxes_b,
-                                    Tensor ans_iou) {
-  IoU3DBoxesIoU3DForwardCUDAKernelLauncher(num_a, boxes_a, num_b, boxes_b,
-                                           ans_iou);
+void iou3d_boxes_overlap_bev_forward_cuda(const int num_a, const Tensor boxes_a,
+                                          const int num_b, const Tensor boxes_b,
+                                          Tensor ans_overlap) {
+  IoU3DBoxesOverlapBevForwardCUDAKernelLauncher(num_a, boxes_a, num_b, boxes_b,
+                                                ans_overlap);
 };
 
 void iou3d_nms3d_forward_cuda(const Tensor boxes, unsigned long long* mask,
@@ -600,9 +600,9 @@ void iou3d_nms3d_normal_forward_cuda(const Tensor boxes,
                                             nms_overlap_thresh);
 };
 
-void iou3d_boxes_iou3d_forward_impl(const int num_a, const Tensor boxes_a,
-                                    const int num_b, const Tensor boxes_b,
-                                    Tensor ans_iou);
+void iou3d_boxes_overlap_bev_forward_impl(const int num_a, const Tensor boxes_a,
+                                          const int num_b, const Tensor boxes_b,
+                                          Tensor ans_overlap);
 
 void iou3d_nms3d_forward_impl(const Tensor boxes, unsigned long long* mask,
                               int boxes_num, float nms_overlap_thresh);
@@ -611,8 +611,8 @@ void iou3d_nms3d_normal_forward_impl(const Tensor boxes,
                                      unsigned long long* mask, int boxes_num,
                                      float nms_overlap_thresh);
 
-REGISTER_DEVICE_IMPL(iou3d_boxes_iou3d_forward_impl, CUDA,
-                     iou3d_boxes_iou3d_forward_cuda);
+REGISTER_DEVICE_IMPL(iou3d_boxes_overlap_bev_forward_impl, CUDA,
+                     iou3d_boxes_overlap_bev_forward_cuda);
 REGISTER_DEVICE_IMPL(iou3d_nms3d_forward_impl, CUDA, iou3d_nms3d_forward_cuda);
 REGISTER_DEVICE_IMPL(iou3d_nms3d_normal_forward_impl, CUDA,
                      iou3d_nms3d_normal_forward_cuda);

mmcv/ops/csrc/pytorch/cuda/iou3d_cuda.cu

@@ -12,11 +12,11 @@ All Rights Reserved 2019-2020.
 #include "iou3d_cuda_kernel.cuh"
 #include "pytorch_cuda_helper.hpp"
 
-void IoU3DBoxesIoU3DForwardCUDAKernelLauncher(const int num_a,
-                                              const Tensor boxes_a,
-                                              const int num_b,
-                                              const Tensor boxes_b,
-                                              Tensor ans_iou) {
+void IoU3DBoxesOverlapBevForwardCUDAKernelLauncher(const int num_a,
+                                                   const Tensor boxes_a,
+                                                   const int num_b,
+                                                   const Tensor boxes_b,
+                                                   Tensor ans_overlap) {
   at::cuda::CUDAGuard device_guard(boxes_a.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
@@ -25,9 +25,9 @@ void IoU3DBoxesIoU3DForwardCUDAKernelLauncher(const int num_a,
               GET_BLOCKS(num_a, THREADS_PER_BLOCK_IOU3D));
   dim3 threads(THREADS_PER_BLOCK_IOU3D, THREADS_PER_BLOCK_IOU3D);
 
-  iou3d_boxes_iou3d_forward_cuda_kernel<<<blocks, threads, 0, stream>>>(
+  iou3d_boxes_overlap_bev_forward_cuda_kernel<<<blocks, threads, 0, stream>>>(
       num_a, boxes_a.data_ptr<float>(), num_b, boxes_b.data_ptr<float>(),
-      ans_iou.data_ptr<float>());
+      ans_overlap.data_ptr<float>());
 
   AT_CUDA_CHECK(cudaGetLastError());
 }

mmcv/ops/csrc/pytorch/iou3d.cpp

@@ -12,11 +12,11 @@ All Rights Reserved 2019-2020.
 
 const int THREADS_PER_BLOCK_NMS = sizeof(unsigned long long) * 8;
 
-void iou3d_boxes_iou3d_forward_impl(const int num_a, const Tensor boxes_a,
-                                    const int num_b, const Tensor boxes_b,
-                                    Tensor ans_iou) {
-  DISPATCH_DEVICE_IMPL(iou3d_boxes_iou3d_forward_impl, num_a, boxes_a, num_b,
-                       boxes_b, ans_iou);
+void iou3d_boxes_overlap_bev_forward_impl(const int num_a, const Tensor boxes_a,
+                                          const int num_b, const Tensor boxes_b,
+                                          Tensor ans_overlap) {
+  DISPATCH_DEVICE_IMPL(iou3d_boxes_overlap_bev_forward_impl, num_a, boxes_a,
+                       num_b, boxes_b, ans_overlap);
 }
 
 void iou3d_nms3d_forward_impl(const Tensor boxes, unsigned long long *mask,
@@ -32,14 +32,16 @@ void iou3d_nms3d_normal_forward_impl(const Tensor boxes,
                        nms_overlap_thresh);
 }
 
-void iou3d_boxes_iou3d_forward(Tensor boxes_a, Tensor boxes_b, Tensor ans_iou) {
+void iou3d_boxes_overlap_bev_forward(Tensor boxes_a, Tensor boxes_b,
+                                     Tensor ans_overlap) {
   // params boxes: (N, 7) [x, y, z, dx, dy, dz, heading]
   // params boxes_b: (M, 5)
   // params ans_overlap: (N, M)
   int num_a = boxes_a.size(0);
   int num_b = boxes_b.size(0);
 
-  iou3d_boxes_iou3d_forward_impl(num_a, boxes_a, num_b, boxes_b, ans_iou);
+  iou3d_boxes_overlap_bev_forward_impl(num_a, boxes_a, num_b, boxes_b,
+                                       ans_overlap);
 }
 
 void iou3d_nms3d_forward(Tensor boxes, Tensor keep, Tensor keep_num,

mmcv/ops/csrc/pytorch/pybind.cpp

@@ -116,7 +116,8 @@ void bbox_overlaps(const Tensor bboxes1, const Tensor bboxes2, Tensor ious,
 void knn_forward(Tensor xyz_tensor, Tensor new_xyz_tensor, Tensor idx_tensor,
                  Tensor dist2_tensor, int b, int n, int m, int nsample);
 
-void iou3d_boxes_iou3d_forward(Tensor boxes_a, Tensor boxes_b, Tensor ans_iou);
+void iou3d_boxes_overlap_bev_forward(Tensor boxes_a, Tensor boxes_b,
+                                     Tensor ans_overlap);
 
 void iou3d_nms3d_forward(Tensor boxes, Tensor keep, Tensor keep_num,
                          float nms_overlap_thresh);
@@ -541,9 +542,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         py::arg("m"), py::arg("nsample"), py::arg("xyz_tensor"),
         py::arg("new_xyz_tensor"), py::arg("idx_tensor"),
         py::arg("dist2_tensor"));
-  m.def("iou3d_boxes_iou3d_forward", &iou3d_boxes_iou3d_forward,
-        "iou3d_boxes_iou3d_forward", py::arg("boxes_a"), py::arg("boxes_b"),
-        py::arg("ans_iou"));
+  m.def("iou3d_boxes_overlap_bev_forward", &iou3d_boxes_overlap_bev_forward,
+        "iou3d_boxes_overlap_bev_forward", py::arg("boxes_a"),
+        py::arg("boxes_b"), py::arg("ans_iou"));
   m.def("iou3d_nms3d_forward", &iou3d_nms3d_forward, "iou3d_nms3d_forward",
         py::arg("boxes"), py::arg("keep"), py::arg("num_out"),
         py::arg("nms_overlap_thresh"));

mmcv/ops/iou3d.py

@@ -8,11 +8,30 @@ from torch import Tensor
 from ..utils import ext_loader
 
 ext_module = ext_loader.load_ext('_ext', [
-    'iou3d_boxes_iou3d_forward', 'iou3d_nms3d_forward',
+    'iou3d_boxes_overlap_bev_forward', 'iou3d_nms3d_forward',
     'iou3d_nms3d_normal_forward'
 ])
 
 
+def boxes_overlap_bev(boxes_a: Tensor, boxes_b: Tensor) -> Tensor:
+    """Calculate boxes BEV overlap.
+
+    Args:
+        boxes_a (torch.Tensor): Input boxes a with shape (M, 7).
+        boxes_b (torch.Tensor): Input boxes b with shape (N, 7).
+
+    Returns:
+        torch.Tensor: BEV overlap result with shape (M, N).
+    """
+    ans_overlap = boxes_a.new_zeros(
+        torch.Size((boxes_a.shape[0], boxes_b.shape[0])))
+    ext_module.iou3d_boxes_overlap_bev_forward(boxes_a.contiguous(),
+                                               boxes_b.contiguous(),
+                                               ans_overlap)
+
+    return ans_overlap
+
+
 def boxes_iou3d(boxes_a: Tensor, boxes_b: Tensor) -> Tensor:
     """Calculate boxes 3D IoU.
 
@@ -21,15 +40,30 @@ def boxes_iou3d(boxes_a: Tensor, boxes_b: Tensor) -> Tensor:
         boxes_b (torch.Tensor): Input boxes b with shape (N, 7).
 
     Returns:
-        torch.Tensor: IoU result with shape (M, N).
+        torch.Tensor: 3D IoU result with shape (M, N).
     """
-    ans_iou = boxes_a.new_zeros(
-        torch.Size((boxes_a.shape[0], boxes_b.shape[0])))
+    assert boxes_a.shape[1] == boxes_b.shape[1] == 7,\
+        'Input boxes shape should be (N, 7)'
 
-    ext_module.iou3d_boxes_iou3d_forward(boxes_a.contiguous(),
-                                         boxes_b.contiguous(), ans_iou)
+    boxes_a_height_max = (boxes_a[:, 2] + boxes_a[:, 5] / 2).view(-1, 1)
+    boxes_a_height_min = (boxes_a[:, 2] - boxes_a[:, 5] / 2).view(-1, 1)
+    boxes_b_height_max = (boxes_b[:, 2] + boxes_b[:, 5] / 2).view(1, -1)
+    boxes_b_height_min = (boxes_b[:, 2] - boxes_b[:, 5] / 2).view(1, -1)
 
-    return ans_iou
+    overlaps_bev = boxes_a.new_zeros(
+        torch.Size((boxes_a.shape[0], boxes_b.shape[0])))
+    ext_module.iou3d_boxes_overlap_bev_forward(boxes_a.contiguous(),
+                                               boxes_b.contiguous(),
+                                               overlaps_bev)
+
+    max_of_min = torch.max(boxes_a_height_min, boxes_b_height_min)
+    min_of_max = torch.min(boxes_a_height_max, boxes_b_height_max)
+    overlaps_h = torch.clamp(min_of_max - max_of_min, min=0)
+    overlaps_3d = overlaps_bev * overlaps_h
+    vol_a = (boxes_a[:, 3] * boxes_a[:, 4] * boxes_a[:, 5]).view(-1, 1)
+    vol_b = (boxes_b[:, 3] * boxes_b[:, 4] * boxes_b[:, 5]).view(1, -1)
+    iou3d = overlaps_3d / torch.clamp(vol_a + vol_b - overlaps_3d, min=1e-6)
+    return iou3d
 
 
 def nms3d(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:

tests/test_ops/test_iou3d.py

@@ -3,7 +3,38 @@ import numpy as np
 import pytest
 import torch
 
-from mmcv.ops import boxes_iou3d, nms3d, nms3d_normal
+from mmcv.ops import boxes_iou3d, boxes_overlap_bev, nms3d, nms3d_normal
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(), reason='requires CUDA support')
+def test_boxes_overlap_bev():
+    np_boxes1 = np.asarray([[1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 0.0],
+                            [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 0.0],
+                            [3.0, 3.0, 3.0, 3.0, 2.0, 2.0, 0.0]],
+                           dtype=np.float32)
+    np_boxes2 = np.asarray([[1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 0.0],
+                            [1.0, 1.0, 1.0, 2.0, 2.0, 2.0, np.pi / 2],
+                            [1.0, 1.0, 1.0, 2.0, 2.0, 2.0, np.pi / 4]],
+                           dtype=np.float32)
+    np_expect_overlaps = np.asarray(
+        [[4.0, 4.0, (8 + 8 * 2**0.5) /
+          (3 + 2 * 2**0.5)], [1.0, 1.0, 1.0], [0.0, 0.0, 0.0]],
+        dtype=np.float32)
+
+    boxes1 = torch.from_numpy(np_boxes1).cuda()
+    boxes2 = torch.from_numpy(np_boxes2).cuda()
+
+    # test for 3 boxes
+    overlaps = boxes_overlap_bev(boxes1, boxes2)
+    assert np.allclose(overlaps.cpu().numpy(), np_expect_overlaps, atol=1e-4)
+
+    # test for many boxes
+    boxes2 = boxes2.repeat_interleave(555, 0)
+
+    overlaps = boxes_overlap_bev(boxes1, boxes2)
+    assert np.allclose(
+        overlaps.cpu().numpy(), np_expect_overlaps.repeat(555, 1), atol=1e-4)
 
 
 @pytest.mark.skipif(
@@ -17,9 +48,10 @@ def test_boxes_iou3d():
                             [1.0, 1.0, 1.0, 2.0, 2.0, 2.0, np.pi / 2],
                             [1.0, 1.0, 1.0, 2.0, 2.0, 2.0, np.pi / 4]],
                            dtype=np.float32)
-    np_expect_ious = np.asarray([[1.0, 1.0, 1.0 / 2**0.5],
-                                 [1.0 / 7, 1.0 / 7, 1.0 / 7], [0.0, 0.0, 0.0]],
-                                dtype=np.float32)
+    np_expect_ious = np.asarray(
+        [[1.0, 1.0, 1.0 / 2**0.5], [1.0 / 15, 1.0 / 15, 1.0 / 15],
+         [0.0, 0.0, 0.0]],
+        dtype=np.float32)
 
     boxes1 = torch.from_numpy(np_boxes1).cuda()
     boxes2 = torch.from_numpy(np_boxes2).cuda()