[Feature] Support LaserMix augmentation (open-mmlab#2302)

* add lasermix

* add prob

* update description

* update

mmdet3d/datasets/seg3d_dataset.py

@@ -295,7 +295,7 @@ def prepare_data(self, idx: int) -> dict:
         if not self.test_mode:
             data_info = self.get_data_info(idx)
             # Pass the dataset to the pipeline during training to support mixed
-            # data augmentation, such as polarmix.
+            # data augmentation, such as polarmix and lasermix.
             data_info['dataset'] = self
             return self.pipeline(data_info)
         else:

mmdet3d/datasets/transforms/__init__.py

@@ -11,8 +11,8 @@
 from .transforms_3d import (AffineResize, BackgroundPointsFilter,
                             GlobalAlignment, GlobalRotScaleTrans,
                             IndoorPatchPointSample, IndoorPointSample,
-                            MultiViewWrapper, ObjectNameFilter, ObjectNoise,
-                            ObjectRangeFilter, ObjectSample,
+                            LaserMix, MultiViewWrapper, ObjectNameFilter,
+                            ObjectNoise, ObjectRangeFilter, ObjectSample,
                             PhotoMetricDistortion3D, PointSample, PointShuffle,
                             PointsRangeFilter, PolarMix, RandomDropPointsColor,
                             RandomFlip3D, RandomJitterPoints, RandomResize3D,
@@ -30,5 +30,5 @@
     'RandomDropPointsColor', 'RandomJitterPoints', 'AffineResize',
     'RandomShiftScale', 'LoadPointsFromDict', 'Resize3D', 'RandomResize3D',
     'MultiViewWrapper', 'PhotoMetricDistortion3D', 'MonoDet3DInferencerLoader',
-    'LidarDet3DInferencerLoader', 'PolarMix'
+    'LidarDet3DInferencerLoader', 'PolarMix', 'LaserMix'
 ]

mmdet3d/datasets/transforms/transforms_3d.py

@@ -2521,3 +2521,148 @@ def __repr__(self) -> str:
         repr_str += f'pre_transform={self.pre_transform}, '
         repr_str += f'prob={self.prob})'
         return repr_str
+
+
+@TRANSFORMS.register_module()
+class LaserMix(BaseTransform):
+    """LaserMix data augmentation.
+
+    The lasermix transform steps are as follows:
+
+        1. Another random point cloud is picked by dataset.
+        2. Divide the point cloud into several regions according to pitch
+           angles and combine the areas crossly.
+
+    Required Keys:
+
+    - points (:obj:`BasePoints`)
+    - pts_semantic_mask (np.int64)
+    - dataset (:obj:`BaseDataset`)
+
+    Modified Keys:
+
+    - points (:obj:`BasePoints`)
+    - pts_semantic_mask (np.int64)
+
+    Args:
+        num_areas (List[int]): A list of area numbers will be divided into.
+        pitch_angles (Sequence[float]): Pitch angles used to divide areas.
+        pre_transform (Sequence[dict], optional): Sequence of transform object
+            or config dict to be composed. Defaults to None.
+        prob (float): The transformation probability. Defaults to 1.0.
+    """
+
+    def __init__(self,
+                 num_areas: List[int],
+                 pitch_angles: Sequence[float],
+                 pre_transform: Optional[Sequence[dict]] = None,
+                 prob: float = 1.0) -> None:
+        assert is_list_of(num_areas, int), \
+            'num_areas should be a list of int.'
+        self.num_areas = num_areas
+
+        assert len(pitch_angles) == 2, \
+            'The length of pitch_angles should be 2, ' \
+            f'but got {len(pitch_angles)}.'
+        assert pitch_angles[1] > pitch_angles[0], \
+            'pitch_angles[1] should be larger than pitch_angles[0].'
+        self.pitch_angles = pitch_angles
+
+        self.prob = prob
+        if pre_transform is None:
+            self.pre_transform = None
+        else:
+            self.pre_transform = Compose(pre_transform)
+
+    def laser_mix_transform(self, input_dict: dict, mix_results: dict) -> dict:
+        """LaserMix transform function.
+
+        Args:
+            input_dict (dict): Result dict from loading pipeline.
+            mix_results (dict): Mixed dict picked from dataset.
+
+        Returns:
+            dict: output dict after transformation.
+        """
+        mix_points = mix_results['points']
+        mix_pts_semantic_mask = mix_results['pts_semantic_mask']
+
+        points = input_dict['points']
+        pts_semantic_mask = input_dict['pts_semantic_mask']
+
+        rho = torch.sqrt(points.coord[:, 0]**2 + points.coord[:, 1]**2)
+        pitch = torch.atan2(points.coord[:, 2], rho)
+        pitch = torch.clip(pitch, self.pitch_angles[0] + 1e-5,
+                           self.pitch_angles[1] - 1e-5)
+
+        mix_rho = torch.sqrt(mix_points.coord[:, 0]**2 +
+                             mix_points.coord[:, 1]**2)
+        mix_pitch = torch.atan2(mix_points.coord[:, 2], mix_rho)
+        mix_pitch = torch.clip(mix_pitch, self.pitch_angles[0] + 1e-5,
+                               self.pitch_angles[1] - 1e-5)
+
+        num_areas = np.random.choice(self.num_areas, size=1)[0]
+        angle_list = np.linspace(self.pitch_angles[1], self.pitch_angles[0],
+                                 num_areas + 1)
+        out_points = []
+        out_pts_semantic_mask = []
+        for i in range(num_areas):
+            # convert angle to radian
+            start_angle = angle_list[i + 1] / 180 * np.pi
+            end_angle = angle_list[i] / 180 * np.pi
+            if i % 2 == 0:  # pick from original point cloud
+                idx = (pitch > start_angle) & (pitch <= end_angle)
+                out_points.append(points[idx])
+                out_pts_semantic_mask.append(pts_semantic_mask[idx.numpy()])
+            else:  # pickle from mixed point cloud
+                idx = (mix_pitch > start_angle) & (mix_pitch <= end_angle)
+                out_points.append(mix_points[idx])
+                out_pts_semantic_mask.append(
+                    mix_pts_semantic_mask[idx.numpy()])
+        out_points = points.cat(out_points)
+        out_pts_semantic_mask = np.concatenate(out_pts_semantic_mask, axis=0)
+        input_dict['points'] = out_points
+        input_dict['pts_semantic_mask'] = out_pts_semantic_mask
+        return input_dict
+
+    def transform(self, input_dict: dict) -> dict:
+        """LaserMix transform function.
+
+        Args:
+            input_dict (dict): Result dict from loading pipeline.
+
+        Returns:
+            dict: output dict after transformation.
+        """
+        if np.random.rand() > self.prob:
+            return input_dict
+
+        assert 'dataset' in input_dict, \
+            '`dataset` is needed to pass through LaserMix, while not found.'
+        dataset = input_dict['dataset']
+
+        # get index of other point cloud
+        index = np.random.randint(0, len(dataset))
+
+        mix_results = dataset.get_data_info(index)
+
+        if self.pre_transform is not None:
+            # pre_transform may also require dataset
+            mix_results.update({'dataset': dataset})
+            # before lasermix need to go through
+            # the necessary pre_transform
+            mix_results = self.pre_transform(mix_results)
+            mix_results.pop('dataset')
+
+        input_dict = self.laser_mix_transform(input_dict, mix_results)
+
+        return input_dict
+
+    def __repr__(self) -> str:
+        """str: Return a string that describes the module."""
+        repr_str = self.__class__.__name__
+        repr_str += f'(num_areas={self.num_areas}, '
+        repr_str += f'pitch_angles={self.pitch_angles}, '
+        repr_str += f'pre_transform={self.pre_transform}, '
+        repr_str += f'prob={self.prob})'
+        return repr_str

tests/test_datasets/test_transforms/test_transforms_3d.py

@@ -8,7 +8,7 @@
 
 from mmdet3d.datasets import (GlobalAlignment, RandomFlip3D,
                               SemanticKITTIDataset)
-from mmdet3d.datasets.transforms import GlobalRotScaleTrans, PolarMix
+from mmdet3d.datasets.transforms import GlobalRotScaleTrans, LaserMix, PolarMix
 from mmdet3d.structures import LiDARPoints
 from mmdet3d.testing import create_data_info_after_loading
 from mmdet3d.utils import register_all_modules
@@ -222,3 +222,128 @@ def test_transform(self):
         results = transform.transform(copy.deepcopy(self.results))
         self.assertTrue(results['points'].shape[0] ==
                         results['pts_semantic_mask'].shape[0])
+
+
+class TestLaserMix(unittest.TestCase):
+
+    def setUp(self):
+        self.pre_transform = [
+            dict(
+                type='LoadPointsFromFile',
+                coord_type='LIDAR',
+                load_dim=4,
+                use_dim=4),
+            dict(
+                type='LoadAnnotations3D',
+                with_bbox_3d=False,
+                with_label_3d=False,
+                with_mask_3d=False,
+                with_seg_3d=True,
+                seg_3d_dtype='np.int32'),
+            dict(type='PointSegClassMapping'),
+        ]
+        classes = ('unlabeled', 'car', 'bicycle', 'motorcycle', 'truck', 'bus',
+                   'person', 'bicyclist', 'motorcyclist', 'road', 'parking',
+                   'sidewalk', 'other-ground', 'building', 'fence',
+                   'vegetation', 'trunck', 'terrian', 'pole', 'traffic-sign')
+        palette = [
+            [174, 199, 232],
+            [152, 223, 138],
+            [31, 119, 180],
+            [255, 187, 120],
+            [188, 189, 34],
+            [140, 86, 75],
+            [255, 152, 150],
+            [214, 39, 40],
+            [197, 176, 213],
+            [148, 103, 189],
+            [196, 156, 148],
+            [23, 190, 207],
+            [247, 182, 210],
+            [219, 219, 141],
+            [255, 127, 14],
+            [158, 218, 229],
+            [44, 160, 44],
+            [112, 128, 144],
+            [227, 119, 194],
+            [82, 84, 163],
+        ]
+        seg_label_mapping = {
+            0: 0,  # "unlabeled"
+            1: 0,  # "outlier" mapped to "unlabeled" --------------mapped
+            10: 1,  # "car"
+            11: 2,  # "bicycle"
+            13: 5,  # "bus" mapped to "other-vehicle" --------------mapped
+            15: 3,  # "motorcycle"
+            16: 5,  # "on-rails" mapped to "other-vehicle" ---------mapped
+            18: 4,  # "truck"
+            20: 5,  # "other-vehicle"
+            30: 6,  # "person"
+            31: 7,  # "bicyclist"
+            32: 8,  # "motorcyclist"
+            40: 9,  # "road"
+            44: 10,  # "parking"
+            48: 11,  # "sidewalk"
+            49: 12,  # "other-ground"
+            50: 13,  # "building"
+            51: 14,  # "fence"
+            52: 0,  # "other-structure" mapped to "unlabeled" ------mapped
+            60: 9,  # "lane-marking" to "road" ---------------------mapped
+            70: 15,  # "vegetation"
+            71: 16,  # "trunk"
+            72: 17,  # "terrain"
+            80: 18,  # "pole"
+            81: 19,  # "traffic-sign"
+            99: 0,  # "other-object" to "unlabeled" ----------------mapped
+            252: 1,  # "moving-car" to "car" ------------------------mapped
+            253: 7,  # "moving-bicyclist" to "bicyclist" ------------mapped
+            254: 6,  # "moving-person" to "person" ------------------mapped
+            255: 8,  # "moving-motorcyclist" to "motorcyclist" ------mapped
+            256: 5,  # "moving-on-rails" mapped to "other-vehic------mapped
+            257: 5,  # "moving-bus" mapped to "other-vehicle" -------mapped
+            258: 4,  # "moving-truck" to "truck" --------------------mapped
+            259: 5  # "moving-other"-vehicle to "other-vehicle"-----mapped
+        }
+        max_label = 259
+        self.dataset = SemanticKITTIDataset(
+            './tests/data/semantickitti/',
+            'semantickitti_infos.pkl',
+            metainfo=dict(
+                classes=classes,
+                palette=palette,
+                seg_label_mapping=seg_label_mapping,
+                max_label=max_label),
+            data_prefix=dict(
+                pts='sequences/00/velodyne',
+                pts_semantic_mask='sequences/00/labels'),
+            pipeline=[],
+            modality=dict(use_lidar=True, use_camera=False))
+        points = np.random.random((100, 4))
+        self.results = {
+            'points': LiDARPoints(points, points_dim=4),
+            'pts_semantic_mask': np.random.randint(0, 20, (100, )),
+            'dataset': self.dataset
+        }
+
+    def test_transform(self):
+        # test assertion for invalid num_areas
+        with self.assertRaises(AssertionError):
+            transform = LaserMix(num_areas=3, pitch_angles=[-20, 0])
+
+        with self.assertRaises(AssertionError):
+            transform = LaserMix(num_areas=[3.0, 4.0], pitch_angles=[-20, 0])
+
+        # test assertion for invalid pitch_angles
+        with self.assertRaises(AssertionError):
+            transform = LaserMix(num_areas=[3, 4], pitch_angles=[-20])
+
+        with self.assertRaises(AssertionError):
+            transform = LaserMix(num_areas=[3, 4], pitch_angles=[0, -20])
+
+        transform = LaserMix(
+            num_areas=[3, 4, 5, 6],
+            pitch_angles=[-20, 0],
+            pre_transform=self.pre_transform)
+        results = transform.transform(copy.deepcopy(self.results))
+        self.assertTrue(results['points'].shape[0] ==
+                        results['pts_semantic_mask'].shape[0])