Merge branch 'master_temp' into indoor_augment

mmdet3d/datasets/registry.py

-from mmdet.utils import Registry
+from mmcv.utils import Registry
 
 OBJECTSAMPLERS = Registry('Object sampler')

mmdet3d/models/__init__.py

 from .anchor_heads import *  # noqa: F401,F403
 from .backbones import *  # noqa: F401,F403
-from .bbox_heads import *  # noqa: F401,F403
 from .builder import (build_backbone, build_detector, build_fusion_layer,
                       build_head, build_loss, build_middle_encoder, build_neck,
                       build_roi_extractor, build_shared_head,
@@ -11,7 +10,7 @@ from .losses import *  # noqa: F401,F403
 from .middle_encoders import *  # noqa: F401,F403
 from .necks import *  # noqa: F401,F403
 from .registry import FUSION_LAYERS, MIDDLE_ENCODERS, VOXEL_ENCODERS
-from .roi_extractors import *  # noqa: F401,F403
+from .roi_heads import *  # noqa: F401,F403
 from .voxel_encoders import *  # noqa: F401,F403
 
 __all__ = [

mmdet3d/models/anchor_heads/__init__.py

 from .boxvelo_head import Anchor3DVeloHead
+from .parta2_rpn_head import PartA2RPNHead
 from .second_head import SECONDHead
 
-__all__ = ['Anchor3DVeloHead', 'SECONDHead']
+__all__ = ['Anchor3DVeloHead', 'SECONDHead', 'PartA2RPNHead']

mmdet3d/models/anchor_heads/boxvelo_head.py

 import numpy as np
 import torch
-from mmcv.cnn import normal_init
+from mmcv.cnn import bias_init_with_prob, normal_init
 
 from mmdet3d.core import box_torch_ops, boxes3d_to_bev_torch_lidar
 from mmdet3d.ops.iou3d.iou3d_utils import nms_gpu, nms_normal_gpu
 from mmdet.models import HEADS
-from ..utils import bias_init_with_prob
 from .second_head import SECONDHead
 
 
-@HEADS.register_module
+@HEADS.register_module()
 class Anchor3DVeloHead(SECONDHead):
     """Anchor-based head for 3D anchor with velocity
     Args:
         in_channels (int): Number of channels in the input feature map.
         feat_channels (int): Number of channels of the feature map.
-        anchor_scales (Iterable): Anchor scales.
-        anchor_ratios (Iterable): Anchor aspect ratios.
-        anchor_strides (Iterable): Anchor strides.
-        anchor_base_sizes (Iterable): Anchor base sizes.
-        target_means (Iterable): Mean values of regression targets.
-        target_stds (Iterable): Std values of regression targets.
         loss_cls (dict): Config of classification loss.
         loss_bbox (dict): Config of localization loss.
     """  # noqa: W605
@@ -31,25 +24,25 @@ class Anchor3DVeloHead(SECONDHead):
                  in_channels,
                  train_cfg,
                  test_cfg,
-                 cache_anchor=False,
                  feat_channels=256,
                  use_direction_classifier=True,
                  encode_bg_as_zeros=False,
                  box_code_size=9,
-                 anchor_generator=dict(type='AnchorGeneratorRange', ),
-                 anchor_range=[0, -39.68, -1.78, 69.12, 39.68, -1.78],
-                 anchor_strides=[2],
-                 anchor_sizes=[[1.6, 3.9, 1.56]],
-                 anchor_rotations=[0, 1.57],
-                 anchor_custom_values=[0, 0],
+                 anchor_generator=dict(
+                     type='Anchor3DRangeGenerator',
+                     range=[0, -39.68, -1.78, 69.12, 39.68, -1.78],
+                     strides=[2],
+                     sizes=[[1.6, 3.9, 1.56]],
+                     rotations=[0, 1.57],
+                     custom_values=[0, 0],
+                     reshape_out=True,
+                 ),
                  assigner_per_size=False,
                  assign_per_class=False,
                  diff_rad_by_sin=True,
                  dir_offset=0,
                  dir_limit_offset=1,
-                 target_means=(.0, .0, .0, .0),
-                 target_stds=(1.0, 1.0, 1.0, 1.0),
-                 bbox_coder=dict(type='Residual3DBoxCoder', ),
+                 bbox_coder=dict(type='DeltaXYZWLHRBBoxCoder'),
                  loss_cls=dict(
                      type='CrossEntropyLoss',
                      use_sigmoid=True,
@@ -58,14 +51,11 @@ class Anchor3DVeloHead(SECONDHead):
                      type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=2.0),
                  loss_dir=dict(type='CrossEntropyLoss', loss_weight=0.2)):
         super().__init__(class_names, in_channels, train_cfg, test_cfg,
-                         cache_anchor, feat_channels, use_direction_classifier,
+                         feat_channels, use_direction_classifier,
                          encode_bg_as_zeros, box_code_size, anchor_generator,
-                         anchor_range, anchor_strides, anchor_sizes,
-                         anchor_rotations, anchor_custom_values,
                          assigner_per_size, assign_per_class, diff_rad_by_sin,
-                         dir_offset, dir_limit_offset, target_means,
-                         target_stds, bbox_coder, loss_cls, loss_bbox,
-                         loss_dir)
+                         dir_offset, dir_limit_offset, bbox_coder, loss_cls,
+                         loss_bbox, loss_dir)
         self.num_classes = num_classes
         # build head layers & losses
         if not self.use_sigmoid_cls:
@@ -131,9 +121,7 @@ class Anchor3DVeloHead(SECONDHead):
                 scores = scores[topk_inds, :]
                 dir_cls_score = dir_cls_score[topk_inds]
 
-            bboxes = self.bbox_coder.decode_torch(anchors, bbox_pred,
-                                                  self.target_means,
-                                                  self.target_stds)
+            bboxes = self.bbox_coder.decode(anchors, bbox_pred)
             mlvl_bboxes.append(bboxes)
             mlvl_scores.append(scores)
             mlvl_dir_scores.append(dir_cls_score)

mmdet3d/models/anchor_heads/parta2_rpn_head.py

+from __future__ import division
+
+import numpy as np
+import torch
+
+from mmdet3d.core import box_torch_ops, boxes3d_to_bev_torch_lidar
+from mmdet3d.ops.iou3d.iou3d_utils import nms_gpu, nms_normal_gpu
+from mmdet.models import HEADS
+from .second_head import SECONDHead
+
+
+@HEADS.register_module()
+class PartA2RPNHead(SECONDHead):
+    """rpn head for PartA2
+
+    Args:
+        class_name (list[str]): name of classes (TODO: to be removed)
+        in_channels (int): Number of channels in the input feature map.
+        train_cfg (dict): train configs
+        test_cfg (dict): test configs
+        feat_channels (int): Number of channels of the feature map.
+        use_direction_classifier (bool): Whether to add a direction classifier.
+        encode_bg_as_zeros (bool): Whether to use sigmoid of softmax
+            (TODO: to be removed)
+        box_code_size (int): The size of box code.
+        anchor_generator(dict): Config dict of anchor generator.
+        assigner_per_size (bool): Whether to do assignment for each separate
+            anchor size.
+        assign_per_class (bool): Whether to do assignment for each class.
+        diff_rad_by_sin (bool): Whether to change the difference into sin
+            difference for box regression loss.
+        dir_offset (float | int): The offset of BEV rotation angles
+            (TODO: may be moved into box coder)
+        dirlimit_offset (float | int): The limited range of BEV rotation angles
+            (TODO: may be moved into box coder)
+        box_coder (dict): Config dict of box coders.
+        loss_cls (dict): Config of classification loss.
+        loss_bbox (dict): Config of localization loss.
+        loss_dir (dict): Config of direction classifier loss.
+    """  # npqa:W293
+
+    def __init__(self,
+                 class_name,
+                 in_channels,
+                 train_cfg,
+                 test_cfg,
+                 feat_channels=256,
+                 use_direction_classifier=True,
+                 encode_bg_as_zeros=False,
+                 box_code_size=7,
+                 anchor_generator=dict(
+                     type='Anchor3DRangeGenerator',
+                     range=[0, -39.68, -1.78, 69.12, 39.68, -1.78],
+                     strides=[2],
+                     sizes=[[1.6, 3.9, 1.56]],
+                     rotations=[0, 1.57],
+                     custom_values=[],
+                     reshape_out=False),
+                 assigner_per_size=False,
+                 assign_per_class=False,
+                 diff_rad_by_sin=True,
+                 dir_offset=0,
+                 dir_limit_offset=1,
+                 bbox_coder=dict(type='DeltaXYZWLHRBBoxCoder'),
+                 loss_cls=dict(
+                     type='CrossEntropyLoss',
+                     use_sigmoid=True,
+                     loss_weight=1.0),
+                 loss_bbox=dict(
+                     type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=2.0),
+                 loss_dir=dict(type='CrossEntropyLoss', loss_weight=0.2)):
+        super().__init__(class_name, in_channels, train_cfg, test_cfg,
+                         feat_channels, use_direction_classifier,
+                         encode_bg_as_zeros, box_code_size, anchor_generator,
+                         assigner_per_size, assign_per_class, diff_rad_by_sin,
+                         dir_offset, dir_limit_offset, bbox_coder, loss_cls,
+                         loss_bbox, loss_dir)
+
+    def get_bboxes(self,
+                   cls_scores,
+                   bbox_preds,
+                   dir_cls_preds,
+                   input_metas,
+                   cfg,
+                   rescale=False):
+        assert len(cls_scores) == len(bbox_preds)
+        assert len(cls_scores) == len(dir_cls_preds)
+        num_levels = len(cls_scores)
+        featmap_sizes = [cls_scores[i].shape[-2:] for i in range(num_levels)]
+        device = cls_scores[0].device
+        mlvl_anchors = self.anchor_generator.grid_anchors(
+            featmap_sizes, device=device)
+        mlvl_anchors = [
+            anchor.reshape(-1, self.box_code_size) for anchor in mlvl_anchors
+        ]
+
+        result_list = []
+        for img_id in range(len(input_metas)):
+            cls_score_list = [
+                cls_scores[i][img_id].detach() for i in range(num_levels)
+            ]
+            bbox_pred_list = [
+                bbox_preds[i][img_id].detach() for i in range(num_levels)
+            ]
+            dir_cls_pred_list = [
+                dir_cls_preds[i][img_id].detach() for i in range(num_levels)
+            ]
+
+            input_meta = input_metas[img_id]
+            proposals = self.get_bboxes_single(cls_score_list, bbox_pred_list,
+                                               dir_cls_pred_list, mlvl_anchors,
+                                               input_meta, cfg, rescale)
+            result_list.append(proposals)
+        return result_list
+
+    def get_bboxes_single(self,
+                          cls_scores,
+                          bbox_preds,
+                          dir_cls_preds,
+                          mlvl_anchors,
+                          input_meta,
+                          cfg,
+                          rescale=False):
+        assert len(cls_scores) == len(bbox_preds) == len(mlvl_anchors)
+        mlvl_bboxes = []
+        mlvl_max_scores = []
+        mlvl_label_pred = []
+        mlvl_dir_scores = []
+        mlvl_cls_score = []
+        for cls_score, bbox_pred, dir_cls_pred, anchors in zip(
+                cls_scores, bbox_preds, dir_cls_preds, mlvl_anchors):
+            assert cls_score.size()[-2:] == bbox_pred.size()[-2:]
+            assert cls_score.size()[-2:] == dir_cls_pred.size()[-2:]
+            dir_cls_pred = dir_cls_pred.permute(1, 2, 0).reshape(-1, 2)
+            dir_cls_score = torch.max(dir_cls_pred, dim=-1)[1]
+
+            cls_score = cls_score.permute(1, 2,
+                                          0).reshape(-1, self.num_classes)
+
+            if self.use_sigmoid_cls:
+                scores = cls_score.sigmoid()
+            else:
+                scores = cls_score.softmax(-1)
+            bbox_pred = bbox_pred.permute(1, 2,
+                                          0).reshape(-1, self.box_code_size)
+
+            nms_pre = cfg.get('nms_pre', -1)
+            if self.use_sigmoid_cls:
+                max_scores, pred_labels = scores.max(dim=1)
+            else:
+                max_scores, pred_labels = scores[:, :-1].max(dim=1)
+            # get topk
+            if nms_pre > 0 and scores.shape[0] > nms_pre:
+                topk_scores, topk_inds = max_scores.topk(nms_pre)
+                anchors = anchors[topk_inds, :]
+                bbox_pred = bbox_pred[topk_inds, :]
+                max_scores = topk_scores
+                cls_score = cls_score[topk_inds, :]
+                dir_cls_score = dir_cls_score[topk_inds]
+                pred_labels = pred_labels[topk_inds]
+
+            bboxes = self.bbox_coder.decode(anchors, bbox_pred)
+            mlvl_bboxes.append(bboxes)
+            mlvl_max_scores.append(max_scores)
+            mlvl_cls_score.append(cls_score)
+            mlvl_label_pred.append(pred_labels)
+            mlvl_dir_scores.append(dir_cls_score)
+
+        mlvl_bboxes = torch.cat(mlvl_bboxes)
+        mlvl_bboxes_for_nms = boxes3d_to_bev_torch_lidar(mlvl_bboxes)
+        mlvl_max_scores = torch.cat(mlvl_max_scores)
+        mlvl_label_pred = torch.cat(mlvl_label_pred)
+        mlvl_dir_scores = torch.cat(mlvl_dir_scores)
+        mlvl_cls_score = torch.cat(
+            mlvl_cls_score)  # shape [k, num_class] before sigmoid
+
+        score_thr = cfg.get('score_thr', 0)
+        result = self.class_agnostic_nms(mlvl_bboxes, mlvl_bboxes_for_nms,
+                                         mlvl_max_scores, mlvl_label_pred,
+                                         mlvl_cls_score, mlvl_dir_scores,
+                                         score_thr, cfg.nms_post, cfg)
+
+        result.update(dict(sample_idx=input_meta['sample_idx']))
+        return result
+
+    def class_agnostic_nms(self, mlvl_bboxes, mlvl_bboxes_for_nms,
+                           mlvl_max_scores, mlvl_label_pred, mlvl_cls_score,
+                           mlvl_dir_scores, score_thr, max_num, cfg):
+        bboxes = []
+        scores = []
+        labels = []
+        dir_scores = []
+        cls_scores = []
+        score_thr_inds = mlvl_max_scores > score_thr
+        _scores = mlvl_max_scores[score_thr_inds]
+        _bboxes_for_nms = mlvl_bboxes_for_nms[score_thr_inds, :]
+        if cfg.use_rotate_nms:
+            nms_func = nms_gpu
+        else:
+            nms_func = nms_normal_gpu
+        selected = nms_func(_bboxes_for_nms, _scores, cfg.nms_thr)
+
+        _mlvl_bboxes = mlvl_bboxes[score_thr_inds, :]
+        _mlvl_dir_scores = mlvl_dir_scores[score_thr_inds]
+        _mlvl_label_pred = mlvl_label_pred[score_thr_inds]
+        _mlvl_cls_score = mlvl_cls_score[score_thr_inds]
+
+        if len(selected) > 0:
+            bboxes.append(_mlvl_bboxes[selected])
+            scores.append(_scores[selected])
+            labels.append(_mlvl_label_pred[selected])
+            cls_scores.append(_mlvl_cls_score[selected])
+            dir_scores.append(_mlvl_dir_scores[selected])
+            dir_rot = box_torch_ops.limit_period(
+                bboxes[-1][..., 6] - self.dir_offset, self.dir_limit_offset,
+                np.pi)
+            bboxes[-1][..., 6] = (
+                dir_rot + self.dir_offset +
+                np.pi * dir_scores[-1].to(bboxes[-1].dtype))
+
+        if bboxes:
+            bboxes = torch.cat(bboxes, dim=0)
+            scores = torch.cat(scores, dim=0)
+            cls_scores = torch.cat(cls_scores, dim=0)
+            labels = torch.cat(labels, dim=0)
+            dir_scores = torch.cat(dir_scores, dim=0)
+            if bboxes.shape[0] > max_num:
+                _, inds = scores.sort(descending=True)
+                inds = inds[:max_num]
+                bboxes = bboxes[inds, :]
+                labels = labels[inds]
+                scores = scores[inds]
+                cls_scores = cls_scores[inds]
+                dir_scores = dir_scores[inds]
+            return dict(
+                box3d_lidar=bboxes.cpu(),
+                scores=scores.cpu(),
+                label_preds=labels.cpu(),
+                cls_preds=cls_scores.cpu(
+                )  # raw scores with shape [max_num, cls_num]
+            )
+        else:
+            return dict(
+                box3d_lidar=mlvl_bboxes.new_zeros([0,
+                                                   self.box_code_size]).cpu(),
+                scores=mlvl_bboxes.new_zeros([0]).cpu(),
+                label_preds=mlvl_bboxes.new_zeros([0]).cpu(),
+                cls_preds=mlvl_bboxes.new_zeros([0, mlvl_cls_score.shape[-1]
+                                                 ]).cpu())

mmdet3d/models/anchor_heads/second_head.py

-from __future__ import division
-
 import numpy as np
 import torch
 import torch.nn as nn
-from mmcv.cnn import normal_init
+from mmcv.cnn import bias_init_with_prob, normal_init
 
 from mmdet3d.core import (PseudoSampler, box_torch_ops,
                           boxes3d_to_bev_torch_lidar, build_anchor_generator,
@@ -12,24 +10,37 @@ from mmdet3d.core import (PseudoSampler, box_torch_ops,
 from mmdet3d.ops.iou3d.iou3d_utils import nms_gpu, nms_normal_gpu
 from mmdet.models import HEADS
 from ..builder import build_loss
-from ..utils import bias_init_with_prob
 from .train_mixins import AnchorTrainMixin
 
 
-@HEADS.register_module
+@HEADS.register_module()
 class SECONDHead(nn.Module, AnchorTrainMixin):
-    """Anchor-based head (RPN, RetinaNet, SSD, etc.).
+    """Anchor-based head for VoxelNet detectors.
+
     Args:
+        class_name (list[str]): name of classes (TODO: to be removed)
         in_channels (int): Number of channels in the input feature map.
+        train_cfg (dict): train configs
+        test_cfg (dict): test configs
         feat_channels (int): Number of channels of the feature map.
-        anchor_scales (Iterable): Anchor scales.
-        anchor_ratios (Iterable): Anchor aspect ratios.
-        anchor_strides (Iterable): Anchor strides.
-        anchor_base_sizes (Iterable): Anchor base sizes.
-        target_means (Iterable): Mean values of regression targets.
-        target_stds (Iterable): Std values of regression targets.
+        use_direction_classifier (bool): Whether to add a direction classifier.
+        encode_bg_as_zeros (bool): Whether to use sigmoid of softmax
+            (TODO: to be removed)
+        box_code_size (int): The size of box code.
+        anchor_generator(dict): Config dict of anchor generator.
+        assigner_per_size (bool): Whether to do assignment for each separate
+            anchor size.
+        assign_per_class (bool): Whether to do assignment for each class.
+        diff_rad_by_sin (bool): Whether to change the difference into sin
+            difference for box regression loss.
+        dir_offset (float | int): The offset of BEV rotation angles
+            (TODO: may be moved into box coder)
+        dirlimit_offset (float | int): The limited range of BEV rotation angles
+            (TODO: may be moved into box coder)
+        box_coder (dict): Config dict of box coders.
         loss_cls (dict): Config of classification loss.
         loss_bbox (dict): Config of localization loss.
+        loss_dir (dict): Config of direction classifier loss.
     """  # noqa: W605
 
     def __init__(self,
@@ -37,25 +48,24 @@ class SECONDHead(nn.Module, AnchorTrainMixin):
                  in_channels,
                  train_cfg,
                  test_cfg,
-                 cache_anchor=False,
                  feat_channels=256,
                  use_direction_classifier=True,
                  encode_bg_as_zeros=False,
                  box_code_size=7,
-                 anchor_generator=dict(type='AnchorGeneratorRange'),
-                 anchor_range=[0, -39.68, -1.78, 69.12, 39.68, -1.78],
-                 anchor_strides=[2],
-                 anchor_sizes=[[1.6, 3.9, 1.56]],
-                 anchor_rotations=[0, 1.57],
-                 anchor_custom_values=[],
+                 anchor_generator=dict(
+                     type='Anchor3DRangeGenerator',
+                     range=[0, -39.68, -1.78, 69.12, 39.68, -1.78],
+                     strides=[2],
+                     sizes=[[1.6, 3.9, 1.56]],
+                     rotations=[0, 1.57],
+                     custom_values=[],
+                     reshape_out=False),
                  assigner_per_size=False,
                  assign_per_class=False,
                  diff_rad_by_sin=True,
                  dir_offset=0,
                  dir_limit_offset=1,
-                 target_means=(.0, .0, .0, .0),
-                 target_stds=(1.0, 1.0, 1.0, 1.0),
-                 bbox_coder=dict(type='Residual3DBoxCoder'),
+                 bbox_coder=dict(type='DeltaXYZWLHRBBoxCoder'),
                  loss_cls=dict(
                      type='CrossEntropyLoss',
                      use_sigmoid=True,
@@ -94,29 +104,9 @@ class SECONDHead(nn.Module, AnchorTrainMixin):
                 ]
 
         # build anchor generator
-        self.anchor_range = anchor_range
-        self.anchor_rotations = anchor_rotations
-        self.anchor_strides = anchor_strides
-        self.anchor_sizes = anchor_sizes
-        self.target_means = target_means
-        self.target_stds = target_stds
-        self.anchor_generators = []
+        self.anchor_generator = build_anchor_generator(anchor_generator)
         # In 3D detection, the anchor stride is connected with anchor size
-        self.num_anchors = (
-            len(self.anchor_rotations) * len(self.anchor_sizes))
-        # if len(self.anchor_sizes) != self.anchor_strides:
-        #     # this means different anchor in the same anchor strides
-        #     anchor_sizes = [self.anchor_sizes]
-        for anchor_stride in self.anchor_strides:
-            anchor_generator.update(
-                anchor_ranges=anchor_range,
-                sizes=self.anchor_sizes,
-                stride=anchor_stride,
-                rotations=anchor_rotations,
-                custom_values=anchor_custom_values,
-                cache_anchor=cache_anchor)
-            self.anchor_generators.append(
-                build_anchor_generator(anchor_generator))
+        self.num_anchors = self.anchor_generator.num_base_anchors
 
         self._init_layers()
         self.use_sigmoid_cls = loss_cls.get('use_sigmoid', False)
@@ -152,7 +142,7 @@ class SECONDHead(nn.Module, AnchorTrainMixin):
     def forward(self, feats):
         return multi_apply(self.forward_single, feats)
 
-    def get_anchors(self, featmap_sizes, input_metas):
+    def get_anchors(self, featmap_sizes, input_metas, device='cuda'):
         """Get anchors according to feature map sizes.
         Args:
             featmap_sizes (list[tuple]): Multi-level feature map sizes.
@@ -161,16 +151,10 @@ class SECONDHead(nn.Module, AnchorTrainMixin):
             tuple: anchors of each image, valid flags of each image
         """
         num_imgs = len(input_metas)
-        num_levels = len(featmap_sizes)
-
         # since feature map sizes of all images are the same, we only compute
         # anchors for one time
-        multi_level_anchors = []
-        for i in range(num_levels):
-            anchors = self.anchor_generators[i].grid_anchors(featmap_sizes[i])
-            if not self.assigner_per_size:
-                anchors = anchors.reshape(-1, anchors.size(-1))
-            multi_level_anchors.append(anchors)
+        multi_level_anchors = self.anchor_generator.grid_anchors(
+            featmap_sizes, device=device)
         anchor_list = [multi_level_anchors for _ in range(num_imgs)]
         return anchor_list
 
@@ -237,16 +221,15 @@ class SECONDHead(nn.Module, AnchorTrainMixin):
              input_metas,
              gt_bboxes_ignore=None):
         featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]
-        assert len(featmap_sizes) == len(self.anchor_generators)
-
-        anchor_list = self.get_anchors(featmap_sizes, input_metas)
+        assert len(featmap_sizes) == self.anchor_generator.num_levels
+        device = cls_scores[0].device
+        anchor_list = self.get_anchors(
+            featmap_sizes, input_metas, device=device)
         label_channels = self.cls_out_channels if self.use_sigmoid_cls else 1
         cls_reg_targets = self.anchor_target_3d(
             anchor_list,
             gt_bboxes,
             input_metas,
-            self.target_means,
-            self.target_stds,
             gt_bboxes_ignore_list=gt_bboxes_ignore,
             gt_labels_list=gt_labels,
             num_classes=self.num_classes,
@@ -288,12 +271,14 @@ class SECONDHead(nn.Module, AnchorTrainMixin):
         assert len(cls_scores) == len(bbox_preds)
         assert len(cls_scores) == len(dir_cls_preds)
         num_levels = len(cls_scores)
-
+        featmap_sizes = [cls_scores[i].shape[-2:] for i in range(num_levels)]
+        device = cls_scores[0].device
+        mlvl_anchors = self.anchor_generator.grid_anchors(
+            featmap_sizes, device=device)
         mlvl_anchors = [
-            self.anchor_generators[i].grid_anchors(
-                cls_scores[i].size()[-2:]).reshape(-1, self.box_code_size)
-            for i in range(num_levels)
+            anchor.reshape(-1, self.box_code_size) for anchor in mlvl_anchors
         ]
+
         result_list = []
         for img_id in range(len(input_metas)):
             cls_score_list = [
@@ -353,9 +338,7 @@ class SECONDHead(nn.Module, AnchorTrainMixin):
                 bbox_pred = bbox_pred[thr_inds]
                 scores = scores[thr_inds]
                 dir_cls_scores = dir_cls_score[thr_inds]
-            bboxes = self.bbox_coder.decode_torch(anchors, bbox_pred,
-                                                  self.target_means,
-                                                  self.target_stds)
+            bboxes = self.bbox_coder.decode(anchors, bbox_pred)
             bboxes_for_nms = boxes3d_to_bev_torch_lidar(bboxes)
             mlvl_bboxes_for_nms.append(bboxes_for_nms)
             mlvl_bboxes.append(bboxes)
@@ -383,6 +366,7 @@ class SECONDHead(nn.Module, AnchorTrainMixin):
             selected_scores = mlvl_scores[selected]
             selected_label_preds = mlvl_label_preds[selected]
             selected_dir_scores = mlvl_dir_scores[selected]
+            # TODO: move dir_offset to box coder
             dir_rot = box_torch_ops.limit_period(
                 selected_bboxes[..., -1] - self.dir_offset,
                 self.dir_limit_offset, np.pi)

mmdet3d/models/anchor_heads/train_mixins.py

@@ -11,8 +11,6 @@ class AnchorTrainMixin(object):
                          anchor_list,
                          gt_bboxes_list,
                          input_metas,
-                         target_means,
-                         target_stds,
                          gt_bboxes_ignore_list=None,
                          gt_labels_list=None,
                          label_channels=1,
@@ -24,8 +22,6 @@ class AnchorTrainMixin(object):
             anchor_list (list[list]): Multi level anchors of each image.
             gt_bboxes_list (list[Tensor]): Ground truth bboxes of each image.
             img_metas (list[dict]): Meta info of each image.
-            target_means (Iterable): Mean value of regression targets.
-            target_stds (Iterable): Std value of regression targets.
 
         Returns:
             tuple
@@ -57,8 +53,6 @@ class AnchorTrainMixin(object):
              gt_bboxes_ignore_list,
              gt_labels_list,
              input_metas,
-             target_means=target_means,
-             target_stds=target_stds,
              label_channels=label_channels,
              num_classes=num_classes,
              sampling=sampling)
@@ -89,8 +83,6 @@ class AnchorTrainMixin(object):
                                 gt_bboxes_ignore,
                                 gt_labels,
                                 input_meta,
-                                target_means,
-                                target_stds,
                                 label_channels=1,
                                 num_classes=1,
                                 sampling=True):
@@ -111,13 +103,12 @@ class AnchorTrainMixin(object):
                     anchor_targets = self.anchor_target_single_assigner(
                         assigner, current_anchors, gt_bboxes[gt_per_cls, :],
                         gt_bboxes_ignore, gt_labels[gt_per_cls], input_meta,
-                        target_means, target_stds, label_channels, num_classes,
-                        sampling)
+                        label_channels, num_classes, sampling)
                 else:
                     anchor_targets = self.anchor_target_single_assigner(
                         assigner, current_anchors, gt_bboxes, gt_bboxes_ignore,
-                        gt_labels, input_meta, target_means, target_stds,
-                        label_channels, num_classes, sampling)
+                        gt_labels, input_meta, label_channels, num_classes,
+                        sampling)
 
                 (labels, label_weights, bbox_targets, bbox_weights,
                  dir_targets, dir_weights, pos_inds, neg_inds) = anchor_targets
@@ -156,8 +147,7 @@ class AnchorTrainMixin(object):
         else:
             return self.anchor_target_single_assigner(
                 self.bbox_assigner, anchors, gt_bboxes, gt_bboxes_ignore,
-                gt_labels, input_meta, target_means, target_stds,
-                label_channels, num_classes, sampling)
+                gt_labels, input_meta, label_channels, num_classes, sampling)
 
     def anchor_target_single_assigner(self,
                                       bbox_assigner,
@@ -166,8 +156,6 @@ class AnchorTrainMixin(object):
                                       gt_bboxes_ignore,
                                       gt_labels,
                                       input_meta,
-                                      target_means,
-                                      target_stds,
                                       label_channels=1,
                                       num_classes=1,
                                       sampling=True):
@@ -188,18 +176,17 @@ class AnchorTrainMixin(object):
             neg_inds = sampling_result.neg_inds
         else:
             pos_inds = torch.nonzero(
-                anchors.new_zeros((anchors.shape[0], ), dtype=torch.long) > 0
+                anchors.new_zeros((anchors.shape[0], ), dtype=torch.bool) > 0
             ).squeeze(-1).unique()
             neg_inds = torch.nonzero(
-                anchors.new_zeros((anchors.shape[0], ), dtype=torch.long) ==
+                anchors.new_zeros((anchors.shape[0], ), dtype=torch.bool) ==
                 0).squeeze(-1).unique()
 
         if gt_labels is not None:
             labels += num_classes
         if len(pos_inds) > 0:
-            pos_bbox_targets = self.bbox_coder.encode_torch(
-                sampling_result.pos_bboxes, sampling_result.pos_gt_bboxes,
-                target_means, target_stds)
+            pos_bbox_targets = self.bbox_coder.encode(
+                sampling_result.pos_bboxes, sampling_result.pos_gt_bboxes)
             pos_dir_targets = get_direction_target(
                 sampling_result.pos_bboxes,
                 pos_bbox_targets,

mmdet3d/models/backbones/second.py

 from functools import partial
 
 import torch.nn as nn
+from mmcv.cnn import build_norm_layer
 from mmcv.runner import load_checkpoint
 
 from mmdet.models import BACKBONES
-from mmdet.ops import build_norm_layer
 
 
 class Empty(nn.Module):
@@ -20,7 +20,7 @@ class Empty(nn.Module):
         return args
 
 
-@BACKBONES.register_module
+@BACKBONES.register_module()
 class SECOND(nn.Module):
     """Compare with RPN, RPNV2 support arbitrary number of stage.
     """
@@ -72,7 +72,7 @@ class SECOND(nn.Module):
 
     def init_weights(self, pretrained=None):
         if isinstance(pretrained, str):
-            from mmdet3d.apis import get_root_logger
+            from mmdet3d.utils import get_root_logger
             logger = get_root_logger()
             load_checkpoint(self, pretrained, strict=False, logger=logger)
 

mmdet3d/models/builder.py

-from mmdet.models.builder import build
-from mmdet.models.registry import (BACKBONES, DETECTORS, HEADS, LOSSES, NECKS,
-                                   ROI_EXTRACTORS, SHARED_HEADS)
+from mmdet.models.builder import (BACKBONES, DETECTORS, HEADS, LOSSES, NECKS,
+                                  ROI_EXTRACTORS, SHARED_HEADS, build)
 from .registry import FUSION_LAYERS, MIDDLE_ENCODERS, VOXEL_ENCODERS
 
 

mmdet3d/models/detectors/__init__.py

@@ -3,10 +3,11 @@ from .mvx_faster_rcnn import (DynamicMVXFasterRCNN, DynamicMVXFasterRCNNV2,
                               DynamicMVXFasterRCNNV3)
 from .mvx_single_stage import MVXSingleStageDetector
 from .mvx_two_stage import MVXTwoStageDetector
+from .parta2 import PartA2
 from .voxelnet import DynamicVoxelNet, VoxelNet
 
 __all__ = [
     'BaseDetector', 'VoxelNet', 'DynamicVoxelNet', 'MVXSingleStageDetector',
     'MVXTwoStageDetector', 'DynamicMVXFasterRCNN', 'DynamicMVXFasterRCNNV2',
-    'DynamicMVXFasterRCNNV3'
+    'DynamicMVXFasterRCNNV3', 'PartA2'
 ]

mmdet3d/models/detectors/base.py

@@ -59,7 +59,7 @@ class BaseDetector(nn.Module, metaclass=ABCMeta):
 
     def init_weights(self, pretrained=None):
         if pretrained is not None:
-            from mmdet3d.apis import get_root_logger
+            from mmdet3d.utils import get_root_logger
             logger = get_root_logger()
             logger.info('load model from: {}'.format(pretrained))
 

mmdet3d/models/detectors/mvx_faster_rcnn.py

@@ -5,7 +5,7 @@ from mmdet.models import DETECTORS
 from .mvx_two_stage import MVXTwoStageDetector
 
 
-@DETECTORS.register_module
+@DETECTORS.register_module()
 class DynamicMVXFasterRCNN(MVXTwoStageDetector):
 
     def __init__(self, **kwargs):
@@ -42,7 +42,7 @@ class DynamicMVXFasterRCNN(MVXTwoStageDetector):
         return points, coors_batch
 
 
-@DETECTORS.register_module
+@DETECTORS.register_module()
 class DynamicMVXFasterRCNNV2(DynamicMVXFasterRCNN):
 
     def __init__(self, **kwargs):
@@ -62,7 +62,7 @@ class DynamicMVXFasterRCNNV2(DynamicMVXFasterRCNN):
         return x
 
 
-@DETECTORS.register_module
+@DETECTORS.register_module()
 class MVXFasterRCNNV2(MVXTwoStageDetector):
 
     def __init__(self, **kwargs):
@@ -84,7 +84,7 @@ class MVXFasterRCNNV2(MVXTwoStageDetector):
         return x
 
 
-@DETECTORS.register_module
+@DETECTORS.register_module()
 class DynamicMVXFasterRCNNV3(DynamicMVXFasterRCNN):
 
     def __init__(self, **kwargs):

mmdet3d/models/detectors/mvx_single_stage.py

@@ -8,7 +8,7 @@ from .. import builder
 from .base import BaseDetector
 
 
-@DETECTORS.register_module
+@DETECTORS.register_module()
 class MVXSingleStageDetector(BaseDetector):
 
     def __init__(self,
@@ -162,7 +162,7 @@ class MVXSingleStageDetector(BaseDetector):
         raise NotImplementedError
 
 
-@DETECTORS.register_module
+@DETECTORS.register_module()
 class DynamicMVXNet(MVXSingleStageDetector):
 
     def __init__(self,
@@ -230,7 +230,7 @@ class DynamicMVXNet(MVXSingleStageDetector):
         return points, coors_batch
 
 
-@DETECTORS.register_module
+@DETECTORS.register_module()
 class DynamicMVXNetV2(DynamicMVXNet):
 
     def __init__(self,
@@ -281,7 +281,7 @@ class DynamicMVXNetV2(DynamicMVXNet):
         return x
 
 
-@DETECTORS.register_module
+@DETECTORS.register_module()
 class DynamicMVXNetV3(DynamicMVXNet):
 
     def __init__(self,

mmdet3d/models/detectors/mvx_two_stage.py

@@ -8,7 +8,7 @@ from .. import builder
 from .base import BaseDetector
 
 
-@DETECTORS.register_module
+@DETECTORS.register_module()
 class MVXTwoStageDetector(BaseDetector):
 
     def __init__(self,

mmdet3d/models/detectors/parta2.py

+import torch
+import torch.nn.functional as F
+
+from mmdet3d.ops import Voxelization
+from mmdet.models import DETECTORS, TwoStageDetector
+from .. import builder
+
+
+@DETECTORS.register_module()
+class PartA2(TwoStageDetector):
+
+    def __init__(self,
+                 voxel_layer,
+                 voxel_encoder,
+                 middle_encoder,
+                 backbone,
+                 neck=None,
+                 rpn_head=None,
+                 roi_head=None,
+                 train_cfg=None,
+                 test_cfg=None,
+                 pretrained=None):
+        super(PartA2, self).__init__(
+            backbone=backbone,
+            neck=neck,
+            rpn_head=rpn_head,
+            roi_head=roi_head,
+            train_cfg=train_cfg,
+            test_cfg=test_cfg,
+            pretrained=pretrained,
+        )
+        self.voxel_layer = Voxelization(**voxel_layer)
+        self.voxel_encoder = builder.build_voxel_encoder(voxel_encoder)
+        self.middle_encoder = builder.build_middle_encoder(middle_encoder)
+
+    def extract_feat(self, points, img_meta):
+        voxels, num_points, coors = self.voxelize(points)
+        voxel_dict = dict(voxels=voxels, num_points=num_points, coors=coors)
+        voxel_features = self.voxel_encoder(voxels, num_points, coors)
+        batch_size = coors[-1, 0].item() + 1
+        feats_dict = self.middle_encoder(voxel_features, coors, batch_size)
+        x = self.backbone(feats_dict['spatial_features'])
+        if self.with_neck:
+            neck_feats = self.neck(x)
+            feats_dict.update({'neck_feats': neck_feats})
+        return feats_dict, voxel_dict
+
+    @torch.no_grad()
+    def voxelize(self, points):
+        voxels, coors, num_points = [], [], []
+        for res in points:
+            res_voxels, res_coors, res_num_points = self.voxel_layer(res)
+            voxels.append(res_voxels)
+            coors.append(res_coors)
+            num_points.append(res_num_points)
+        voxels = torch.cat(voxels, dim=0)
+        num_points = torch.cat(num_points, dim=0)
+        coors_batch = []
+        for i, coor in enumerate(coors):
+            coor_pad = F.pad(coor, (1, 0), mode='constant', value=i)
+            coors_batch.append(coor_pad)
+        coors_batch = torch.cat(coors_batch, dim=0)
+        return voxels, num_points, coors_batch
+
+    def forward_train(self,
+                      points,
+                      img_meta,
+                      gt_bboxes_3d,
+                      gt_labels_3d,
+                      gt_bboxes_ignore=None,
+                      proposals=None):
+        # TODO: complete it
+        feats_dict, voxels_dict = self.extract_feat(points, img_meta)
+
+        losses = dict()
+
+        if self.with_rpn:
+            rpn_outs = self.rpn_head(feats_dict['neck_feats'])
+            rpn_loss_inputs = rpn_outs + (gt_bboxes_3d, gt_labels_3d, img_meta)
+            rpn_losses = self.rpn_head.loss(
+                *rpn_loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore)
+            losses.update(rpn_losses)
+
+            proposal_cfg = self.train_cfg.get('rpn_proposal',
+                                              self.test_cfg.rpn)
+            proposal_inputs = rpn_outs + (img_meta, proposal_cfg)
+            proposal_list = self.rpn_head.get_bboxes(*proposal_inputs)
+        else:
+            proposal_list = proposals  # noqa: F841
+
+        return losses
+
+    def forward_test(self, **kwargs):
+        return self.simple_test(**kwargs)
+
+    def forward(self, return_loss=True, **kwargs):
+        if return_loss:
+            return self.forward_train(**kwargs)
+        else:
+            return self.forward_test(**kwargs)
+
+    def simple_test(self,
+                    points,
+                    img_meta,
+                    gt_bboxes_3d=None,
+                    proposals=None,
+                    rescale=False):
+        feats_dict, voxels_dict = self.extract_feat(points, img_meta)
+        # TODO: complete it
+        if proposals is None:
+            proposal_list = self.simple_test_rpn(feats_dict['neck_feats'],
+                                                 img_meta, self.test_cfg.rpn)
+        else:
+            proposal_list = proposals
+
+        return self.roi_head.simple_test(
+            feats_dict, proposal_list, img_meta, rescale=rescale)

mmdet3d/models/detectors/voxelnet.py

@@ -6,7 +6,7 @@ from mmdet.models import DETECTORS, SingleStageDetector
 from .. import builder
 
 
-@DETECTORS.register_module
+@DETECTORS.register_module()
 class VoxelNet(SingleStageDetector):
 
     def __init__(self,
@@ -88,7 +88,7 @@ class VoxelNet(SingleStageDetector):
         return bbox_list
 
 
-@DETECTORS.register_module
+@DETECTORS.register_module()
 class DynamicVoxelNet(VoxelNet):
 
     def __init__(self,

mmdet3d/models/fusion_layers/point_fusion.py

 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from mmcv.cnn import xavier_init
+from mmcv.cnn import ConvModule, xavier_init
 
-from mmdet.ops import ConvModule
 from ..registry import FUSION_LAYERS
 
 
@@ -96,7 +95,7 @@ def point_sample(
     return point_features.squeeze().t()
 
 
-@FUSION_LAYERS.register_module
+@FUSION_LAYERS.register_module()
 class PointFusion(nn.Module):
     """Fuse image features from fused single scale features
     """
@@ -235,9 +234,8 @@ class PointFusion(nn.Module):
             pts.new_tensor(img_meta['pcd_trans'])
             if 'pcd_trans' in img_meta.keys() else 0)
         pcd_rotate_mat = (
-            pts.new_tensor(img_meta['pcd_rotation'])
-            if 'pcd_rotation' in img_meta.keys() else
-            torch.eye(3).type_as(pts).to(pts.device))
+            pts.new_tensor(img_meta['pcd_rotation']) if 'pcd_rotation'
+            in img_meta.keys() else torch.eye(3).type_as(pts).to(pts.device))
         img_scale_factor = (
             img_meta['scale_factor']
             if 'scale_factor' in img_meta.keys() else 1)

mmdet3d/models/losses/__init__.py

-from mmdet.models.losses import FocalLoss, SmoothL1Loss
+from mmdet.models.losses import FocalLoss, SmoothL1Loss, binary_cross_entropy
 
-__all__ = ['FocalLoss', 'SmoothL1Loss']
+__all__ = ['FocalLoss', 'SmoothL1Loss', 'binary_cross_entropy']

mmdet3d/models/middle_encoders/__init__.py

 from .pillar_scatter import PointPillarsScatter
 from .sparse_encoder import SparseEncoder
+from .sparse_unet import SparseUNet
 
-__all__ = ['PointPillarsScatter', 'SparseEncoder']
+__all__ = ['PointPillarsScatter', 'SparseEncoder', 'SparseUNet']

mmdet3d/models/middle_encoders/pillar_scatter.py

@@ -4,7 +4,7 @@ from torch import nn
 from ..registry import MIDDLE_ENCODERS
 
 
-@MIDDLE_ENCODERS.register_module
+@MIDDLE_ENCODERS.register_module()
 class PointPillarsScatter(nn.Module):
 
     def __init__(self, in_channels, output_shape):