divide augment into flip and global_rot_scale

mmdet3d/datasets/pipelines/indoor_augment.py

@@ -20,18 +20,57 @@ def _rotz(t):
     return rot_mat
 
 
-# TODO: merge outdoor indoor transform
 @PIPELINES.register_module()
-class IndoorAugment(object):
-    """Indoor Augment.
+class IndoorFlipData(object):
+    """Indoor Flip Data
 
-    Augment sunrgbd and scannet data.
+    Flip point_cloud and groundtruth boxes.
+
+    Args:
+        seed (int): Numpy random seed.
+    """
+
+    def __init__(self, seed=None):
+        if seed is not None:
+            np.random.seed(seed)
+
+    def __call__(self, results):
+        point_cloud = results.get('point_cloud', None)
+        gt_boxes = results.get('gt_boxes', None)
+        name = 'scannet' if gt_boxes.shape[1] == 6 else 'sunrgbd'
+        if np.random.random() > 0.5:
+            # Flipping along the YZ plane
+            point_cloud[:, 0] = -1 * point_cloud[:, 0]
+            gt_boxes[:, 0] = -1 * gt_boxes[:, 0]
+            if name == 'sunrgbd':
+                gt_boxes[:, 6] = np.pi - gt_boxes[:, 6]
+            results['gt_boxes'] = gt_boxes
+
+        if name == 'scannet' and np.random.random() > 0.5:
+            # Flipping along the XZ plane
+            point_cloud[:, 1] = -1 * point_cloud[:, 1]
+            gt_boxes[:, 1] = -1 * gt_boxes[:, 1]
+            results['gt_boxes'] = gt_boxes
+        results['point_cloud'] = point_cloud
+
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        return repr_str
+
+
+# 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:
         seed (int): Numpy random seed.
-        use_flip (bool): Whether to use flip.
         use_rotate (bool): Whether to use rotate.
-        use_shuffle (bool): Whether to use shuffle.
         use_color (bool): Whether to use color.
         use_height (bool): Whether to use height.
         rot_range (float): Range of rotation.
@@ -41,27 +80,23 @@ class IndoorAugment(object):
 
     def __init__(self,
                  seed=None,
-                 use_flip=True,
                  use_rotate=True,
-                 use_shuffle=True,
                  use_color=False,
                  use_scale=True,
                  use_height=True,
                  rot_range=1 / 3,
                  scale_range=0.3,
-                 mean_color=[0.5, 0.5, 0.5]):
+                 color_mean=[0.5, 0.5, 0.5]):
         if seed is not None:
             np.random.seed(seed)
 
-        self.use_flip = use_flip
         self.use_rotate = use_rotate
-        self.use_shuffle = use_shuffle
         self.use_color = use_color
         self.use_scale = use_scale
         self.use_height = use_height
         self.rot_range = rot_range
         self.scale_range = scale_range
-        self.mean_color = mean_color
+        self.color_mean = color_mean
 
     def _rotate_aligned_boxes(self, input_boxes, rot_mat):
         """Rotate Aligned Boxes.
@@ -101,19 +136,6 @@ class IndoorAugment(object):
         gt_boxes = results.get('gt_boxes', None)
         name = 'scannet' if gt_boxes.shape[1] == 6 else 'sunrgbd'
 
-        if self.use_flip:
-            if np.random.random() > 0.5:
-                # Flipping along the YZ plane
-                point_cloud[:, 0] = -1 * point_cloud[:, 0]
-                gt_boxes[:, 0] = -1 * gt_boxes[:, 0]
-                if name == 'sunrgbd':
-                    gt_boxes[:, 6] = np.pi - gt_boxes[:, 6]
-
-            if name == 'scannet' and np.random.random() > 0.5:
-                # Flipping along the XZ plane
-                point_cloud[:, 1] = -1 * point_cloud[:, 1]
-                gt_boxes[:, 1] = -1 * gt_boxes[:, 1]
-
         if self.use_rotate:
             rot_angle = (np.random.random() * self.rot_range * np.pi
                          ) - np.pi * self.rot_range / 2  # -30 ~ +30 degree
@@ -128,12 +150,9 @@ class IndoorAugment(object):
                                           np.transpose(rot_mat))
                 gt_boxes[:, 6] -= rot_angle
 
-        if self.use_shuffle:
-            np.random.shuffle(point_cloud)
-
         # Augment RGB color
         if self.use_color:
-            rgb_color = point_cloud[:, 3:6] + self.mean_color
+            rgb_color = point_cloud[:, 3:6] + self.color_mean
             rgb_color *= (1 + 0.4 * np.random.random(3) - 0.2
                           )  # brightness change for each channel
             rgb_color += (0.1 * np.random.random(3) - 0.05
@@ -145,7 +164,7 @@ class IndoorAugment(object):
             # randomly drop out 30% of the points' colors
             rgb_color *= np.expand_dims(
                 np.random.random(point_cloud.shape[0]) > 0.3, -1)
-            point_cloud[:, 3:6] = rgb_color - self.mean_color
+            point_cloud[:, 3:6] = rgb_color - self.color_mean
 
         if self.use_scale:
             # Augment point cloud scale: 0.85x-1.15x
@@ -165,8 +184,6 @@ class IndoorAugment(object):
     def __repr__(self):
         repr_str = self.__class__.__name__
         repr_str += '(use_rotate={})'.format(self.use_rotate)
-        repr_str += '(use_flip={})'.format(self.use_flip)
-        repr_str += '(use_rotate={})'.format(self.use_shuffle)
         repr_str += '(use_color={})'.format(self.use_color)
         repr_str += '(use_scale={})'.format(self.use_scale)
         repr_str += '(use_height={})'.format(self.use_height)

tests/test_indoor_augment.py

 import numpy as np
 
-from mmdet3d.datasets.pipelines.indoor_augment import IndoorAugment
+from mmdet3d.datasets.pipelines.indoor_augment import (IndoorFlipData,
+                                                       IndoorGlobalRotScale)
 
 
-def test_indoor_augment():
-    sunrgbd_augment = IndoorAugment(0, True, True, False, False, True, True)
+def test_indoor_flip_data():
+    sunrgbd_indoor_flip_data = IndoorFlipData(0)
+    sunrgbd_results = dict()
+    sunrgbd_results['point_cloud'] = 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_boxes'] = 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_point_cloud = sunrgbd_results.get('point_cloud', None)
+    sunrgbd_gt_boxes = sunrgbd_results.get('gt_boxes', None)
+
+    expected_sunrgbd_point_cloud = np.array(
+        [[-1.02828765, 3.65790772, 0.1972947, 1.61959505],
+         [0.39597902, 1.05465031, -0.74920434, 0.673096]])
+    expected_sunrgbd_gt_boxes = 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_point_cloud, expected_sunrgbd_point_cloud)
+    assert np.allclose(sunrgbd_gt_boxes, expected_sunrgbd_gt_boxes)
+
+    scannet_indoor_flip_data = IndoorFlipData(0)
+    scannet_results = dict()
+    scannet_results['point_cloud'] = 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_boxes'] = 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_point_cloud = scannet_results.get('point_cloud', None)
+    scannet_gt_boxes = scannet_results.get('gt_boxes', None)
+    expected_scannet_point_cloud = np.array(
+        [[-1.6110241, 0.16903955, 0.5811581, 0.5989725],
+         [-1.3978075, -0.42035791, 0.38729519, 0.40510958]])
+    expected_scannet_gt_boxes = 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_point_cloud, expected_scannet_point_cloud)
+    assert np.allclose(scannet_gt_boxes, expected_scannet_gt_boxes)
+
+
+def test_global_rot_scale():
+    sunrgbd_augment = IndoorGlobalRotScale(0, True, False, True, True)
     sunrgbd_results = dict()
     sunrgbd_results['point_cloud'] = np.array(
         [[1.02828765e+00, 3.65790772e+00, 1.97294697e-01, 1.61959505e+00],
@@ -22,23 +77,23 @@ def test_indoor_augment():
     sunrgbd_point_cloud = sunrgbd_results.get('point_cloud', None)
     sunrgbd_gt_boxes = sunrgbd_results.get('gt_boxes', None)
     expected_sunrgbd_point_cloud = np.array(
-        [[-1.87572197, 3.4384955, 0.2033771, 1.66952557],
-         [0.15494677, 1.15088388, -0.77230157, 0.69384689]])
+        [[0.89427376, 3.94489646, 0.21003141, 1.72415094],
+         [-0.47835783, 1.09972989, -0.79757058, 0.71654893]])
     expected_sunrgbd_gt_boxes = np.array([[
-        -0.45341026, 0.99208554, -1.01260705, 0.63438248, 0.59022589,
-        0.89963334, -0.15403838
+        0.17080999, 1.11345031, -1.04573864, 0.65513891, 0.60953755,
+        0.92906854, 3.01916788
     ],
                                           [
-                                              0.13067981, 1.50574494,
-                                              -1.05946338, 1.54722899,
-                                              1.68777428, 0.65598247,
-                                              4.49867047
+                                              -0.55427876, 1.45912611,
+                                              -1.09412807, 1.59785293,
+                                              1.74299674, 0.67744563,
+                                              -1.63354097
                                           ]])
     assert np.allclose(sunrgbd_point_cloud, expected_sunrgbd_point_cloud)
     assert np.allclose(sunrgbd_gt_boxes, expected_sunrgbd_gt_boxes)
 
-    scannet_augment = IndoorAugment(
-        0, True, True, True, False, False, False, rot_range=1 / 18)
+    scannet_augment = IndoorGlobalRotScale(
+        0, True, False, False, True, rot_range=1 / 18)
     scannet_results = dict()
     scannet_results['point_cloud'] = np.array(
         [[1.6110241e+00, -1.6903955e-01, 5.8115810e-01, 5.9897250e-01],
@@ -52,12 +107,12 @@ def test_indoor_augment():
     scannet_point_cloud = scannet_results.get('point_cloud', None)
     scannet_gt_boxes = scannet_results.get('gt_boxes', None)
     expected_scannet_point_cloud = np.array(
-        [[-1.61379665, 0.14011924, 0.5811581, 0.5989725],
-         [-1.39004371, -0.44535946, 0.38729519, 0.40510958]])
+        [[1.61240576, -0.15530836, 0.5811581, 0.5989725],
+         [1.39417555, 0.43225122, 0.38729519, 0.40510958]])
     expected_scannet_gt_boxes = np.array([[
-        -0.55030367, -0.49196554, 0.65688646, 0.66436803, 0.61192608, 0.5163464
+        0.55491157, 0.48676213, 0.65688646, 0.65879754, 0.60584609, 0.5163464
     ], [
-        0.06281816, -1.70307376, 0.60348618, 0.66448129, 0.73241497, 0.64667457
+        -0.04677942, 1.70358975, 0.60348618, 0.65777559, 0.72636927, 0.64667457
     ]])
     assert np.allclose(scannet_point_cloud, expected_scannet_point_cloud)
     assert np.allclose(scannet_gt_boxes, expected_scannet_gt_boxes)