add sample_points data_processor from PointRCNN

pcdet/datasets/processor/data_processor.py

@@ -62,6 +62,34 @@ class DataProcessor(object):
         data_dict['voxel_num_points'] = num_points
         return data_dict
 
+    def sample_points(self, data_dict=None, config=None):
+        if data_dict is None:
+            return partial(self.sample_points, config=config)
+
+        num_points = config.NUM_POINTS[self.mode]
+        if num_points == -1:
+            return data_dict
+
+        points = data_dict['points']
+        if num_points < len(points):
+            pts_depth = np.linalg.norm(points[:, 0:3], axis=1)
+            pts_near_flag = pts_depth < 40.0
+            far_idxs_choice = np.where(pts_near_flag == 0)[0]
+            near_idxs = np.where(pts_near_flag == 1)[0]
+            near_idxs_choice = np.random.choice(near_idxs, num_points - len(far_idxs_choice), replace=False)
+
+            choice = np.concatenate((near_idxs_choice, far_idxs_choice), axis=0) \
+                if len(far_idxs_choice) > 0 else near_idxs_choice
+            np.random.shuffle(choice)
+        else:
+            choice = np.arange(0, len(points), dtype=np.int32)
+            if num_points > len(points):
+                extra_choice = np.random.choice(choice, num_points - len(points), replace=False)
+                choice = np.concatenate((choice, extra_choice), axis=0)
+            np.random.shuffle(choice)
+        data_dict['points'] = points[choice]
+        return data_dict
+
     def forward(self, data_dict):
         """
         Args: