Merge branch 'refactor-anchor_coder' into 'master'

Refactor anchor generator and box coder

See merge request open-mmlab/mmdet.3d!3

mmdet3d/datasets/builder.py

 import copy
 
+from mmcv.utils import build_from_cfg
+
 from mmdet.datasets import DATASETS, ConcatDataset, RepeatDataset
-from mmdet.utils import build_from_cfg
 from .dataset_wrappers import RepeatFactorDataset
 
 

mmdet3d/datasets/kitti_dataset.py

@@ -8,8 +8,8 @@ import torch
 import torch.utils.data as torch_data
 
 from mmdet.datasets import DATASETS
-from mmdet.datasets.pipelines import Compose
 from ..core.bbox import box_np_ops
+from .pipelines import Compose
 from .utils import remove_dontcare
 
 
@@ -274,7 +274,7 @@ class KittiDataset(torch_data.Dataset):
                                                   out)
         return result_files
 
-    def evaluate(self, result_files, eval_types=None):
+    def evaluate(self, result_files, logger=None, eval_types=None):
         from mmdet3d.core.evaluation import kitti_eval
         gt_annos = [info['annos'] for info in self.kitti_infos]
         if eval_types == 'img_bbox':
@@ -283,7 +283,7 @@ class KittiDataset(torch_data.Dataset):
         else:
             ap_result_str, ap_dict = kitti_eval(gt_annos, result_files,
                                                 self.class_names)
-        return ap_result_str, ap_dict
+        return ap_dict
 
     def bbox2result_kitti(self, net_outputs, class_names, out=None):
         if out:

mmdet3d/datasets/nuscenes2d_dataset.py

-from pycocotools.coco import COCO
-
-from mmdet3d.core.evaluation.coco_utils import getImgIds
-from mmdet.datasets import DATASETS, CocoDataset
-
-
-@DATASETS.register_module
-class NuScenes2DDataset(CocoDataset):
-
-    CLASSES = ('car', 'truck', 'trailer', 'bus', 'construction_vehicle',
-               'bicycle', 'motorcycle', 'pedestrian', 'traffic_cone',
-               'barrier')
-
-    def load_annotations(self, ann_file):
-        if not self.class_names:
-            self.class_names = self.CLASSES
-        self.coco = COCO(ann_file)
-        # send class_names into the get id
-        # in case we only need to train on several classes
-        # by default self.class_names = CLASSES
-        self.cat_ids = self.coco.getCatIds(catNms=self.class_names)
-
-        self.cat2label = {
-            cat_id: i  # + 1 rm +1 here thus the 0-79 are fg, 80 is bg
-            for i, cat_id in enumerate(self.cat_ids)
-        }
-        # send cat ids to the get img id
-        # in case we only need to train on several classes
-        if len(self.cat_ids) < len(self.CLASSES):
-            self.img_ids = getImgIds(self.coco, catIds=self.cat_ids)
-        else:
-            self.img_ids = self.coco.getImgIds()
-        img_infos = []
-        for i in self.img_ids:
-            info = self.coco.loadImgs([i])[0]
-            info['filename'] = info['file_name']
-            img_infos.append(info)
-        return img_infos

mmdet3d/datasets/nuscenes_dataset.py

@@ -9,8 +9,8 @@ import torch.utils.data as torch_data
 from nuscenes.utils.data_classes import Box as NuScenesBox
 
 from mmdet.datasets import DATASETS
-from mmdet.datasets.pipelines import Compose
 from ..core.bbox import box_np_ops
+from .pipelines import Compose
 
 
 @DATASETS.register_module

mmdet3d/datasets/pipelines/__init__.py

+from mmdet.datasets.pipelines import Compose
+from .dbsampler import DataBaseSampler, MMDataBaseSampler
+from .formating import DefaultFormatBundle, DefaultFormatBundle3D
+from .loading import LoadMultiViewImageFromFiles, LoadPointsFromFile
 from .train_aug import (GlobalRotScale, ObjectNoise, ObjectRangeFilter,
                         ObjectSample, PointShuffle, PointsRangeFilter,
                         RandomFlip3D)
 
 __all__ = [
     'ObjectSample', 'RandomFlip3D', 'ObjectNoise', 'GlobalRotScale',
-    'PointShuffle', 'ObjectRangeFilter', 'PointsRangeFilter', 'Collect3D'
+    'PointShuffle', 'ObjectRangeFilter', 'PointsRangeFilter', 'Collect3D',
+    'Compose', 'LoadMultiViewImageFromFiles', 'LoadPointsFromFile',
+    'DefaultFormatBundle', 'DefaultFormatBundle3D', 'DataBaseSampler',
+    'MMDataBaseSampler'
 ]

mmdet3d/datasets/pipelines/dbsampler.py

@@ -68,7 +68,7 @@ class DataBaseSampler(object):
             db_infos = pickle.load(f)
 
         # filter database infos
-        from mmdet3d.apis import get_root_logger
+        from mmdet.apis import get_root_logger
         logger = get_root_logger()
         for k, v in db_infos.items():
             logger.info(f'load {len(v)} {k} database infos')

mmdet3d/datasets/pipelines/formating.py

 import numpy as np
 from mmcv.parallel import DataContainer as DC
 
+from mmdet.datasets.builder import PIPELINES
 from mmdet.datasets.pipelines import to_tensor
-from mmdet.datasets.registry import PIPELINES
 
 PIPELINES._module_dict.pop('DefaultFormatBundle')
 

mmdet3d/datasets/pipelines/loading.py

@@ -3,7 +3,7 @@ import os.path as osp
 import mmcv
 import numpy as np
 
-from mmdet.datasets.registry import PIPELINES
+from mmdet.datasets.builder import PIPELINES
 
 
 @PIPELINES.register_module

mmdet3d/datasets/pipelines/train_aug.py

 import numpy as np
+from mmcv.utils import build_from_cfg
 
 from mmdet3d.core.bbox import box_np_ops
-from mmdet3d.utils import build_from_cfg
+from mmdet.datasets.builder import PIPELINES
 from mmdet.datasets.pipelines import RandomFlip
-from mmdet.datasets.registry import PIPELINES
 from ..registry import OBJECTSAMPLERS
 from .data_augment_utils import noise_per_object_v3_
 

mmdet3d/datasets/registry.py

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

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
 
 
@@ -15,12 +14,6 @@ class Anchor3DVeloHead(SECONDHead):
     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/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
 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/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/fusion_layers/point_fusion.py

@@ -235,9 +235,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/registry.py

-from mmdet.utils import Registry
+from mmcv.utils import Registry
 
 VOXEL_ENCODERS = Registry('voxel_encoder')
 MIDDLE_ENCODERS = Registry('middle_encoder')

mmdet3d/models/utils/__init__.py

-from mmdet.models.utils import ResLayer, bias_init_with_prob
-
-__all__ = ['bias_init_with_prob', 'ResLayer']

mmdet3d/models/utils/weight_init.py

-import numpy as np
-import torch.nn as nn
-
-
-def xavier_init(module, gain=1, bias=0, distribution='normal'):
-    assert distribution in ['uniform', 'normal']
-    if distribution == 'uniform':
-        nn.init.xavier_uniform_(module.weight, gain=gain)
-    else:
-        nn.init.xavier_normal_(module.weight, gain=gain)
-    if hasattr(module, 'bias'):
-        nn.init.constant_(module.bias, bias)
-
-
-def normal_init(module, mean=0, std=1, bias=0):
-    nn.init.normal_(module.weight, mean, std)
-    if hasattr(module, 'bias'):
-        nn.init.constant_(module.bias, bias)
-
-
-def uniform_init(module, a=0, b=1, bias=0):
-    nn.init.uniform_(module.weight, a, b)
-    if hasattr(module, 'bias'):
-        nn.init.constant_(module.bias, bias)
-
-
-def kaiming_init(module,
-                 mode='fan_out',
-                 nonlinearity='relu',
-                 bias=0,
-                 distribution='normal'):
-    assert distribution in ['uniform', 'normal']
-    if distribution == 'uniform':
-        nn.init.kaiming_uniform_(
-            module.weight, mode=mode, nonlinearity=nonlinearity)
-    else:
-        nn.init.kaiming_normal_(
-            module.weight, mode=mode, nonlinearity=nonlinearity)
-    if hasattr(module, 'bias'):
-        nn.init.constant_(module.bias, bias)
-
-
-def bias_init_with_prob(prior_prob):
-    """ initialize conv/fc bias value according to giving probablity"""
-    bias_init = float(-np.log((1 - prior_prob) / prior_prob))
-    return bias_init

mmdet3d/ops/iou3d/src/iou3d.cpp

-#include <torch/serialize/tensor.h>
-#include <torch/extension.h>
-#include <vector>
 #include <cuda.h>
 #include <cuda_runtime_api.h>
+#include <torch/extension.h>
+#include <torch/serialize/tensor.h>
 
-#define CHECK_CUDA(x) AT_CHECK(x.type().is_cuda(), #x, " must be a CUDAtensor ")
-#define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x, " must be contiguous ")
-#define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x)
-
-#define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0))
+#include <vector>
 
-#define CHECK_ERROR(ans) { gpuAssert((ans), __FILE__, __LINE__); }
-inline void gpuAssert(cudaError_t code, const char *file, int line, bool abort=true)
-{
-   if (code != cudaSuccess)
-   {
-      fprintf(stderr,"GPUassert: %s %s %d\n", cudaGetErrorString(code), file, line);
+#define CHECK_CUDA(x) \
+  TORCH_CHECK(x.device().is_cuda(), #x, " must be a CUDAtensor ")
+#define CHECK_CONTIGUOUS(x) \
+  TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous ")
+#define CHECK_INPUT(x) \
+  CHECK_CUDA(x);       \
+  CHECK_CONTIGUOUS(x)
+
+#define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0))
+
+#define CHECK_ERROR(ans) \
+  { gpuAssert((ans), __FILE__, __LINE__); }
+inline void gpuAssert(cudaError_t code, const char *file, int line,
+                      bool abort = true) {
+  if (code != cudaSuccess) {
+    fprintf(stderr, "GPUassert: %s %s %d\n", cudaGetErrorString(code), file,
+            line);
     if (abort) exit(code);
   }
 }
 
 const int THREADS_PER_BLOCK_NMS = sizeof(unsigned long long) * 8;
 
-
-void boxesoverlapLauncher(const int num_a, const float *boxes_a, const int num_b, const float *boxes_b, float *ans_overlap);
-void boxesioubevLauncher(const int num_a, const float *boxes_a, const int num_b, const float *boxes_b, float *ans_iou);
-void nmsLauncher(const float *boxes, unsigned long long * mask, int boxes_num, float nms_overlap_thresh);
-void nmsNormalLauncher(const float *boxes, unsigned long long * mask, int boxes_num, float nms_overlap_thresh);
-
-int boxes_overlap_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans_overlap){
+void boxesoverlapLauncher(const int num_a, const float *boxes_a,
+                          const int num_b, const float *boxes_b,
+                          float *ans_overlap);
+void boxesioubevLauncher(const int num_a, const float *boxes_a, const int num_b,
+                         const float *boxes_b, float *ans_iou);
+void nmsLauncher(const float *boxes, unsigned long long *mask, int boxes_num,
+                 float nms_overlap_thresh);
+void nmsNormalLauncher(const float *boxes, unsigned long long *mask,
+                       int boxes_num, float nms_overlap_thresh);
+
+int boxes_overlap_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b,
+                          at::Tensor ans_overlap) {
   // params boxes_a: (N, 5) [x1, y1, x2, y2, ry]
   // params boxes_b: (M, 5)
   // params ans_overlap: (N, M)
@@ -40,16 +51,18 @@ int boxes_overlap_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans
   int num_a = boxes_a.size(0);
   int num_b = boxes_b.size(0);
 
-    const float * boxes_a_data = boxes_a.data<float>();
-    const float * boxes_b_data = boxes_b.data<float>();
-    float * ans_overlap_data = ans_overlap.data<float>();
+  const float *boxes_a_data = boxes_a.data_ptr<float>();
+  const float *boxes_b_data = boxes_b.data_ptr<float>();
+  float *ans_overlap_data = ans_overlap.data_ptr<float>();
 
-    boxesoverlapLauncher(num_a, boxes_a_data, num_b, boxes_b_data, ans_overlap_data);
+  boxesoverlapLauncher(num_a, boxes_a_data, num_b, boxes_b_data,
+                       ans_overlap_data);
 
   return 1;
 }
 
-int boxes_iou_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans_iou){
+int boxes_iou_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b,
+                      at::Tensor ans_iou) {
   // params boxes_a: (N, 5) [x1, y1, x2, y2, ry]
   // params boxes_b: (M, 5)
   // params ans_overlap: (N, M)
@@ -61,16 +74,16 @@ int boxes_iou_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans_iou
   int num_a = boxes_a.size(0);
   int num_b = boxes_b.size(0);
 
-    const float * boxes_a_data = boxes_a.data<float>();
-    const float * boxes_b_data = boxes_b.data<float>();
-    float * ans_iou_data = ans_iou.data<float>();
+  const float *boxes_a_data = boxes_a.data_ptr<float>();
+  const float *boxes_b_data = boxes_b.data_ptr<float>();
+  float *ans_iou_data = ans_iou.data_ptr<float>();
 
   boxesioubevLauncher(num_a, boxes_a_data, num_b, boxes_b_data, ans_iou_data);
 
   return 1;
 }
 
-int nms_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh){
+int nms_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh) {
   // params boxes: (N, 5) [x1, y1, x2, y2, ry]
   // params keep: (N)
 
@@ -78,21 +91,24 @@ int nms_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh){
   CHECK_CONTIGUOUS(keep);
 
   int boxes_num = boxes.size(0);
-    const float * boxes_data = boxes.data<float>();
-    long * keep_data = keep.data<long>();
+  const float *boxes_data = boxes.data_ptr<float>();
+  long *keep_data = keep.data_ptr<long>();
 
   const int col_blocks = DIVUP(boxes_num, THREADS_PER_BLOCK_NMS);
 
   unsigned long long *mask_data = NULL;
-    CHECK_ERROR(cudaMalloc((void**)&mask_data, boxes_num * col_blocks * sizeof(unsigned long long)));
+  CHECK_ERROR(cudaMalloc((void **)&mask_data,
+                         boxes_num * col_blocks * sizeof(unsigned long long)));
   nmsLauncher(boxes_data, mask_data, boxes_num, nms_overlap_thresh);
 
   // unsigned long long mask_cpu[boxes_num * col_blocks];
-    // unsigned long long *mask_cpu = new unsigned long long [boxes_num * col_blocks];
+  // unsigned long long *mask_cpu = new unsigned long long [boxes_num *
+  // col_blocks];
   std::vector<unsigned long long> mask_cpu(boxes_num * col_blocks);
 
-//    printf("boxes_num=%d, col_blocks=%d\n", boxes_num, col_blocks);
-    CHECK_ERROR(cudaMemcpy(&mask_cpu[0], mask_data, boxes_num * col_blocks * sizeof(unsigned long long),
+  //    printf("boxes_num=%d, col_blocks=%d\n", boxes_num, col_blocks);
+  CHECK_ERROR(cudaMemcpy(&mask_cpu[0], mask_data,
+                         boxes_num * col_blocks * sizeof(unsigned long long),
                          cudaMemcpyDeviceToHost));
 
   cudaFree(mask_data);
@@ -102,25 +118,25 @@ int nms_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh){
 
   int num_to_keep = 0;
 
-    for (int i = 0; i < boxes_num; i++){
+  for (int i = 0; i < boxes_num; i++) {
     int nblock = i / THREADS_PER_BLOCK_NMS;
     int inblock = i % THREADS_PER_BLOCK_NMS;
 
-        if (!(remv_cpu[nblock] & (1ULL << inblock))){
+    if (!(remv_cpu[nblock] & (1ULL << inblock))) {
       keep_data[num_to_keep++] = i;
       unsigned long long *p = &mask_cpu[0] + i * col_blocks;
-            for (int j = nblock; j < col_blocks; j++){
+      for (int j = nblock; j < col_blocks; j++) {
         remv_cpu[j] |= p[j];
       }
     }
   }
-    if ( cudaSuccess != cudaGetLastError() ) printf( "Error!\n" );
+  if (cudaSuccess != cudaGetLastError()) printf("Error!\n");
 
   return num_to_keep;
 }
 
-
-int nms_normal_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh){
+int nms_normal_gpu(at::Tensor boxes, at::Tensor keep,
+                   float nms_overlap_thresh) {
   // params boxes: (N, 5) [x1, y1, x2, y2, ry]
   // params keep: (N)
 
@@ -128,21 +144,24 @@ int nms_normal_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh){
   CHECK_CONTIGUOUS(keep);
 
   int boxes_num = boxes.size(0);
-    const float * boxes_data = boxes.data<float>();
-    long * keep_data = keep.data<long>();
+  const float *boxes_data = boxes.data_ptr<float>();
+  long *keep_data = keep.data_ptr<long>();
 
   const int col_blocks = DIVUP(boxes_num, THREADS_PER_BLOCK_NMS);
 
   unsigned long long *mask_data = NULL;
-    CHECK_ERROR(cudaMalloc((void**)&mask_data, boxes_num * col_blocks * sizeof(unsigned long long)));
+  CHECK_ERROR(cudaMalloc((void **)&mask_data,
+                         boxes_num * col_blocks * sizeof(unsigned long long)));
   nmsNormalLauncher(boxes_data, mask_data, boxes_num, nms_overlap_thresh);
 
   // unsigned long long mask_cpu[boxes_num * col_blocks];
-    // unsigned long long *mask_cpu = new unsigned long long [boxes_num * col_blocks];
+  // unsigned long long *mask_cpu = new unsigned long long [boxes_num *
+  // col_blocks];
   std::vector<unsigned long long> mask_cpu(boxes_num * col_blocks);
 
-//    printf("boxes_num=%d, col_blocks=%d\n", boxes_num, col_blocks);
-    CHECK_ERROR(cudaMemcpy(&mask_cpu[0], mask_data, boxes_num * col_blocks * sizeof(unsigned long long),
+  //    printf("boxes_num=%d, col_blocks=%d\n", boxes_num, col_blocks);
+  CHECK_ERROR(cudaMemcpy(&mask_cpu[0], mask_data,
+                         boxes_num * col_blocks * sizeof(unsigned long long),
                          cudaMemcpyDeviceToHost));
 
   cudaFree(mask_data);
@@ -152,27 +171,26 @@ int nms_normal_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh){
 
   int num_to_keep = 0;
 
-    for (int i = 0; i < boxes_num; i++){
+  for (int i = 0; i < boxes_num; i++) {
     int nblock = i / THREADS_PER_BLOCK_NMS;
     int inblock = i % THREADS_PER_BLOCK_NMS;
 
-        if (!(remv_cpu[nblock] & (1ULL << inblock))){
+    if (!(remv_cpu[nblock] & (1ULL << inblock))) {
       keep_data[num_to_keep++] = i;
       unsigned long long *p = &mask_cpu[0] + i * col_blocks;
-            for (int j = nblock; j < col_blocks; j++){
+      for (int j = nblock; j < col_blocks; j++) {
         remv_cpu[j] |= p[j];
       }
     }
   }
-    if ( cudaSuccess != cudaGetLastError() ) printf( "Error!\n" );
+  if (cudaSuccess != cudaGetLastError()) printf("Error!\n");
 
   return num_to_keep;
 }
 
-
-
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("boxes_overlap_bev_gpu", &boxes_overlap_bev_gpu, "oriented boxes overlap");
+  m.def("boxes_overlap_bev_gpu", &boxes_overlap_bev_gpu,
+        "oriented boxes overlap");
   m.def("boxes_iou_bev_gpu", &boxes_iou_bev_gpu, "oriented boxes iou");
   m.def("nms_gpu", &nms_gpu, "oriented nms gpu");
   m.def("nms_normal_gpu", &nms_normal_gpu, "nms gpu");

mmdet3d/ops/roiaware_pool3d/src/points_in_boxes_cpu.cpp

-//Modified from
-//https://github.com/sshaoshuai/PCDet/blob/master/pcdet/ops/roiaware_pool3d/src/roiaware_pool3d_kernel.cu
-//Points in boxes cpu
-//Written by Shaoshuai Shi
-//All Rights Reserved 2019.
+// Modified from
+// https://github.com/sshaoshuai/PCDet/blob/master/pcdet/ops/roiaware_pool3d/src/roiaware_pool3d_kernel.cu
+// Points in boxes cpu
+// Written by Shaoshuai Shi
+// All Rights Reserved 2019.
 
-
-#include <torch/serialize/tensor.h>
-#include <torch/extension.h>
 #include <assert.h>
-
 #include <math.h>
 #include <stdio.h>
+#include <torch/extension.h>
+#include <torch/serialize/tensor.h>
 
-#define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x, " must be contiguous ")
+#define CHECK_CONTIGUOUS(x) \
+  TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous ")
 // #define DEBUG
 
-
-inline void lidar_to_local_coords_cpu(float shift_x, float shift_y, float rz, float &local_x, float &local_y){
+inline void lidar_to_local_coords_cpu(float shift_x, float shift_y, float rz,
+                                      float &local_x, float &local_y) {
   // should rotate pi/2 + alpha to translate LiDAR to local
   float rot_angle = rz + M_PI / 2;
   float cosa = cos(rot_angle), sina = sin(rot_angle);
@@ -24,10 +23,11 @@ inline void lidar_to_local_coords_cpu(float shift_x, float shift_y, float rz, fl
   local_y = shift_x * sina + shift_y * cosa;
 }
 
-
-inline int check_pt_in_box3d_cpu(const float *pt, const float *box3d, float &local_x, float &local_y){
+inline int check_pt_in_box3d_cpu(const float *pt, const float *box3d,
+                                 float &local_x, float &local_y) {
   // param pt: (x, y, z)
-    // param box3d: (cx, cy, cz, w, l, h, rz) in LiDAR coordinate, cz in the bottom center
+  // param box3d: (cx, cy, cz, w, l, h, rz) in LiDAR coordinate, cz in the
+  // bottom center
   float x = pt[0], y = pt[1], z = pt[2];
   float cx = box3d[0], cy = box3d[1], cz = box3d[2];
   float w = box3d[3], l = box3d[4], h = box3d[5], rz = box3d[6];
@@ -35,15 +35,16 @@ inline int check_pt_in_box3d_cpu(const float *pt, const float *box3d, float &loc
 
   if (fabsf(z - cz) > h / 2.0) return 0;
   lidar_to_local_coords_cpu(x - cx, y - cy, rz, local_x, local_y);
-    float in_flag = (local_x > -l / 2.0) & (local_x < l / 2.0) & (local_y > -w / 2.0) & (local_y < w / 2.0);
+  float in_flag = (local_x > -l / 2.0) & (local_x < l / 2.0) &
+                  (local_y > -w / 2.0) & (local_y < w / 2.0);
   return in_flag;
 }
 
-
-int points_in_boxes_cpu(at::Tensor boxes_tensor, at::Tensor pts_tensor, at::Tensor pts_indices_tensor){
-    // params boxes: (N, 7) [x, y, z, w, l, h, rz] in LiDAR coordinate, z is the bottom center, each box DO NOT overlaps
-    // params pts: (npoints, 3) [x, y, z] in LiDAR coordinate
-    // params pts_indices: (N, npoints)
+int points_in_boxes_cpu(at::Tensor boxes_tensor, at::Tensor pts_tensor,
+                        at::Tensor pts_indices_tensor) {
+  // params boxes: (N, 7) [x, y, z, w, l, h, rz] in LiDAR coordinate, z is the
+  // bottom center, each box DO NOT overlaps params pts: (npoints, 3) [x, y, z]
+  // in LiDAR coordinate params pts_indices: (N, npoints)
 
   CHECK_CONTIGUOUS(boxes_tensor);
   CHECK_CONTIGUOUS(pts_tensor);
@@ -52,14 +53,15 @@ int points_in_boxes_cpu(at::Tensor boxes_tensor, at::Tensor pts_tensor, at::Tens
   int boxes_num = boxes_tensor.size(0);
   int pts_num = pts_tensor.size(0);
 
-    const float *boxes = boxes_tensor.data<float>();
-    const float *pts = pts_tensor.data<float>();
-    int *pts_indices = pts_indices_tensor.data<int>();
+  const float *boxes = boxes_tensor.data_ptr<float>();
+  const float *pts = pts_tensor.data_ptr<float>();
+  int *pts_indices = pts_indices_tensor.data_ptr<int>();
 
   float local_x = 0, local_y = 0;
-    for (int i = 0; i < boxes_num; i++){
-        for (int j = 0; j < pts_num; j++){
-            int cur_in_flag = check_pt_in_box3d_cpu(pts + j * 3, boxes + i * 7, local_x, local_y);
+  for (int i = 0; i < boxes_num; i++) {
+    for (int j = 0; j < pts_num; j++) {
+      int cur_in_flag =
+          check_pt_in_box3d_cpu(pts + j * 3, boxes + i * 7, local_x, local_y);
       pts_indices[i * pts_num + j] = cur_in_flag;
     }
   }