[feature]: Add multi-group head of CenterPoint (#49)

* Add modules.

* Add test_center_head.

* Add docstring.

* Change comments.

* Add dcn_head.

* Add doc_string.

* Add get_targets.

* Can use_get_targets.

* get_targets results aligned.

* Use box_structure.

* Use get_targets_single.

* Add docstring.

* Fix dcn_center_head unittest.

* Delete unnecessary unittest.

* Add docstring.

* Change format.

* Add circle_nms.

* Change structure of mg_head.

* Add bbox coder for centerpoint.

* Add docstrings.

* Add docstrings.

* Add get_bboxes and unittest.

* Change docstring.

* Add img_metas.

* Change bbox coder unittest.

* Add task_detections.

* Change docstring.

* Change circle_nms to cpu.

* Change test_nms.

* Change score_th, chang keep to long type.

* Change docstring and unittest.

* Remove unnecessary things.

* Move gaussian.

* move clip_sigmoid, change dict.

* Change config.

* Change test_heads.

* Move weight initialization to init_weights func.

* Remove loc_loss_element adn==nd num_postive.

* Change bboxes to the right format.

* Change loss and bbox order.

* Update test_heads.

* Change loss.

* Change names in mg_head, change head unittest.

* Remove centerpoint_focal_loss, change docstring.

* Change topK default to 80.

* Change boxes in test_nms. Change task_boxes defaults to None.

* Fix rotate nms bug.

* Change docstring.

* Add docstring for get_task_detection and loss.

* Remove gaussian funcs, change mg_head.

* Change gaussianfocalloss to mean.

* change centerpoint_bbox_coder '/' to torch.div, fix centerhead unittest.

* Change div to '/'

* Change order in centerpoint_coder, change names, change dcn layer.

* Fix import in __init__

* Add gaussian unittest.

* Remove np ops in mg_head.

* Update docstring.

* Fix docstring use config to build head.

* Remove **kwargs

* Remove unnecessary codes, change order of bboxes.

* Remove '\' in args and pdb, change loss_bbox.

* Fix test_heads unittest.

* Remove unnecessary comments

* Change bbox order in rotate nms.

* Remove unnecessary attributes

* Change name, remove float

mmdet3d/core/__init__.py

@@ -2,5 +2,6 @@ from .anchor import *  # noqa: F401, F403
 from .bbox import *  # noqa: F401, F403
 from .evaluation import *  # noqa: F401, F403
 from .post_processing import *  # noqa: F401, F403
+from .utils import *  # noqa: F401, F403
 from .visualizer import *  # noqa: F401, F403
 from .voxel import *  # noqa: F401, F403

mmdet3d/core/bbox/coders/__init__.py

 from mmdet.core.bbox import build_bbox_coder
 from .anchor_free_bbox_coder import AnchorFreeBBoxCoder
+from .centerpoint_bbox_coders import CenterPointBBoxCoder
 from .delta_xyzwhlr_bbox_coder import DeltaXYZWLHRBBoxCoder
 from .partial_bin_based_bbox_coder import PartialBinBasedBBoxCoder
 
 __all__ = [
     'build_bbox_coder', 'DeltaXYZWLHRBBoxCoder', 'PartialBinBasedBBoxCoder',
-    'AnchorFreeBBoxCoder'
+    'CenterPointBBoxCoder', 'AnchorFreeBBoxCoder'
 ]

mmdet3d/core/bbox/coders/centerpoint_bbox_coders.py

+import torch
+
+from mmdet.core.bbox import BaseBBoxCoder
+from mmdet.core.bbox.builder import BBOX_CODERS
+
+
+@BBOX_CODERS.register_module()
+class CenterPointBBoxCoder(BaseBBoxCoder):
+    """Bbox coder for CenterPoint.
+
+    Args:
+        pc_range (list[float]): Range of point cloud.
+        out_size_factor (int): Downsample factor of the model.
+        voxel_size (list[float]): Size of voxel.
+        post_center_range (list[float]): Limit of the center.
+            Default: None.
+        max_num (int): Max number to be kept. Default: 100.
+        score_threshold (float): Threshold to filter boxes based on score.
+            Default: None.
+        code_size (int): Code size of bboxes. Default: 9
+    """
+
+    def __init__(self,
+                 pc_range,
+                 out_size_factor,
+                 voxel_size,
+                 post_center_range=None,
+                 max_num=100,
+                 score_threshold=None,
+                 code_size=9):
+
+        self.pc_range = pc_range
+        self.out_size_factor = out_size_factor
+        self.voxel_size = voxel_size
+        self.post_center_range = post_center_range
+        self.max_num = max_num
+        self.score_threshold = score_threshold
+        self.code_size = code_size
+
+    def _gather_feat(self, feats, inds, feat_masks=None):
+        """Given feats and indexes, returns the gathered feats.
+
+        Args:
+            feats (torch.Tensor): Features to be transposed and gathered
+                with the shape of [B, 2, W, H].
+            inds (torch.Tensor): Indexes with the shape of [B, N].
+            feat_masks (torch.Tensor): Mask of the feats. Default: None.
+
+        Returns:
+            torch.Tensor: Gathered feats.
+        """
+        dim = feats.size(2)
+        inds = inds.unsqueeze(2).expand(inds.size(0), inds.size(1), dim)
+        feats = feats.gather(1, inds)
+        if feat_masks is not None:
+            feat_masks = feat_masks.unsqueeze(2).expand_as(feats)
+            feats = feats[feat_masks]
+            feats = feats.view(-1, dim)
+        return feats
+
+    def _topk(self, scores, K=80):
+        """Get indexes based on scores.
+
+        Args:
+            scores (torch.Tensor): scores with the shape of [B, N, W, H].
+            K (int): Number to be kept. Defaults to 80.
+
+        Returns:
+            tuple[torch.Tensor]
+                torch.Tensor: Selected scores with the shape of [B, K].
+                torch.Tensor: Selected indexes with the shape of [B, K].
+                torch.Tensor: Selected classes with the shape of [B, K].
+                torch.Tensor: Selected y coord with the shape of [B, K].
+                torch.Tensor: Selected x coord with the shape of [B, K].
+        """
+        batch, cat, height, width = scores.size()
+
+        topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), K)
+
+        topk_inds = topk_inds % (height * width)
+        topk_ys = (topk_inds.float() /
+                   torch.tensor(width, dtype=torch.float)).int().float()
+        topk_xs = (topk_inds % width).int().float()
+
+        topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), K)
+        topk_clses = (topk_ind / torch.tensor(K, dtype=torch.float)).int()
+        topk_inds = self._gather_feat(topk_inds.view(batch, -1, 1),
+                                      topk_ind).view(batch, K)
+        topk_ys = self._gather_feat(topk_ys.view(batch, -1, 1),
+                                    topk_ind).view(batch, K)
+        topk_xs = self._gather_feat(topk_xs.view(batch, -1, 1),
+                                    topk_ind).view(batch, K)
+
+        return topk_score, topk_inds, topk_clses, topk_ys, topk_xs
+
+    def _transpose_and_gather_feat(self, feat, ind):
+        """Given feats and indexes, returns the transposed and gathered feats.
+
+        Args:
+            feat (torch.Tensor): Features to be transposed and gathered
+                with the shape of [B, 2, W, H].
+            ind (torch.Tensor): Indexes with the shape of [B, N].
+
+        Returns:
+            torch.Tensor: Transposed and gathered feats.
+        """
+        feat = feat.permute(0, 2, 3, 1).contiguous()
+        feat = feat.view(feat.size(0), -1, feat.size(3))
+        feat = self._gather_feat(feat, ind)
+        return feat
+
+    def encode(self):
+        pass
+
+    def decode(self,
+               heat,
+               rot_sine,
+               rot_cosine,
+               hei,
+               dim,
+               vel,
+               reg=None,
+               task_id=-1):
+        """Decode bboxes.
+
+        Args:
+            heat (torch.Tensor): Heatmap with the shape of [B, N, W, H].
+            rot_sine (torch.Tensor): Sine of rotation with the shape of
+                [B, 1, W, H].
+            rot_cosine (torch.Tensor): Cosine of rotation with the shape of
+                [B, 1, W, H].
+            hei (torch.Tensor): Height of the boxes with the shape
+                of [B, 1, W, H].
+            dim (torch.Tensor): Dim of the boxes with the shape of
+                [B, 1, W, H].
+            vel (torch.Tensor): Velocity with the shape of [B, 1, W, H].
+            reg (torch.Tensor): Regression value of the boxes in 2D with
+                the shape of [B, 2, W, H]. Default: None.
+            task_id (int): Index of task. Default: -1.
+
+        Returns:
+            list[dict]: Decoded boxes.
+        """
+        batch, cat, _, _ = heat.size()
+
+        scores, inds, clses, ys, xs = self._topk(heat, K=self.max_num)
+
+        if reg is not None:
+            reg = self._transpose_and_gather_feat(reg, inds)
+            reg = reg.view(batch, self.max_num, 2)
+            xs = xs.view(batch, self.max_num, 1) + reg[:, :, 0:1]
+            ys = ys.view(batch, self.max_num, 1) + reg[:, :, 1:2]
+        else:
+            xs = xs.view(batch, self.max_num, 1) + 0.5
+            ys = ys.view(batch, self.max_num, 1) + 0.5
+
+        # rotation value and direction label
+        rot_sine = self._transpose_and_gather_feat(rot_sine, inds)
+        rot_sine = rot_sine.view(batch, self.max_num, 1)
+
+        rot_cosine = self._transpose_and_gather_feat(rot_cosine, inds)
+        rot_cosine = rot_cosine.view(batch, self.max_num, 1)
+        rot = torch.atan2(rot_sine, rot_cosine)
+
+        # height in the bev
+        hei = self._transpose_and_gather_feat(hei, inds)
+        hei = hei.view(batch, self.max_num, 1)
+
+        # dim of the box
+        dim = self._transpose_and_gather_feat(dim, inds)
+        dim = dim.view(batch, self.max_num, 3)
+
+        # class label
+        clses = clses.view(batch, self.max_num).float()
+        scores = scores.view(batch, self.max_num)
+
+        xs = xs.view(
+            batch, self.max_num,
+            1) * self.out_size_factor * self.voxel_size[0] + self.pc_range[0]
+        ys = ys.view(
+            batch, self.max_num,
+            1) * self.out_size_factor * self.voxel_size[1] + self.pc_range[1]
+
+        if vel is None:  # KITTI FORMAT
+            final_box_preds = torch.cat([xs, ys, hei, dim, rot], dim=2)
+        else:  # exist velocity, nuscene format
+            vel = self._transpose_and_gather_feat(vel, inds)
+            vel = vel.view(batch, self.max_num, 2)
+            final_box_preds = torch.cat([xs, ys, hei, dim, rot, vel], dim=2)
+
+        final_scores = scores
+        final_preds = clses
+
+        # use score threshold
+        if self.score_threshold is not None:
+            thresh_mask = final_scores > self.score_threshold
+
+        if self.post_center_range is not None:
+            self.post_center_range = torch.tensor(
+                self.post_center_range, device=heat.device)
+            mask = (final_box_preds[..., :3] >=
+                    self.post_center_range[:3]).all(2)
+            mask &= (final_box_preds[..., :3] <=
+                     self.post_center_range[3:]).all(2)
+
+            predictions_dicts = []
+            for i in range(batch):
+                cmask = mask[i, :]
+                if self.score_threshold:
+                    cmask &= thresh_mask[i]
+
+                boxes3d = final_box_preds[i, cmask]
+                scores = final_scores[i, cmask]
+                labels = final_preds[i, cmask]
+                predictions_dict = {
+                    'bboxes': boxes3d,
+                    'scores': scores,
+                    'labels': labels
+                }
+
+                predictions_dicts.append(predictions_dict)
+        else:
+            raise NotImplementedError(
+                'Need to reorganize output as a batch, only '
+                'support post_center_range is not None for now!')
+
+        return predictions_dicts

mmdet3d/core/post_processing/__init__.py

 from mmdet.core.post_processing import (merge_aug_bboxes, merge_aug_masks,
                                         merge_aug_proposals, merge_aug_scores,
                                         multiclass_nms)
-from .box3d_nms import aligned_3d_nms, box3d_multiclass_nms
+from .box3d_nms import aligned_3d_nms, box3d_multiclass_nms, circle_nms
 from .merge_augs import merge_aug_bboxes_3d
 
 __all__ = [
     'multiclass_nms', 'merge_aug_proposals', 'merge_aug_bboxes',
     'merge_aug_scores', 'merge_aug_masks', 'box3d_multiclass_nms',
-    'aligned_3d_nms', 'merge_aug_bboxes_3d'
+    'aligned_3d_nms', 'merge_aug_bboxes_3d', 'circle_nms'
 ]

mmdet3d/core/post_processing/box3d_nms.py

+import numba
+import numpy as np
 import torch
 
 from mmdet3d.ops.iou3d.iou3d_utils import nms_gpu, nms_normal_gpu
@@ -134,3 +136,46 @@ def aligned_3d_nms(boxes, scores, classes, thresh):
 
     indices = boxes.new_tensor(pick, dtype=torch.long)
     return indices
+
+
+@numba.jit(nopython=True)
+def circle_nms(dets, thresh, post_max_size=83):
+    """Circular NMS.
+
+    An object is only counted as positive if no other center
+    with a higher confidence exists within a radius r using a
+    bird-eye view distance metric.
+
+    Args:
+        dets (torch.Tensor): Detection results with the shape of [N, 3].
+        thresh (float): Value of threshold.
+        post_max_size (int): Max number of prediction to be kept. Defaults
+            to 83
+
+    Returns:
+        torch.Tensor: Indexes of the detections to be kept.
+    """
+    x1 = dets[:, 0]
+    y1 = dets[:, 1]
+    scores = dets[:, 2]
+    order = scores.argsort()[::-1].astype(np.int32)  # highest->lowest
+    ndets = dets.shape[0]
+    suppressed = np.zeros((ndets), dtype=np.int32)
+    keep = []
+    for _i in range(ndets):
+        i = order[_i]  # start with highest score box
+        if suppressed[
+                i] == 1:  # if any box have enough iou with this, remove it
+            continue
+        keep.append(i)
+        for _j in range(_i + 1, ndets):
+            j = order[_j]
+            if suppressed[j] == 1:
+                continue
+            # calculate center distance between i and j box
+            dist = (x1[i] - x1[j])**2 + (y1[i] - y1[j])**2
+
+            # ovr = inter / areas[j]
+            if dist <= thresh:
+                suppressed[j] = 1
+    return keep[:post_max_size]

mmdet3d/core/utils/__init__.py

+from .gaussian import draw_heatmap_gaussian, gaussian_2d, gaussian_radius
+
+__all__ = ['gaussian_2d', 'gaussian_radius', 'draw_heatmap_gaussian']

mmdet3d/core/utils/gaussian.py

+import numpy as np
+import torch
+
+
+def gaussian_2d(shape, sigma=1):
+    """Generate gaussian map.
+
+    Args:
+        shape (list[int]): Shape of the map.
+        sigma (float): Sigma to generate gaussian map.
+            Defaults to 1.
+
+    Returns:
+        np.ndarray: Generated gaussian map.
+    """
+    m, n = [(ss - 1.) / 2. for ss in shape]
+    y, x = np.ogrid[-m:m + 1, -n:n + 1]
+
+    h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
+    h[h < np.finfo(h.dtype).eps * h.max()] = 0
+    return h
+
+
+def draw_heatmap_gaussian(heatmap, center, radius, k=1):
+    """Get gaussian masked heatmap.
+
+    Args:
+        heatmap (torch.Tensor): Heatmap to be masked.
+        center (torch.Tensor): Center coord of the heatmap.
+        radius (int): Radius of gausian.
+        K (int): Multiple of masked_gaussian. Defaults to 1.
+
+    Returns:
+        torch.Tensor: Masked heatmap.
+    """
+    diameter = 2 * radius + 1
+    gaussian = gaussian_2d((diameter, diameter), sigma=diameter / 6)
+
+    x, y = int(center[0]), int(center[1])
+
+    height, width = heatmap.shape[0:2]
+
+    left, right = min(x, radius), min(width - x, radius + 1)
+    top, bottom = min(y, radius), min(height - y, radius + 1)
+
+    masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right]
+    masked_gaussian = torch.from_numpy(
+        gaussian[radius - top:radius + bottom,
+                 radius - left:radius + right]).to(heatmap.device,
+                                                   torch.float32)
+    if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0:
+        torch.max(masked_heatmap, masked_gaussian * k, out=masked_heatmap)
+    return heatmap
+
+
+def gaussian_radius(det_size, min_overlap=0.5):
+    """Get radius of gaussian.
+
+    Args:
+        det_size (tuple[torch.Tensor]): Size of the detection result.
+        min_overlap (float): Gaussian_overlap. Defaults to 0.5.
+
+    Returns:
+        torch.Tensor: Computed radius.
+    """
+    height, width = det_size
+
+    a1 = 1
+    b1 = (height + width)
+    c1 = width * height * (1 - min_overlap) / (1 + min_overlap)
+    sq1 = torch.sqrt(b1**2 - 4 * a1 * c1)
+    r1 = (b1 + sq1) / 2
+
+    a2 = 4
+    b2 = 2 * (height + width)
+    c2 = (1 - min_overlap) * width * height
+    sq2 = torch.sqrt(b2**2 - 4 * a2 * c2)
+    r2 = (b2 + sq2) / 2
+
+    a3 = 4 * min_overlap
+    b3 = -2 * min_overlap * (height + width)
+    c3 = (min_overlap - 1) * width * height
+    sq3 = torch.sqrt(b3**2 - 4 * a3 * c3)
+    r3 = (b3 + sq3) / 2
+    return min(r1, r2, r3)

mmdet3d/models/roi_heads/bbox_heads/__init__.py

@@ -3,10 +3,11 @@ from mmdet.models.roi_heads.bbox_heads import (BBoxHead, ConvFCBBoxHead,
                                                Shared2FCBBoxHead,
                                                Shared4Conv1FCBBoxHead)
 from .h3d_bbox_head import H3DBboxHead
+from .multi_group_head import CenterHead
 from .parta2_bbox_head import PartA2BboxHead
 
 __all__ = [
     'BBoxHead', 'ConvFCBBoxHead', 'Shared2FCBBoxHead',
-    'Shared4Conv1FCBBoxHead', 'DoubleConvFCBBoxHead', 'H3DBboxHead',
-    'PartA2BboxHead'
+    'Shared4Conv1FCBBoxHead', 'DoubleConvFCBBoxHead', 'PartA2BboxHead',
+    'H3DBboxHead', 'CenterHead'
 ]

mmdet3d/models/utils/__init__.py

+from .clip_sigmoid import clip_sigmoid
+
+__all__ = ['clip_sigmoid']

mmdet3d/models/utils/clip_sigmoid.py

+import torch
+
+
+def clip_sigmoid(x, eps=1e-4):
+    """Sigmoid function for input feature.
+
+    Args:
+        x (torch.Tensor): Input feature map with the shape of [B, N, H, W].
+        eps (float): Lower bound of the range to be clamped to. Defaults
+            to 1e-4.
+
+    Returns:
+        torch.Tensor: Feature map after sigmoid.
+    """
+    y = torch.clamp(x.sigmoid_(), min=eps, max=1 - eps)
+    return y

mmdet3d/ops/iou3d/iou3d_utils.py

@@ -22,24 +22,32 @@ def boxes_iou_bev(boxes_a, boxes_b):
     return ans_iou
 
 
-def nms_gpu(boxes, scores, thresh):
-    """Non maximum suppression on GPU.
+def nms_gpu(boxes, scores, thresh, pre_maxsize=None, post_max_size=None):
+    """Nms function with gpu implementation.
 
     Args:
-        boxes (torch.Tensor): Input boxes with shape (N, 5).
-        scores (torch.Tensor): Scores of predicted boxes with shape (N).
-        thresh (torch.Tensor): Threshold of non maximum suppression.
+        boxes (torch.Tensor): Input boxes with the shape of [N, 5]
+            ([x1, y1, x2, y2, ry]).
+        scores (torch.Tensor): Scores of boxes with the shape of [N].
+        thresh (int): Threshold.
+        pre_maxsize (int): Max size of boxes before nms. Default: None.
+        post_maxsize (int): Max size of boxes after nms. Default: None.
 
     Returns:
-        torch.Tensor: Remaining indices with scores in descending order.
+        torch.Tensor: Indexes after nms.
     """
     order = scores.sort(0, descending=True)[1]
 
+    if pre_maxsize is not None:
+        order = order[:pre_maxsize]
     boxes = boxes[order].contiguous()
 
     keep = torch.zeros(boxes.size(0), dtype=torch.long)
     num_out = iou3d_cuda.nms_gpu(boxes, keep, thresh, boxes.device.index)
-    return order[keep[:num_out].cuda(boxes.device)].contiguous()
+    keep = order[keep[:num_out].cuda(boxes.device)].contiguous()
+    if post_max_size is not None:
+        keep = keep[:post_max_size]
+    return keep
 
 
 def nms_normal_gpu(boxes, scores, thresh):

tests/test_box_coders.py → tests/test_bbox_coders.py

@@ -323,3 +323,31 @@ def test_anchor_free_box_coder():
     assert dir_res_norm.shape == torch.Size([2, 256, 12])
     assert dir_res.shape == torch.Size([2, 256, 12])
     assert size.shape == torch.Size([2, 256, 3])
+
+
+def test_centerpoint_bbox_coder():
+    bbox_coder_cfg = dict(
+        type='CenterPointBBoxCoder',
+        post_center_range=[-61.2, -61.2, -10.0, 61.2, 61.2, 10.0],
+        max_num=500,
+        score_threshold=0.1,
+        pc_range=[-51.2, -51.2],
+        out_size_factor=4,
+        voxel_size=[0.2, 0.2])
+
+    bbox_coder = build_bbox_coder(bbox_coder_cfg)
+
+    batch_dim = torch.rand([2, 3, 128, 128])
+    batch_hei = torch.rand([2, 1, 128, 128])
+    batch_hm = torch.rand([2, 2, 128, 128])
+    batch_reg = torch.rand([2, 2, 128, 128])
+    batch_rotc = torch.rand([2, 1, 128, 128])
+    batch_rots = torch.rand([2, 1, 128, 128])
+    batch_vel = torch.rand([2, 2, 128, 128])
+
+    temp = bbox_coder.decode(batch_hm, batch_rots, batch_rotc, batch_hei,
+                             batch_dim, batch_vel, batch_reg, 5)
+    for i in range(len(temp)):
+        assert temp[i]['bboxes'].shape == torch.Size([500, 9])
+        assert temp[i]['scores'].shape == torch.Size([500])
+        assert temp[i]['labels'].shape == torch.Size([500])

tests/test_heads.py

@@ -8,6 +8,7 @@ from os.path import dirname, exists, join
 from mmdet3d.core.bbox import (Box3DMode, DepthInstance3DBoxes,
                                LiDARInstance3DBoxes)
 from mmdet3d.models.builder import build_head
+from mmdet.apis import set_random_seed
 
 
 def _setup_seed(seed):
@@ -689,6 +690,199 @@ def test_h3d_head():
     assert ret_dict['primitive_sem_matching_loss'] >= 0
 
 
+def test_center_head():
+    tasks = [
+        dict(num_class=1, class_names=['car']),
+        dict(num_class=2, class_names=['truck', 'construction_vehicle']),
+        dict(num_class=2, class_names=['bus', 'trailer']),
+        dict(num_class=1, class_names=['barrier']),
+        dict(num_class=2, class_names=['motorcycle', 'bicycle']),
+        dict(num_class=2, class_names=['pedestrian', 'traffic_cone']),
+    ]
+    bbox_cfg = dict(
+        type='CenterPointBBoxCoder',
+        post_center_range=[-61.2, -61.2, -10.0, 61.2, 61.2, 10.0],
+        max_num=500,
+        score_threshold=0.1,
+        pc_range=[-51.2, -51.2],
+        out_size_factor=8,
+        voxel_size=[0.2, 0.2])
+    train_cfg = dict(
+        grid_size=[1024, 1024, 40],
+        point_cloud_range=[-51.2, -51.2, -5., 51.2, 51.2, 3.],
+        voxel_size=[0.1, 0.1, 0.2],
+        out_size_factor=8,
+        dense_reg=1,
+        gaussian_overlap=0.1,
+        max_objs=500,
+        code_weights=[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.2, 0.2, 1.0, 1.0],
+        min_radius=2)
+    test_cfg = dict(
+        post_center_limit_range=[-61.2, -61.2, -10.0, 61.2, 61.2, 10.0],
+        max_per_img=500,
+        max_pool_nms=False,
+        min_radius=[4, 12, 10, 1, 0.85, 0.175],
+        post_max_size=83,
+        score_threshold=0.1,
+        pc_range=[-51.2, -51.2],
+        out_size_factor=8,
+        voxel_size=[0.2, 0.2],
+        nms_type='circle')
+    center_head_cfg = dict(
+        type='CenterHead',
+        in_channels=sum([256, 256]),
+        tasks=tasks,
+        train_cfg=train_cfg,
+        test_cfg=test_cfg,
+        bbox_coder=bbox_cfg,
+        common_heads=dict(
+            reg=(2, 2), height=(1, 2), dim=(3, 2), rot=(2, 2), vel=(2, 2)),
+        share_conv_channel=64,
+        norm_bbox=True)
+
+    center_head = build_head(center_head_cfg)
+
+    x = torch.rand([2, 512, 128, 128])
+    output = center_head([x])
+    for i in range(6):
+        assert output[i][0]['reg'].shape == torch.Size([2, 2, 128, 128])
+        assert output[i][0]['height'].shape == torch.Size([2, 1, 128, 128])
+        assert output[i][0]['dim'].shape == torch.Size([2, 3, 128, 128])
+        assert output[i][0]['rot'].shape == torch.Size([2, 2, 128, 128])
+        assert output[i][0]['vel'].shape == torch.Size([2, 2, 128, 128])
+        assert output[i][0]['heatmap'].shape == torch.Size(
+            [2, tasks[i]['num_class'], 128, 128])
+
+    # test get_bboxes
+    img_metas = [
+        dict(box_type_3d=LiDARInstance3DBoxes),
+        dict(box_type_3d=LiDARInstance3DBoxes)
+    ]
+    ret_lists = center_head.get_bboxes(output, img_metas)
+    for ret_list in ret_lists:
+        assert ret_list[0].tensor.shape[0] <= 500
+        assert ret_list[1].shape[0] <= 500
+        assert ret_list[2].shape[0] <= 500
+
+
+def test_dcn_center_head():
+    if not torch.cuda.is_available():
+        pytest.skip('test requires GPU and CUDA')
+    set_random_seed(0)
+    tasks = [
+        dict(num_class=1, class_names=['car']),
+        dict(num_class=2, class_names=['truck', 'construction_vehicle']),
+        dict(num_class=2, class_names=['bus', 'trailer']),
+        dict(num_class=1, class_names=['barrier']),
+        dict(num_class=2, class_names=['motorcycle', 'bicycle']),
+        dict(num_class=2, class_names=['pedestrian', 'traffic_cone']),
+    ]
+    voxel_size = [0.2, 0.2, 8]
+    dcn_center_head_cfg = dict(
+        type='CenterHead',
+        mode='3d',
+        in_channels=sum([128, 128, 128]),
+        tasks=[
+            dict(num_class=1, class_names=['car']),
+            dict(num_class=2, class_names=['truck', 'construction_vehicle']),
+            dict(num_class=2, class_names=['bus', 'trailer']),
+            dict(num_class=1, class_names=['barrier']),
+            dict(num_class=2, class_names=['motorcycle', 'bicycle']),
+            dict(num_class=2, class_names=['pedestrian', 'traffic_cone']),
+        ],
+        common_heads={
+            'reg': (2, 2),
+            'height': (1, 2),
+            'dim': (3, 2),
+            'rot': (2, 2),
+            'vel': (2, 2)
+        },
+        share_conv_channel=64,
+        bbox_coder=dict(
+            type='CenterPointBBoxCoder',
+            post_center_range=[-61.2, -61.2, -10.0, 61.2, 61.2, 10.0],
+            max_num=500,
+            score_threshold=0.1,
+            pc_range=[-51.2, -51.2],
+            out_size_factor=4,
+            voxel_size=voxel_size[:2],
+            code_size=9),
+        seperate_head=dict(
+            type='DCNSeperateHead',
+            dcn_config=dict(
+                type='DCN',
+                in_channels=64,
+                out_channels=64,
+                kernel_size=3,
+                padding=1,
+                groups=4,
+                bias=True),
+            init_bias=-2.19,
+            final_kernel=3),
+        loss_cls=dict(type='GaussianFocalLoss', reduction='mean'),
+        loss_bbox=dict(type='L1Loss', reduction='none', loss_weight=0.25),
+        norm_bbox=True)
+    # model training and testing settings
+    train_cfg = dict(
+        grid_size=[512, 512, 1],
+        point_cloud_range=[-51.2, -51.2, -5., 51.2, 51.2, 3.],
+        voxel_size=voxel_size,
+        out_size_factor=4,
+        dense_reg=1,
+        gaussian_overlap=0.1,
+        max_objs=500,
+        min_radius=2,
+        code_weights=[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.2, 0.2, 1.0, 1.0])
+
+    test_cfg = dict(
+        post_center_limit_range=[-61.2, -61.2, -10.0, 61.2, 61.2, 10.0],
+        max_per_img=500,
+        max_pool_nms=False,
+        min_radius=[4, 12, 10, 1, 0.85, 0.175],
+        post_max_size=83,
+        score_threshold=0.1,
+        pc_range=[-51.2, -51.2],
+        out_size_factor=4,
+        voxel_size=voxel_size[:2],
+        nms_type='circle')
+    dcn_center_head_cfg.update(train_cfg=train_cfg, test_cfg=test_cfg)
+
+    dcn_center_head = build_head(dcn_center_head_cfg).cuda()
+
+    x = torch.ones([2, 384, 128, 128]).cuda()
+    output = dcn_center_head([x])
+    for i in range(6):
+        assert output[i][0]['reg'].shape == torch.Size([2, 2, 128, 128])
+        assert output[i][0]['height'].shape == torch.Size([2, 1, 128, 128])
+        assert output[i][0]['dim'].shape == torch.Size([2, 3, 128, 128])
+        assert output[i][0]['rot'].shape == torch.Size([2, 2, 128, 128])
+        assert output[i][0]['vel'].shape == torch.Size([2, 2, 128, 128])
+        assert output[i][0]['heatmap'].shape == torch.Size(
+            [2, tasks[i]['num_class'], 128, 128])
+
+    # Test loss.
+    gt_bboxes_0 = LiDARInstance3DBoxes(torch.rand([10, 9]).cuda(), box_dim=9)
+    gt_bboxes_1 = LiDARInstance3DBoxes(torch.rand([20, 9]).cuda(), box_dim=9)
+    gt_labels_0 = torch.randint(1, 11, [10]).cuda()
+    gt_labels_1 = torch.randint(1, 11, [20]).cuda()
+    gt_bboxes_3d = [gt_bboxes_0, gt_bboxes_1]
+    gt_labels_3d = [gt_labels_0, gt_labels_1]
+    loss = dcn_center_head.loss(gt_bboxes_3d, gt_labels_3d, output)
+    loss_sum = torch.sum(torch.stack([item for _, item in loss.items()]))
+    assert torch.isclose(loss_sum, torch.tensor(21972.1230))
+
+    # test get_bboxes
+    img_metas = [
+        dict(box_type_3d=LiDARInstance3DBoxes),
+        dict(box_type_3d=LiDARInstance3DBoxes)
+    ]
+    ret_lists = dcn_center_head.get_bboxes(output, img_metas)
+    for ret_list in ret_lists:
+        assert ret_list[0].tensor.shape[0] <= 500
+        assert ret_list[1].shape[0] <= 500
+        assert ret_list[2].shape[0] <= 500
+
+
 def test_ssd3d_head():
     if not torch.cuda.is_available():
         pytest.skip('test requires GPU and torch+cuda')

tests/test_nms.py

+import numpy as np
 import torch
 
 
@@ -55,3 +56,19 @@ def test_aligned_3d_nms():
     ])
 
     assert torch.all(pick == expected_pick)
+
+
+def test_circle_nms():
+    from mmdet3d.core.post_processing import circle_nms
+    boxes = torch.tensor([[-11.1100, 2.1300, 0.8823],
+                          [-11.2810, 2.2422, 0.8914],
+                          [-10.3966, -0.3198, 0.8643],
+                          [-10.2906, -13.3159,
+                           0.8401], [5.6518, 9.9791, 0.8271],
+                          [-11.2652, 13.3637, 0.8267],
+                          [4.7768, -13.0409, 0.7810], [5.6621, 9.0422, 0.7753],
+                          [-10.5561, 18.9627, 0.7518],
+                          [-10.5643, 13.2293, 0.7200]])
+    keep = circle_nms(boxes.numpy(), 0.175)
+    expected_keep = [1, 2, 3, 4, 5, 6, 7, 8, 9]
+    assert np.all(keep == expected_keep)

tests/test_utils.py

+import torch
+
+from mmdet3d.core import draw_heatmap_gaussian
+
+
+def test_gaussian():
+    heatmap = torch.zeros((128, 128))
+    ct_int = torch.tensor([64, 64], dtype=torch.int32)
+    radius = 2
+    draw_heatmap_gaussian(heatmap, ct_int, radius)
+    assert torch.isclose(torch.sum(heatmap), torch.tensor(4.3505), atol=1e-3)