finish indoor augment

mmdet3d/datasets/pipelines/indoor_augment.py

@@ -20,7 +20,61 @@ def _rotz(t):
     return rot_mat
 
 
-def _rotate_aligned_boxes(input_boxes, rot_mat):
+# TODO: merge outdoor indoor transform
+@PIPELINES.register_module()
+class IndoorAugment(object):
+    """Indoor Augment.
+
+    Augment sunrgbd and scannet data.
+
+    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.
+        scale_range (float): Range of scale.
+        (List[float]): Mean color of the point cloud.
+    """
+
+    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]):
+        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
+
+    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, np.transpose(rot_mat))
 
@@ -42,83 +96,78 @@ def _rotate_aligned_boxes(input_boxes, rot_mat):
 
         return np.concatenate([new_centers, new_lengths], axis=1)
 
-
-@PIPELINES.register_module()
-class IndoorFlipData(object):
-    """Indoor Flip Data
-
-    Flip point_cloud and groundtruth boxes.
-
-    Args:
-        name (str): name of the dataset.
-    """
-
-    def __init__(self, name):
-        assert name in ['scannet', 'sunrgbd']
-        self.name = name
-
     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 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 self.name == 'sunrgbd':
+                if name == 'sunrgbd':
                     gt_boxes[:, 6] = np.pi - gt_boxes[:, 6]
-            results['gt_boxes'] = gt_boxes
 
-        if self.name == 'scannet' and np.random.random() > 0.5:
+            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__
-        repr_str += '(dataset_name={})'.format(self.name)
-        return repr_str
-
-
-@PIPELINES.register_module()
-class IndoorRotateData(object):
-    """Indoor Rotate Data
 
-    Rotate point_cloud and groundtruth boxes.
-
-    Args:
-        name (str): name of the dataset.
-    """
-
-    def __init__(self, name):
-        assert name in ['scannet', 'sunrgbd']
-        self.name = name
-        self.rot_range = np.pi / 3
-
-    def __call__(self, results):
-        point_cloud = results.get('point_cloud', None)
-        gt_boxes = results.get('gt_boxes', None)
-        rot_angle = (np.random.random() *
-                     self.rot_range) - self.rot_range / 2  # -30 ~ +30 degree
+        if self.use_rotate:
+            rot_angle = (np.random.random() * self.rot_range * np.pi
+                         ) - np.pi * self.rot_range / 2  # -30 ~ +30 degree
             rot_mat = _rotz(rot_angle)
             point_cloud[:, 0:3] = np.dot(point_cloud[:, 0:3],
                                          np.transpose(rot_mat))
 
-        if self.name == 'scannet':
-            gt_boxes = _rotate_aligned_boxes(gt_boxes, rot_mat)
+            if name == 'scannet':
+                gt_boxes = self._rotate_aligned_boxes(gt_boxes, rot_mat)
             else:
-            gt_boxes[:, 0:3] = np.dot(gt_boxes[:, 0:3], np.transpose(rot_mat))
+                gt_boxes[:, 0:3] = np.dot(gt_boxes[:, 0:3],
+                                          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 *= (1 + 0.4 * np.random.random(3) - 0.2
+                          )  # brightness change for each channel
+            rgb_color += (0.1 * np.random.random(3) - 0.05
+                          )  # color shift for each channel
+            rgb_color += np.expand_dims(
+                (0.05 * np.random.random(point_cloud.shape[0]) - 0.025),
+                -1)  # jittering on each pixel
+            rgb_color = np.clip(rgb_color, 0, 1)
+            # 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
+
+        if self.use_scale:
+            # Augment point cloud scale: 0.85x-1.15x
+            scale_ratio = np.random.random(
+            ) * self.scale_range + 1 - self.scale_range / 2
+            scale_ratio = np.expand_dims(np.tile(scale_ratio, 3), 0)
+            point_cloud[:, 0:3] *= scale_ratio
+            gt_boxes[:, 0:3] *= scale_ratio
+            gt_boxes[:, 3:6] *= scale_ratio
+            if self.use_height:
+                point_cloud[:, -1] *= scale_ratio[0, 0]
+
         results['point_cloud'] = point_cloud
         results['gt_boxes'] = gt_boxes
         return results
 
     def __repr__(self):
         repr_str = self.__class__.__name__
-        repr_str += '(dataset_name={})'.format(self.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)
         return repr_str

tests/test_indoor_augment.py

 import numpy as np
 
-from mmdet3d.datasets.pipelines.indoor_augment import (IndoorFlipData,
-                                                       IndoorRotateData)
+from mmdet3d.datasets.pipelines.indoor_augment import IndoorAugment
 
 
-def test_indoor_flip_data():
-    sunrgbd_flip_data = IndoorFlipData('sunrgbd')
+def test_indoor_augment():
+    sunrgbd_augment = IndoorAugment(0, True, True, False, False, True, True)
     sunrgbd_results = dict()
     sunrgbd_results['point_cloud'] = np.array(
         [[1.02828765e+00, 3.65790772e+00, 1.97294697e-01, 1.61959505e+00],
@@ -18,51 +17,28 @@ def test_indoor_flip_data():
                                                 1.500956, 1.637298, 0.636364,
                                                 -1.58242359
                                             ]])
-    sunrgbd_results = sunrgbd_flip_data(sunrgbd_results)
+
+    sunrgbd_results = sunrgbd_augment(sunrgbd_results)
     sunrgbd_point_cloud = sunrgbd_results.get('point_cloud', None)
     sunrgbd_gt_boxes = sunrgbd_results.get('gt_boxes', None)
-    assert sunrgbd_point_cloud.shape == (2, 4)
-    assert sunrgbd_gt_boxes.shape == (2, 7)
-
-    scannet_flip_data = IndoorFlipData('scannet')
-    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_flip_data(scannet_results)
-    scannet_point_cloud = scannet_results.get('point_cloud', None)
-    scannet_gt_boxes = scannet_results.get('gt_boxes', None)
-    assert scannet_point_cloud.shape == (2, 4)
-    assert scannet_gt_boxes.shape == (2, 6)
-
-
-def test_indoor_rotate_data():
-    sunrgbd_indoor_rotate_data = IndoorRotateData('sunrgbd')
-
-    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
+    expected_sunrgbd_point_cloud = np.array(
+        [[-1.87572197, 3.4384955, 0.2033771, 1.66952557],
+         [0.15494677, 1.15088388, -0.77230157, 0.69384689]])
+    expected_sunrgbd_gt_boxes = np.array([[
+        -0.45341026, 0.99208554, -1.01260705, 0.63438248, 0.59022589,
+        0.89963334, -0.15403838
     ],
                                           [
-                                                -0.449953, 1.395455, -1.027778,
-                                                1.500956, 1.637298, 0.636364,
-                                                -1.58242359
+                                              0.13067981, 1.50574494,
+                                              -1.05946338, 1.54722899,
+                                              1.68777428, 0.65598247,
+                                              4.49867047
                                           ]])
-    sunrgbd_results = sunrgbd_indoor_rotate_data(sunrgbd_results)
-    sunrgbd_point_cloud = sunrgbd_results.get('point_cloud', None)
-    sunrgbd_gt_boxes = sunrgbd_results.get('gt_boxes', None)
-    assert sunrgbd_point_cloud.shape == (2, 4)
-    assert sunrgbd_gt_boxes.shape == (2, 7)
+    assert np.allclose(sunrgbd_point_cloud, expected_sunrgbd_point_cloud)
+    assert np.allclose(sunrgbd_gt_boxes, expected_sunrgbd_gt_boxes)
 
-    scannet_indoor_rotate_data = IndoorRotateData('scannet')
+    scannet_augment = IndoorAugment(
+        0, True, True, True, False, False, False, 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],
@@ -72,8 +48,16 @@ def test_indoor_rotate_data():
     ], [
         -0.03226406, 1.70392646, 0.60348618, 0.65165804, 0.72084366, 0.64667457
     ]])
-    scannet_results = scannet_indoor_rotate_data(scannet_results)
+    scannet_results = scannet_augment(scannet_results)
     scannet_point_cloud = scannet_results.get('point_cloud', None)
     scannet_gt_boxes = scannet_results.get('gt_boxes', None)
-    assert scannet_point_cloud.shape == (2, 4)
-    assert scannet_gt_boxes.shape == (2, 6)
+    expected_scannet_point_cloud = np.array(
+        [[-1.61379665, 0.14011924, 0.5811581, 0.5989725],
+         [-1.39004371, -0.44535946, 0.38729519, 0.40510958]])
+    expected_scannet_gt_boxes = np.array([[
+        -0.55030367, -0.49196554, 0.65688646, 0.66436803, 0.61192608, 0.5163464
+    ], [
+        0.06281816, -1.70307376, 0.60348618, 0.66448129, 0.73241497, 0.64667457
+    ]])
+    assert np.allclose(scannet_point_cloud, expected_scannet_point_cloud)
+    assert np.allclose(scannet_gt_boxes, expected_scannet_gt_boxes)