load_points_from_mult_sweeps (#67)

* Add unittest.

* Change to 'defaults to'.

* Add replace=False, remove T op, add num_sweeps=10 unittest.

* Change lyft_infos.pkl.

* Add nus unittest.

* Add nus unittest.

* Change nus_info.pkl.

* Add test_mode.

mmdet3d/datasets/pipelines/loading.py

@@ -71,21 +71,37 @@ class LoadPointsFromMultiSweeps(object):
     This is usually used for nuScenes dataset to utilize previous sweeps.
 
     Args:
-        sweeps_num (int): number of sweeps. Defaults to 10.
-        load_dim (int): dimension number of the loaded points. Defaults to 5.
+        sweeps_num (int): Number of sweeps. Defaults to 10.
+        load_dim (int): Dimension number of the loaded points. Defaults to 5.
+        use_dim (list[int]): Which dimension to use. Defaults to [0, 1, 2, 4].
         file_client_args (dict): Config dict of file clients, refer to
             https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py
             for more details. Defaults to dict(backend='disk').
+        pad_empty_sweeps (bool): Whether to repeat keyframe when
+            sweeps is empty. Defaults to False.
+        remove_close (bool): Whether to remove close points.
+            Defaults to False.
+        test_mode (bool): If test_model=True used for testing, it will not
+            randomly sample sweeps but select the nearest N frames.
+            Defaults to False.
     """
 
     def __init__(self,
                  sweeps_num=10,
                  load_dim=5,
-                 file_client_args=dict(backend='disk')):
+                 use_dim=[0, 1, 2, 4],
+                 file_client_args=dict(backend='disk'),
+                 pad_empty_sweeps=False,
+                 remove_close=False,
+                 test_mode=False):
         self.load_dim = load_dim
         self.sweeps_num = sweeps_num
+        self.use_dim = use_dim
         self.file_client_args = file_client_args.copy()
         self.file_client = None
+        self.pad_empty_sweeps = pad_empty_sweeps
+        self.remove_close = remove_close
+        self.test_mode = test_mode
 
     def _load_points(self, pts_filename):
         """Private function to load point clouds data.
@@ -109,6 +125,22 @@ class LoadPointsFromMultiSweeps(object):
                 points = np.fromfile(pts_filename, dtype=np.float32)
         return points
 
+    def _remove_close(self, points, radius=1.0):
+        """Removes point too close within a certain radius from origin.
+
+        Args:
+            points (np.ndarray): Sweep points.
+            radius (float): Radius below which points are removed.
+                Defaults to 1.0.
+
+        Returns:
+            np.ndarray: Points after removing.
+        """
+        x_filt = np.abs(points[:, 0]) < radius
+        y_filt = np.abs(points[:, 1]) < radius
+        not_close = np.logical_not(np.logical_and(x_filt, y_filt))
+        return points[not_close, :]
+
     def __call__(self, results):
         """Call function to load multi-sweep point clouds from files.
 
@@ -123,25 +155,37 @@ class LoadPointsFromMultiSweeps(object):
                 - points (np.ndarray): Multi-sweep point cloud arrays.
         """
         points = results['points']
-        points[:, 3] /= 255
         points[:, 4] = 0
         sweep_points_list = [points]
         ts = results['timestamp']
-
-        for idx, sweep in enumerate(results['sweeps']):
-            if idx >= self.sweeps_num:
-                break
-            points_sweep = self._load_points(sweep['data_path'])
-            points_sweep = np.copy(points_sweep).reshape(-1, self.load_dim)
-            sweep_ts = sweep['timestamp'] / 1e6
-            points_sweep[:, 3] /= 255
-            points_sweep[:, :3] = points_sweep[:, :3] @ sweep[
-                'sensor2lidar_rotation'].T
-            points_sweep[:, :3] += sweep['sensor2lidar_translation']
-            points_sweep[:, 4] = ts - sweep_ts
-            sweep_points_list.append(points_sweep)
-
-        points = np.concatenate(sweep_points_list, axis=0)[:, [0, 1, 2, 4]]
+        if self.pad_empty_sweeps and len(results['sweeps']) == 0:
+            for i in range(self.sweeps_num):
+                if self.remove_close:
+                    sweep_points_list.append(self._remove_close(points))
+                else:
+                    sweep_points_list.append(points)
+        else:
+            if len(results['sweeps']) <= self.sweeps_num:
+                choices = np.arange(len(results['sweeps']))
+            elif self.test_mode:
+                choices = np.arange(self.sweeps_num)
+            else:
+                choices = np.random.choice(
+                    len(results['sweeps']), self.sweeps_num, replace=False)
+            for idx in choices:
+                sweep = results['sweeps'][idx]
+                points_sweep = self._load_points(sweep['data_path'])
+                points_sweep = np.copy(points_sweep).reshape(-1, self.load_dim)
+                if self.remove_close:
+                    points_sweep = self._remove_close(points_sweep)
+                sweep_ts = sweep['timestamp'] / 1e6
+                points_sweep[:, :3] = points_sweep[:, :3] @ sweep[
+                    'sensor2lidar_rotation'].T
+                points_sweep[:, :3] += sweep['sensor2lidar_translation']
+                points_sweep[:, 4] = ts - sweep_ts
+                sweep_points_list.append(points_sweep)
+
+        points = np.concatenate(sweep_points_list, axis=0)[:, self.use_dim]
         results['points'] = points
         return results
 

tests/test_dataset/test_nuscene_dataset.py

+import numpy as np
+
+from mmdet3d.datasets import NuScenesDataset
+
+
+def test_getitem():
+    np.random.seed(0)
+    point_cloud_range = [-50, -50, -5, 50, 50, 3]
+    file_client_args = dict(backend='disk')
+    class_names = [
+        'car', 'truck', 'trailer', 'bus', 'construction_vehicle', 'bicycle',
+        'motorcycle', 'pedestrian', 'traffic_cone', 'barrier'
+    ]
+    pipeline = [
+        dict(
+            type='LoadPointsFromFile',
+            load_dim=5,
+            use_dim=5,
+            file_client_args=file_client_args),
+        dict(
+            type='LoadPointsFromMultiSweeps',
+            sweeps_num=2,
+            file_client_args=file_client_args),
+        dict(
+            type='MultiScaleFlipAug3D',
+            img_scale=(1333, 800),
+            pts_scale_ratio=1,
+            flip=False,
+            transforms=[
+                dict(
+                    type='GlobalRotScaleTrans',
+                    rot_range=[0, 0],
+                    scale_ratio_range=[1., 1.],
+                    translation_std=[0, 0, 0]),
+                dict(type='RandomFlip3D'),
+                dict(
+                    type='PointsRangeFilter',
+                    point_cloud_range=point_cloud_range),
+                dict(
+                    type='DefaultFormatBundle3D',
+                    class_names=class_names,
+                    with_label=False),
+                dict(type='Collect3D', keys=['points'])
+            ])
+    ]
+
+    nus_dataset = NuScenesDataset(
+        'tests/data/nuscenes/nus_info.pkl',
+        pipeline,
+        'tests/data/nuscenes',
+        test_mode=True)
+    data = nus_dataset[0]
+    assert data['img_metas'][0].data['flip'] is False
+    assert data['img_metas'][0].data['pcd_horizontal_flip'] is False
+    assert data['points'][0]._data.shape == (100, 4)
+
+    data = nus_dataset[1]
+    assert data['img_metas'][0].data['flip'] is False
+    assert data['img_metas'][0].data['pcd_horizontal_flip'] is False
+    assert data['points'][0]._data.shape == (597, 4)

tests/test_pipeline/test_load_points_from_multi_sweeps.py

+import numpy as np
+
+from mmdet3d.datasets.pipelines.loading import LoadPointsFromMultiSweeps
+
+
+def test_load_points_from_multi_sweeps():
+    np.random.seed(0)
+
+    file_client_args = dict(backend='disk')
+    load_points_from_multi_sweeps_1 = LoadPointsFromMultiSweeps(
+        sweeps_num=9,
+        use_dim=[0, 1, 2, 3, 4],
+        file_client_args=file_client_args)
+
+    load_points_from_multi_sweeps_2 = LoadPointsFromMultiSweeps(
+        sweeps_num=9,
+        use_dim=[0, 1, 2, 3, 4],
+        file_client_args=file_client_args,
+        pad_empty_sweeps=True,
+        remove_close=True)
+
+    load_points_from_multi_sweeps_3 = LoadPointsFromMultiSweeps(
+        sweeps_num=9,
+        use_dim=[0, 1, 2, 3, 4],
+        file_client_args=file_client_args,
+        pad_empty_sweeps=True,
+        remove_close=True,
+        test_mode=True)
+
+    points = np.random.random([100, 5]) * 2
+
+    input_results = dict(points=points, sweeps=[], timestamp=None)
+    results = load_points_from_multi_sweeps_1(input_results)
+    assert results['points'].shape == (100, 5)
+
+    input_results = dict(points=points, sweeps=[], timestamp=None)
+    results = load_points_from_multi_sweeps_2(input_results)
+    assert results['points'].shape == (775, 5)
+
+    sensor2lidar_rotation = np.array(
+        [[9.99999967e-01, 1.13183067e-05, 2.56845368e-04],
+         [-1.12839618e-05, 9.99999991e-01, -1.33719456e-04],
+         [-2.56846879e-04, 1.33716553e-04, 9.99999958e-01]])
+    sensor2lidar_translation = np.array([-0.0009198, -0.03964854, -0.00190136])
+    sweep = dict(
+        data_path='tests/data/nuscenes/sweeps/LIDAR_TOP/'
+        'n008-2018-09-18-12-07-26-0400__LIDAR_TOP__'
+        '1537287083900561.pcd.bin',
+        sensor2lidar_rotation=sensor2lidar_rotation,
+        sensor2lidar_translation=sensor2lidar_translation,
+        timestamp=0)
+
+    input_results = dict(points=points, sweeps=[sweep], timestamp=1.0)
+    results = load_points_from_multi_sweeps_1(input_results)
+    assert results['points'].shape == (500, 5)
+
+    input_results = dict(points=points, sweeps=[sweep], timestamp=1.0)
+    results = load_points_from_multi_sweeps_2(input_results)
+    assert results['points'].shape == (451, 5)
+
+    input_results = dict(points=points, sweeps=[sweep] * 10, timestamp=1.0)
+    results = load_points_from_multi_sweeps_2(input_results)
+    assert results['points'].shape == (3259, 5)
+
+    input_results = dict(points=points, sweeps=[sweep] * 10, timestamp=1.0)
+    results = load_points_from_multi_sweeps_3(input_results)
+    assert results['points'].shape == (3259, 5)