Support PV-RCNN++ frameworks, support VectorPool aggregation

serkankartal · Dec 26, 2021 · 1378979 · 1378979
1 parent 183d353
commit 1378979
Show file tree

Hide file tree

Showing 11 changed files with 1,553 additions and 37 deletions.
diff --git a/pcdet/models/backbones_3d/pfe/voxel_set_abstraction.py b/pcdet/models/backbones_3d/pfe/voxel_set_abstraction.py
@@ -139,38 +139,31 @@ def __init__(self, model_cfg, voxel_size, point_cloud_range, num_bev_features=No
             if src_name in ['bev', 'raw_points']:
                 continue
             self.downsample_times_map[src_name] = SA_cfg[src_name].DOWNSAMPLE_FACTOR
-            mlps = SA_cfg[src_name].MLPS
-            for k in range(len(mlps)):
-                mlps[k] = [mlps[k][0]] + mlps[k]
-            cur_layer = pointnet2_stack_modules.StackSAModuleMSG(
-                radii=SA_cfg[src_name].POOL_RADIUS,
-                nsamples=SA_cfg[src_name].NSAMPLE,
-                mlps=mlps,
-                use_xyz=True,
-                pool_method='max_pool',
+
+            if SA_cfg[src_name].get('INPUT_CHANNELS', None) is None:
+                input_channels = SA_cfg[src_name].MLPS[0][0] \
+                    if isinstance(SA_cfg[src_name].MLPS[0], list) else SA_cfg[src_name].MLPS[0]
+            else:
+                input_channels = SA_cfg[src_name]['INPUT_CHANNELS']
+
+            cur_layer, cur_num_c_out = pointnet2_stack_modules.build_local_aggregation_module(
+                input_channels=input_channels, config=SA_cfg[src_name]
             )
             self.SA_layers.append(cur_layer)
             self.SA_layer_names.append(src_name)
 
-            c_in += sum([x[-1] for x in mlps])
+            c_in += cur_num_c_out
 
         if 'bev' in self.model_cfg.FEATURES_SOURCE:
             c_bev = num_bev_features
             c_in += c_bev
 
         if 'raw_points' in self.model_cfg.FEATURES_SOURCE:
-            mlps = SA_cfg['raw_points'].MLPS
-            for k in range(len(mlps)):
-                mlps[k] = [num_rawpoint_features - 3] + mlps[k]
-
-            self.SA_rawpoints = pointnet2_stack_modules.StackSAModuleMSG(
-                radii=SA_cfg['raw_points'].POOL_RADIUS,
-                nsamples=SA_cfg['raw_points'].NSAMPLE,
-                mlps=mlps,
-                use_xyz=True,
-                pool_method='max_pool'
+            self.SA_rawpoints, cur_num_c_out = pointnet2_stack_modules.build_local_aggregation_module(
+                input_channels=num_rawpoint_features - 3, config=SA_cfg['raw_points']
             )
-            c_in += sum([x[-1] for x in mlps])
+
+            c_in += cur_num_c_out
 
         self.vsa_point_feature_fusion = nn.Sequential(
             nn.Linear(c_in, self.model_cfg.NUM_OUTPUT_FEATURES, bias=False),

diff --git a/pcdet/models/detectors/__init__.py b/pcdet/models/detectors/__init__.py
@@ -8,6 +8,7 @@
 from .caddn import CaDDN
 from .voxel_rcnn import VoxelRCNN
 from .centerpoint import CenterPoint
+from .pv_rcnn_plusplus import PVRCNNPlusPlus
 
 __all__ = {
     'Detector3DTemplate': Detector3DTemplate,
@@ -19,7 +20,8 @@
     'SECONDNetIoU': SECONDNetIoU,
     'CaDDN': CaDDN,
     'VoxelRCNN': VoxelRCNN,
-    'CenterPoint': CenterPoint
+    'CenterPoint': CenterPoint,
+    'PVRCNNPlusPlus': PVRCNNPlusPlus
 }
 
 

diff --git a/pcdet/models/roi_heads/pvrcnn_head.py b/pcdet/models/roi_heads/pvrcnn_head.py
@@ -10,21 +10,12 @@ def __init__(self, input_channels, model_cfg, num_class=1, **kwargs):
         super().__init__(num_class=num_class, model_cfg=model_cfg)
         self.model_cfg = model_cfg
 
-        mlps = self.model_cfg.ROI_GRID_POOL.MLPS
-        for k in range(len(mlps)):
-            mlps[k] = [input_channels] + mlps[k]
-
-        self.roi_grid_pool_layer = pointnet2_stack_modules.StackSAModuleMSG(
-            radii=self.model_cfg.ROI_GRID_POOL.POOL_RADIUS,
-            nsamples=self.model_cfg.ROI_GRID_POOL.NSAMPLE,
-            mlps=mlps,
-            use_xyz=True,
-            pool_method=self.model_cfg.ROI_GRID_POOL.POOL_METHOD,
+        self.roi_grid_pool_layer, num_c_out = pointnet2_stack_modules.build_local_aggregation_module(
+            input_channels=input_channels, config=self.model_cfg.ROI_GRID_POOL
         )
 
         GRID_SIZE = self.model_cfg.ROI_GRID_POOL.GRID_SIZE
-        c_out = sum([x[-1] for x in mlps])
-        pre_channel = GRID_SIZE * GRID_SIZE * GRID_SIZE * c_out
+        pre_channel = GRID_SIZE * GRID_SIZE * GRID_SIZE * num_c_out
 
         shared_fc_list = []
         for k in range(0, self.model_cfg.SHARED_FC.__len__()):
@@ -150,9 +141,11 @@ def forward(self, batch_dict):
             batch_dict, nms_config=self.model_cfg.NMS_CONFIG['TRAIN' if self.training else 'TEST']
         )
         if self.training:
-            targets_dict = self.assign_targets(batch_dict)
-            batch_dict['rois'] = targets_dict['rois']
-            batch_dict['roi_labels'] = targets_dict['roi_labels']
+            targets_dict = batch_dict.get('roi_targets_dict', None)
+            if targets_dict is None:
+                targets_dict = self.assign_targets(batch_dict)
+                batch_dict['rois'] = targets_dict['rois']
+                batch_dict['roi_labels'] = targets_dict['roi_labels']
 
         # RoI aware pooling
         pooled_features = self.roi_grid_pool(batch_dict)  # (BxN, 6x6x6, C)