Merge branch 'indoor_augment' into 'master'

Indoor augment

See merge request open-mmlab/mmdet.3d!11

mmdet3d/datasets/__init__.py

@@ -5,11 +5,12 @@ from .kitti2d_dataset import Kitti2DDataset
 from .kitti_dataset import KittiDataset
 from .loader import DistributedGroupSampler, GroupSampler, build_dataloader
 from .nuscenes_dataset import NuScenesDataset
-from .pipelines import (GlobalRotScale, IndoorLoadAnnotations3D,
-                        IndoorLoadPointsFromFile, IndoorPointsColorNormalize,
-                        ObjectNoise, ObjectRangeFilter, ObjectSample,
-                        PointSample, PointShuffle, PointsRangeFilter,
-                        RandomFlip3D)
+from .pipelines import (GlobalRotScale, IndoorFlipData, IndoorGlobalRotScale,
+                        IndoorLoadAnnotations3D, IndoorLoadPointsFromFile,
+                        IndoorPointSample, IndoorPointsColorJitter,
+                        IndoorPointsColorNormalize, ObjectNoise,
+                        ObjectRangeFilter, ObjectSample, PointShuffle,
+                        PointsRangeFilter, RandomFlip3D)
 
 __all__ = [
     'KittiDataset', 'GroupSampler', 'DistributedGroupSampler',
@@ -17,6 +18,7 @@ __all__ = [
     'CocoDataset', 'Kitti2DDataset', 'NuScenesDataset', 'ObjectSample',
     'RandomFlip3D', 'ObjectNoise', 'GlobalRotScale', 'PointShuffle',
     'ObjectRangeFilter', 'PointsRangeFilter', 'Collect3D',
-    'IndoorLoadPointsFromFile', 'IndoorPointsColorNormalize', 'PointSample',
-    'IndoorLoadAnnotations3D'
+    'IndoorLoadPointsFromFile', 'IndoorPointsColorNormalize',
+    'IndoorPointSample', 'IndoorLoadAnnotations3D', 'IndoorPointsColorJitter',
+    'IndoorGlobalRotScale', 'IndoorFlipData'
 ]

mmdet3d/datasets/pipelines/__init__.py

 from mmdet.datasets.pipelines import Compose
 from .dbsampler import DataBaseSampler, MMDataBaseSampler
 from .formating import DefaultFormatBundle, DefaultFormatBundle3D
+from .indoor_augment import (IndoorFlipData, IndoorGlobalRotScale,
+                             IndoorPointsColorJitter)
 from .indoor_loading import (IndoorLoadAnnotations3D, IndoorLoadPointsFromFile,
                              IndoorPointsColorNormalize)
-from .indoor_sample import PointSample
+from .indoor_sample import IndoorPointSample
 from .loading import LoadMultiViewImageFromFiles, LoadPointsFromFile
 from .train_aug import (GlobalRotScale, ObjectNoise, ObjectRangeFilter,
                         ObjectSample, PointShuffle, PointsRangeFilter,
@@ -14,6 +16,8 @@ __all__ = [
     'PointShuffle', 'ObjectRangeFilter', 'PointsRangeFilter', 'Collect3D',
     'Compose', 'LoadMultiViewImageFromFiles', 'LoadPointsFromFile',
     'DefaultFormatBundle', 'DefaultFormatBundle3D', 'DataBaseSampler',
+    'IndoorGlobalRotScale', 'IndoorPointsColorJitter', 'IndoorFlipData',
     'MMDataBaseSampler', 'IndoorLoadPointsFromFile',
-    'IndoorPointsColorNormalize', 'IndoorLoadAnnotations3D', 'PointSample'
+    'IndoorPointsColorNormalize', 'IndoorLoadAnnotations3D',
+    'IndoorPointSample'
 ]

mmdet3d/datasets/pipelines/indoor_augment.py

+import numpy as np
+
+from mmdet.datasets.builder import PIPELINES
+
+
+@PIPELINES.register_module()
+class IndoorFlipData(object):
+    """Indoor Flip Data.
+
+    Flip point cloud and ground truth boxes.
+    The point cloud will ve flipped along the yz plane
+    and the xz plane with a certain probability.
+
+    Args:
+        flip_ratio_yz (float): Probability of being flipped along yz plane.
+            Default: 0.5.
+        flip_ratio_xz (float): Probability of being flipped along xz plane.
+            Default: 0.5.
+    """
+
+    def __init__(self, flip_ratio_yz=0.5, flip_ratio_xz=0.5):
+        self.flip_ratio_yz = flip_ratio_yz
+        self.flip_ratio_xz = flip_ratio_xz
+
+    def __call__(self, results):
+        points = results['points']
+        gt_bboxes_3d = results['gt_bboxes_3d']
+        aligned = True if gt_bboxes_3d.shape[1] == 6 else False
+        if np.random.random() < self.flip_ratio_yz:
+            # Flipping along the YZ plane
+            points[:, 0] = -1 * points[:, 0]
+            gt_bboxes_3d[:, 0] = -1 * gt_bboxes_3d[:, 0]
+            if not aligned:
+                gt_bboxes_3d[:, 6] = np.pi - gt_bboxes_3d[:, 6]
+            results['flip_yz'] = True
+
+        if aligned and np.random.random() < self.flip_ratio_xz:
+            # Flipping along the XZ plane
+            points[:, 1] = -1 * points[:, 1]
+            gt_bboxes_3d[:, 1] = -1 * gt_bboxes_3d[:, 1]
+            results['flip_xz'] = True
+
+        results['points'] = points
+        results['gt_bboxes_3d'] = gt_bboxes_3d
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += '(flip_ratio_yz={})'.format(self.flip_ratio_yz)
+        repr_str += '(flip_ratio_xz={})'.format(self.flip_ratio_xz)
+        return repr_str
+
+
+@PIPELINES.register_module()
+class IndoorPointsColorJitter(object):
+    """Indoor Points Color Jitter.
+
+    Randomly change the brightness and color of the point cloud, and
+    drop out the points' colors with a certain range and probability.
+
+    Args:
+        color_mean (List[float]): Mean color of the point cloud.
+            Default: [0.5, 0.5, 0.5].
+        bright_range (List[float]): Range of brightness.
+            Default: [0.8, 1.2].
+        color_shift_range (List[float]): Range of color shift.
+            Default: [0.95, 1.05].
+        jitter_range (List[float]): Range of jittering.
+            Default: [-0.025, 0.025].
+        drop_prob (float): Probability to drop out points' color.
+            Default: 0.3
+    """
+
+    def __init__(self,
+                 color_mean=[0.5, 0.5, 0.5],
+                 bright_range=[0.8, 1.2],
+                 color_shift_range=[0.95, 1.05],
+                 jitter_range=[-0.025, 0.025],
+                 drop_prob=0.3):
+        self.color_mean = color_mean
+        self.bright_range = bright_range
+        self.color_shift_range = color_shift_range
+        self.jitter_range = jitter_range
+        self.drop_prob = drop_prob
+
+    def __call__(self, results):
+        points = results['points']
+        assert points.shape[1] >= 6, \
+            f'Expect points have channel >=6, got {points.shape[1]}.'
+        rgb_color = points[:, 3:6] + self.color_mean
+        # brightness change for each channel
+        rgb_color *= np.random.uniform(self.bright_range[0],
+                                       self.bright_range[1], 3)
+        # color shift for each channel
+        rgb_color += np.random.uniform(self.color_shift_range[0],
+                                       self.color_shift_range[1], 3)
+        # jittering on each pixel
+        rgb_color += np.expand_dims(
+            np.random.uniform(self.jitter_range[0], self.jitter_range[1]), -1)
+        rgb_color = np.clip(rgb_color, 0, 1)
+        # randomly drop out points' colors
+        rgb_color *= np.expand_dims(
+            np.random.random(points.shape[0]) > self.drop_prob, -1)
+        points[:, 3:6] = rgb_color - self.color_mean
+        results['points'] = points
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += '(color_mean={})'.format(self.color_mean)
+        repr_str += '(bright_range={})'.format(self.bright_range)
+        repr_str += '(color_shift_range={})'.format(self.color_shift_range)
+        repr_str += '(jitter_range={})'.format(self.jitter_range)
+        repr_str += '(drop_prob={})'.format(self.drop_prob)
+
+
+# TODO: merge outdoor indoor transform.
+# TODO: try transform noise.
+@PIPELINES.register_module()
+class IndoorGlobalRotScale(object):
+    """Indoor Global Rotate Scale.
+
+    Augment sunrgbd and scannet data with global rotating and scaling.
+
+    Args:
+        use_height (bool): Whether to use height.
+            Default: True.
+        rot_range (List[float]): Range of rotation.
+            Default: None.
+        scale_range (List[float]): Range of scale.
+            Default: None.
+    """
+
+    def __init__(self, use_height=True, rot_range=None, scale_range=None):
+        self.use_height = use_height
+        self.rot_range = rot_range
+        self.scale_range = scale_range
+
+    def _rotz(self, t):
+        """Rotate About Z.
+
+        Rotation about the z-axis.
+
+        Args:
+            t (float): Angle of rotation.
+
+        Returns:
+            rot_mat (ndarray): Matrix of rotation.
+        """
+        c = np.cos(t)
+        s = np.sin(t)
+        rot_mat = np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])
+        return rot_mat
+
+    def _rotate_aligned_boxes(self, input_boxes, rot_mat):
+        """Rotate Aligned Boxes.
+
+        Rotate function for the aligned boxes.
+
+        Args:
+            input_boxes (ndarray): 3D boxes.
+            rot_mat (ndarray): Rotation matrix.
+
+        Returns:
+            rotated_boxes (ndarry): 3D boxes after rotation.
+        """
+        centers, lengths = input_boxes[:, 0:3], input_boxes[:, 3:6]
+        new_centers = np.dot(centers, rot_mat.T)
+
+        dx, dy = lengths[:, 0] / 2.0, lengths[:, 1] / 2.0
+        new_x = np.zeros((dx.shape[0], 4))
+        new_y = np.zeros((dx.shape[0], 4))
+
+        for i, corner in enumerate([(-1, -1), (1, -1), (1, 1), (-1, 1)]):
+            corners = np.zeros((dx.shape[0], 3))
+            corners[:, 0] = corner[0] * dx
+            corners[:, 1] = corner[1] * dy
+            corners = np.dot(corners, rot_mat.T)
+            new_x[:, i] = corners[:, 0]
+            new_y[:, i] = corners[:, 1]
+
+        new_dx = 2.0 * np.max(new_x, 1)
+        new_dy = 2.0 * np.max(new_y, 1)
+        new_lengths = np.stack((new_dx, new_dy, lengths[:, 2]), axis=1)
+
+        return np.concatenate([new_centers, new_lengths], axis=1)
+
+    def __call__(self, results):
+        points = results['points']
+        gt_bboxes_3d = results['gt_bboxes_3d']
+        aligned = True if gt_bboxes_3d.shape[1] == 6 else False
+
+        if self.rot_range is not None:
+            assert len(self.rot_range) == 2, \
+                f'Expect length of rot range =2, ' \
+                f'got {len(self.rot_range)}.'
+            rot_angle = np.random.uniform(self.rot_range[0], self.rot_range[1])
+            rot_mat = self._rotz(rot_angle)
+            points[:, :3] = np.dot(points[:, :3], rot_mat.T)
+            if aligned:
+                gt_bboxes_3d = self._rotate_aligned_boxes(
+                    gt_bboxes_3d, rot_mat)
+            else:
+                gt_bboxes_3d[:, :3] = np.dot(gt_bboxes_3d[:, :3], rot_mat.T)
+                gt_bboxes_3d[:, 6] -= rot_angle
+
+        if self.scale_range is not None:
+            assert len(self.scale_range) == 2, \
+                f'Expect length of scale range =2, ' \
+                f'got {len(self.scale_range)}.'
+            # Augment point cloud scale
+            scale_ratio = np.random.uniform(self.scale_range[0],
+                                            self.scale_range[1])
+
+            points[:, :3] *= scale_ratio
+            gt_bboxes_3d[:, :3] *= scale_ratio
+            gt_bboxes_3d[:, 3:6] *= scale_ratio
+            if self.use_height:
+                points[:, -1] *= scale_ratio
+
+        results['points'] = points
+        results['gt_bboxes_3d'] = gt_bboxes_3d
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += '(use_height={})'.format(self.use_height)
+        repr_str += '(rot_range={})'.format(self.rot_range)
+        repr_str += '(scale_range={})'.format(self.scale_range)
+        return repr_str

mmdet3d/datasets/pipelines/indoor_sample.py

@@ -4,7 +4,7 @@ from mmdet.datasets.builder import PIPELINES
 
 
 @PIPELINES.register_module()
-class PointSample(object):
+class IndoorPointSample(object):
     """Point Sample.
 
     Sampling data to a certain number.
@@ -46,7 +46,7 @@ class PointSample(object):
             return points[choices]
 
     def __call__(self, results):
-        points = results.get('points', None)
+        points = results['points']
         points, choices = self.points_random_sampling(
             points, self.num_points, return_choices=True)
         pts_instance_mask = results.get('pts_instance_mask', None)

tests/test_indoor_augment.py

+import numpy as np
+
+from mmdet3d.datasets.pipelines import IndoorFlipData, IndoorGlobalRotScale
+
+
+def test_indoor_flip_data():
+    np.random.seed(0)
+    sunrgbd_indoor_flip_data = IndoorFlipData(1, 1)
+    sunrgbd_results = dict()
+    sunrgbd_results['points'] = np.array(
+        [[1.02828765e+00, 3.65790772e+00, 1.97294697e-01, 1.61959505e+00],
+         [-3.95979017e-01, 1.05465031e+00, -7.49204338e-01, 6.73096001e-01]])
+    sunrgbd_results['gt_bboxes_3d'] = np.array([[
+        0.213684, 1.036364, -0.982323, 0.61541, 0.572574, 0.872728, 3.07028526
+    ],
+                                                [
+                                                    -0.449953, 1.395455,
+                                                    -1.027778, 1.500956,
+                                                    1.637298, 0.636364,
+                                                    -1.58242359
+                                                ]])
+    sunrgbd_results = sunrgbd_indoor_flip_data(sunrgbd_results)
+    sunrgbd_points = sunrgbd_results['points']
+    sunrgbd_gt_bboxes_3d = sunrgbd_results['gt_bboxes_3d']
+
+    expected_sunrgbd_points = np.array(
+        [[-1.02828765, 3.65790772, 0.1972947, 1.61959505],
+         [0.39597902, 1.05465031, -0.74920434, 0.673096]])
+    expected_sunrgbd_gt_bboxes_3d = np.array([[
+        -0.213684, 1.036364, -0.982323, 0.61541, 0.572574, 0.872728, 0.07130739
+    ], [
+        0.449953, 1.395455, -1.027778, 1.500956, 1.637298, 0.636364, 4.72401624
+    ]])
+    assert np.allclose(sunrgbd_points, expected_sunrgbd_points)
+    assert np.allclose(sunrgbd_gt_bboxes_3d, expected_sunrgbd_gt_bboxes_3d)
+
+    np.random.seed(0)
+    scannet_indoor_flip_data = IndoorFlipData(1, 1)
+    scannet_results = dict()
+    scannet_results['points'] = np.array(
+        [[1.6110241e+00, -1.6903955e-01, 5.8115810e-01, 5.9897250e-01],
+         [1.3978075e+00, 4.2035791e-01, 3.8729519e-01, 4.0510958e-01]])
+    scannet_results['gt_bboxes_3d'] = np.array([[
+        0.55903838, 0.48201692, 0.65688646, 0.65370704, 0.60029864, 0.5163464
+    ], [
+        -0.03226406, 1.70392646, 0.60348618, 0.65165804, 0.72084366, 0.64667457
+    ]])
+    scannet_results = scannet_indoor_flip_data(scannet_results)
+    scannet_points = scannet_results['points']
+    scannet_gt_bboxes_3d = scannet_results['gt_bboxes_3d']
+
+    expected_scannet_points = np.array(
+        [[-1.6110241, 0.16903955, 0.5811581, 0.5989725],
+         [-1.3978075, -0.42035791, 0.38729519, 0.40510958]])
+    expected_scannet_gt_bboxes_3d = np.array([[
+        -0.55903838, -0.48201692, 0.65688646, 0.65370704, 0.60029864, 0.5163464
+    ], [
+        0.03226406, -1.70392646, 0.60348618, 0.65165804, 0.72084366, 0.64667457
+    ]])
+    assert np.allclose(scannet_points, expected_scannet_points)
+    assert np.allclose(scannet_gt_bboxes_3d, expected_scannet_gt_bboxes_3d)
+
+
+def test_global_rot_scale():
+    np.random.seed(0)
+    sunrgbd_augment = IndoorGlobalRotScale(
+        True, rot_range=[-np.pi / 6, np.pi / 6], scale_range=[0.85, 1.15])
+    sunrgbd_results = dict()
+    sunrgbd_results['points'] = np.array(
+        [[1.02828765e+00, 3.65790772e+00, 1.97294697e-01, 1.61959505e+00],
+         [-3.95979017e-01, 1.05465031e+00, -7.49204338e-01, 6.73096001e-01]])
+    sunrgbd_results['gt_bboxes_3d'] = np.array([[
+        0.213684, 1.036364, -0.982323, 0.61541, 0.572574, 0.872728, 3.07028526
+    ],
+                                                [
+                                                    -0.449953, 1.395455,
+                                                    -1.027778, 1.500956,
+                                                    1.637298, 0.636364,
+                                                    -1.58242359
+                                                ]])
+
+    sunrgbd_results = sunrgbd_augment(sunrgbd_results)
+    sunrgbd_points = sunrgbd_results['points']
+    sunrgbd_gt_bboxes_3d = sunrgbd_results['gt_bboxes_3d']
+
+    expected_sunrgbd_points = np.array(
+        [[0.89427376, 3.94489646, 0.21003141, 1.72415094],
+         [-0.47835783, 1.09972989, -0.79757058, 0.71654893]])
+    expected_sunrgbd_gt_bboxes_3d = np.array([[
+        0.17080999, 1.11345031, -1.04573864, 0.65513891, 0.60953755,
+        0.92906854, 3.01916788
+    ],
+                                              [
+                                                  -0.55427876, 1.45912611,
+                                                  -1.09412807, 1.59785293,
+                                                  1.74299674, 0.67744563,
+                                                  -1.63354097
+                                              ]])
+    assert np.allclose(sunrgbd_points, expected_sunrgbd_points)
+    assert np.allclose(sunrgbd_gt_bboxes_3d, expected_sunrgbd_gt_bboxes_3d)
+
+    np.random.seed(0)
+    scannet_augment = IndoorGlobalRotScale(
+        True, rot_range=[-np.pi * 1 / 36, np.pi * 1 / 36], scale_range=None)
+    scannet_results = dict()
+    scannet_results['points'] = np.array(
+        [[1.6110241e+00, -1.6903955e-01, 5.8115810e-01, 5.9897250e-01],
+         [1.3978075e+00, 4.2035791e-01, 3.8729519e-01, 4.0510958e-01]])
+    scannet_results['gt_bboxes_3d'] = np.array([[
+        0.55903838, 0.48201692, 0.65688646, 0.65370704, 0.60029864, 0.5163464
+    ], [
+        -0.03226406, 1.70392646, 0.60348618, 0.65165804, 0.72084366, 0.64667457
+    ]])
+    scannet_results = scannet_augment(scannet_results)
+    scannet_points = scannet_results['points']
+    scannet_gt_bboxes_3d = scannet_results['gt_bboxes_3d']
+
+    expected_scannet_points = np.array(
+        [[1.61240576, -0.15530836, 0.5811581, 0.5989725],
+         [1.39417555, 0.43225122, 0.38729519, 0.40510958]])
+    expected_scannet_gt_bboxes_3d = np.array([[
+        0.55491157, 0.48676213, 0.65688646, 0.65879754, 0.60584609, 0.5163464
+    ], [
+        -0.04677942, 1.70358975, 0.60348618, 0.65777559, 0.72636927, 0.64667457
+    ]])
+    assert np.allclose(scannet_points, expected_scannet_points)
+    assert np.allclose(scannet_gt_bboxes_3d, expected_scannet_gt_bboxes_3d)

tests/test_indoor_loading.py

@@ -65,7 +65,7 @@ def test_load_annotations3D():
     scannet_results['pts_instance_mask_path'] = osp.join(
         data_path, f'{scan_name}_ins_label.npy')
     scannet_results['pts_semantic_mask_path'] = osp.join(
-        data_path, scan_name + '_sem_label.npy')
+        data_path, f'{scan_name}_sem_label.npy')
     scannet_results['info'] = scannet_info
     scannet_results['gt_bboxes_3d'] = scannet_gt_bboxes_3d
     scannet_results['gt_labels'] = scannet_gt_labels

tests/test_indoor_sample.py

 import numpy as np
 
-from mmdet3d.datasets.pipelines import PointSample
+from mmdet3d.datasets.pipelines import IndoorPointSample
 
 
 def test_indoor_sample():
     np.random.seed(0)
-    scannet_sample_points = PointSample(5)
+    scannet_sample_points = IndoorPointSample(5)
     scannet_results = dict()
     scannet_points = np.array([[1.0719866, -0.7870435, 0.8408122, 0.9196809],
                                [1.103661, 0.81065744, 2.6616862, 2.7405548],
@@ -24,11 +24,9 @@ def test_indoor_sample():
     scannet_pts_semantic_mask = np.array([38, 1, 1, 40, 0, 40, 1, 1, 1, 0])
     scannet_results['pts_semantic_mask'] = scannet_pts_semantic_mask
     scannet_results = scannet_sample_points(scannet_results)
-    scannet_points_result = scannet_results.get('points', None)
-    scannet_instance_labels_result = scannet_results.get(
-        'pts_instance_mask', None)
-    scannet_semantic_labels_result = scannet_results.get(
-        'pts_semantic_mask', None)
+    scannet_points_result = scannet_results['points']
+    scannet_instance_labels_result = scannet_results['pts_instance_mask']
+    scannet_semantic_labels_result = scannet_results['pts_semantic_mask']
     scannet_choices = np.array([2, 8, 4, 9, 1])
     assert np.allclose(scannet_points[scannet_choices], scannet_points_result)
     assert np.all(scannet_pts_instance_mask[scannet_choices] ==
@@ -37,7 +35,7 @@ def test_indoor_sample():
                   scannet_semantic_labels_result)
 
     np.random.seed(0)
-    sunrgbd_sample_points = PointSample(5)
+    sunrgbd_sample_points = IndoorPointSample(5)
     sunrgbd_results = dict()
     sunrgbd_point_cloud = np.array(
         [[-1.8135729e-01, 1.4695230e+00, -1.2780589e+00, 7.8938007e-03],
@@ -53,6 +51,6 @@ def test_indoor_sample():
     sunrgbd_results['points'] = sunrgbd_point_cloud
     sunrgbd_results = sunrgbd_sample_points(sunrgbd_results)
     sunrgbd_choices = np.array([2, 8, 4, 9, 1])
-    sunrgbd_points_result = sunrgbd_results.get('points', None)
+    sunrgbd_points_result = sunrgbd_results['points']
     assert np.allclose(sunrgbd_point_cloud[sunrgbd_choices],
                        sunrgbd_points_result)