move max_sweeps filter to PointFeatureEncoder with new config FILTER_…

…SWEEPS
serkankartal · Jan 5, 2022 · ac5198a · ac5198a
1 parent 38c89dd
commit ac5198a
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/pcdet/datasets/dataset.py b/pcdet/datasets/dataset.py
@@ -142,15 +142,6 @@ def prepare_data(self, data_dict):
             if data_dict.get('gt_boxes2d', None) is not None:
                 data_dict['gt_boxes2d'] = data_dict['gt_boxes2d'][selected]
 
-        if 'timestamp' in self.dataset_cfg.POINT_FEATURE_ENCODING.get('src_feature_list'):
-            if data_dict.get('points', None) is not None:
-                max_sweeps = self.dataset_cfg.get('MAX_SWEEPS', 1)
-                idx = self.dataset_cfg.POINT_FEATURE_ENCODING.get('src_feature_list').index('timestamp')
-                dt = np.round(data_dict['points'][:, idx], 2)
-                if np.unique(dt).shape[0] == max_sweeps:
-                    max_dt = sorted(np.unique(dt))[max_sweeps-1]
-                    data_dict['points'] = data_dict['points'][dt <= max_dt]
-
         if data_dict.get('points', None) is not None:
             data_dict = self.point_feature_encoder.forward(data_dict)
 

diff --git a/pcdet/datasets/processor/point_feature_encoder.py b/pcdet/datasets/processor/point_feature_encoder.py
@@ -30,6 +30,14 @@ def forward(self, data_dict):
             data_dict['points']
         )
         data_dict['use_lead_xyz'] = use_lead_xyz
+
+        if self.point_encoding_config.get('FILTER_SWEEPS', False) and 'timestamp' in self.src_feature_list:
+            max_sweeps = self.point_encoding_config.MAX_SWEEPS
+            idx = self.src_feature_list.index('timestamp')
+            dt = np.round(data_dict['points'][:, idx], 2)
+            max_dt = sorted(np.unique(dt))[min(len(np.unique(dt))-1, max_sweeps-1)]
+            data_dict['points'] = data_dict['points'][dt <= max_dt]
+
         return data_dict
 
     def absolute_coordinates_encoding(self, points=None):
@@ -44,4 +52,5 @@ def absolute_coordinates_encoding(self, points=None):
             idx = self.src_feature_list.index(x)
             point_feature_list.append(points[:, idx:idx+1])
         point_features = np.concatenate(point_feature_list, axis=1)
+
         return point_features, True