support to USE_CENTER_TO_FILTER in mask_points_and_boxes_outside_range

chenshi3 · Aug 20, 2022 · 5e21e42 · 5e21e42
1 parent 6bd8be7
commit 5e21e42
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 8 deletions.
diff --git a/pcdet/datasets/processor/data_processor.py b/pcdet/datasets/processor/data_processor.py
@@ -86,7 +86,7 @@ def mask_points_and_boxes_outside_range(self, data_dict=None, config=None):
         if data_dict.get('gt_boxes', None) is not None and config.REMOVE_OUTSIDE_BOXES and self.training:
             mask = box_utils.mask_boxes_outside_range_numpy(
                 data_dict['gt_boxes'], self.point_cloud_range, min_num_corners=config.get('min_num_corners', 1), 
-                filter_z=config.get('FILTER_Z', False)
+                use_center_to_filter=config.get('USE_CENTER_TO_FILTER', True)
             )
             data_dict['gt_boxes'] = data_dict['gt_boxes'][mask]
         return data_dict

diff --git a/pcdet/utils/box_utils.py b/pcdet/utils/box_utils.py
@@ -90,7 +90,7 @@ def corners_rect_to_camera(corners):
     return camera_rect
 
 
-def mask_boxes_outside_range_numpy(boxes, limit_range, min_num_corners=1, filter_z=True):
+def mask_boxes_outside_range_numpy(boxes, limit_range, min_num_corners=1, use_center_to_filter=True):
     """
     Args:
         boxes: (N, 7) [x, y, z, dx, dy, dz, heading, ...], (x, y, z) is the box center
@@ -102,13 +102,14 @@ def mask_boxes_outside_range_numpy(boxes, limit_range, min_num_corners=1, filter
     """
     if boxes.shape[1] > 7:
         boxes = boxes[:, 0:7]
-    corners = boxes_to_corners_3d(boxes)  # (N, 8, 3)
-    if filter_z:
-        mask = ((corners >= limit_range[0:3]) & (corners <= limit_range[3:6])).all(axis=2)
+    if use_center_to_filter:
+        box_centers = boxes[:, 0:3]
+        mask = ((box_centers >= limit_range[0:3]) & (box_centers <= limit_range[3:6])).all(axis=-1)
     else:
+        corners = boxes_to_corners_3d(boxes)  # (N, 8, 3)
         corners = corners[:, :, 0:2]
         mask = ((corners >= limit_range[0:2]) & (corners <= limit_range[3:5])).all(axis=2)
-    mask = mask.sum(axis=1) >= min_num_corners  # (N)
+        mask = mask.sum(axis=1) >= min_num_corners  # (N)
 
     return mask
 

diff --git a/tools/cfgs/dataset_configs/waymo_dataset.yaml b/tools/cfgs/dataset_configs/waymo_dataset.yaml
@@ -62,7 +62,7 @@ POINT_FEATURE_ENCODING: {
 DATA_PROCESSOR:
     - NAME: mask_points_and_boxes_outside_range
       REMOVE_OUTSIDE_BOXES: True
-      FILTER_Z: False
+      USE_CENTER_TO_FILTER: True
 
     - NAME: shuffle_points
       SHUFFLE_ENABLED: {

diff --git a/tools/cfgs/dataset_configs/waymo_dataset_multiframe.yaml b/tools/cfgs/dataset_configs/waymo_dataset_multiframe.yaml
@@ -72,7 +72,7 @@ POINT_FEATURE_ENCODING: {
 DATA_PROCESSOR:
     - NAME: mask_points_and_boxes_outside_range
       REMOVE_OUTSIDE_BOXES: True
-      FILTER_Z: False
+      USE_CENTER_TO_FILTER: True
 
     - NAME: shuffle_points
       SHUFFLE_ENABLED: {