finish indoor rotate, shuffle and related test unit

mmdet3d/datasets/__init__.py

-from mmdet.datasets.registry import DATASETS
-from .builder import build_dataset
-from .dataset_wrappers import RepeatFactorDataset
-from .kitti2d_dataset import Kitti2DDataset
-from .kitti_dataset import KittiDataset
-from .loader import DistributedGroupSampler, GroupSampler, build_dataloader
-from .nuscenes2d_dataset import NuScenes2DDataset
-from .nuscenes_dataset import NuScenesDataset
-
-__all__ = [
-    'KittiDataset', 'GroupSampler', 'DistributedGroupSampler',
-    'build_dataloader', 'RepeatFactorDataset', 'DATASETS', 'build_dataset',
-    'CocoDataset', 'Kitti2DDataset', 'NuScenesDataset', 'NuScenes2DDataset'
-]
+# from mmdet.datasets.registry import DATASETS
+# from .builder import build_dataset
+# from .dataset_wrappers import RepeatFactorDataset
+# from .kitti2d_dataset import Kitti2DDataset
+# from .kitti_dataset import KittiDataset
+# from .loader import DistributedGroupSampler, GroupSampler, build_dataloader
+# from .nuscenes2d_dataset import NuScenes2DDataset
+# from .nuscenes_dataset import NuScenesDataset
+#
+# __all__ = [
+#     'KittiDataset', 'GroupSampler', 'DistributedGroupSampler',
+#     'build_dataloader', 'RepeatFactorDataset', 'DATASETS', 'build_dataset',
+#     'CocoDataset', 'Kitti2DDataset', 'NuScenesDataset', 'NuScenes2DDataset'
+# ]

mmdet3d/datasets/pipelines/indoor_augment.py

@@ -3,11 +3,51 @@ import numpy as np
 from mmdet.datasets.registry import PIPELINES
 
 
+def _rotz(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(input_boxes, rot_mat):
+    centers, lengths = input_boxes[:, 0:3], input_boxes[:, 3:6]
+    new_centers = np.dot(centers, np.transpose(rot_mat))
+
+    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, crnr in enumerate([(-1, -1), (1, -1), (1, 1), (-1, 1)]):
+        crnrs = np.zeros((dx.shape[0], 3))
+        crnrs[:, 0] = crnr[0] * dx
+        crnrs[:, 1] = crnr[1] * dy
+        crnrs = np.dot(crnrs, np.transpose(rot_mat))
+        new_x[:, i] = crnrs[:, 0]
+        new_y[:, i] = crnrs[:, 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)
+
+
 @PIPELINES.register_module()
 class IndoorFlipData(object):
     """Indoor Flip Data
 
-    Flip the points and groundtruth boxes.
+    Flip points and groundtruth boxes.
 
     Args:
         name (str): name of the dataset.
@@ -18,25 +58,82 @@ class IndoorFlipData(object):
         self.name = name
 
     def __call__(self, results):
-        point_cloud = results.get('point_cloud', None)
+        points = results.get('points', None)
         gt_boxes = results.get('gt_boxes', None)
         if np.random.random() > 0.5:
             # Flipping along the YZ plane
-            point_cloud[:, 0] = -1 * point_cloud[:, 0]
+            points[:, 0] = -1 * points[:, 0]
             gt_boxes[:, 0] = -1 * gt_boxes[:, 0]
             if self.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:
             # Flipping along the XZ plane
-            point_cloud[:, 1] = -1 * point_cloud[:, 1]
+            points[:, 1] = -1 * points[:, 1]
             gt_boxes[:, 1] = -1 * gt_boxes[:, 1]
             results['gt_boxes'] = gt_boxes
-        results['point_cloud'] = point_cloud
+        results['points'] = points
+
+        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 points 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):
+        points = results.get('points', None)
+        gt_boxes = results.get('gt_boxes', None)
+        rot_angle = (np.random.random() * np.pi /
+                     3) - np.pi / 6  # -30 ~ +30 degree
+        rot_mat = _rotz(rot_angle)
+        points[:, 0:3] = np.dot(points[:, 0:3], np.transpose(rot_mat))
+        if self.name == 'scannet':
+            gt_boxes = _rotate_aligned_boxes(gt_boxes, rot_mat)
+        else:
+            gt_boxes[:, 0:3] = np.dot(gt_boxes[:, 0:3], np.transpose(rot_mat))
+            gt_boxes[:, 6] -= rot_angle
+        results['points'] = points
+        results['gt_boxes'] = gt_boxes
         return results
 
     def __repr__(self):
         repr_str = self.__class__.__name__
         repr_str += '(dataset_name={})'.format(self.name)
         return repr_str
+
+
+@PIPELINES.register_module()
+class IndoorShuffleData(object):
+    """Indoor Shuffle Data
+
+    Shuffle points.
+    """
+
+    def __init__(self):
+        pass
+
+    def __call__(self, results):
+        points = results.get('points')
+        np.random.shuffle(points)
+        results['points'] = points
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        return repr_str

mmdet3d/datasets/pipelines/indoor_loading.py

@@ -37,28 +37,26 @@ class IndoorLoadData(object):
             gt_classes = np.zeros((1, 1))
             gt_boxes_mask = np.zeros((1, 1))
         if self.name == 'scannet':
-            scan_name = info['point_cloud']['lidar_idx']
-            point_cloud = self._get_lidar(scan_name, data_path)
+            scan_name = info['points']['lidar_idx']
+            points = self._get_lidar(scan_name, data_path)
             instance_labels = self._get_instance_label(scan_name, data_path)
             semantic_labels = self._get_semantic_label(scan_name, data_path)
         else:
-            point_cloud = np.load(
+            points = np.load(
                 osp.join(data_path, 'lidar',
-                         '%06d.npz' % info['point_cloud']['lidar_idx']))['pc']
+                         '%06d.npz' % info['points']['lidar_idx']))['pc']
         if not self.use_color:
-            point_cloud = point_cloud[:, 0:3]  # do not use color for now
-            pcl_color = point_cloud[:, 3:6]
+            points = points[:, 0:3]  # do not use color for now
+            pcl_color = points[:, 3:6]
         else:
-            point_cloud = point_cloud[:, 0:6]
-            pcl_color = point_cloud[:, 3:6]
-            point_cloud[:, 3:] = (point_cloud[:, 3:] -
-                                  np.array(self.mean_color)) / 256.0
+            points = points[:, 0:6]
+            pcl_color = points[:, 3:6]
+            points[:, 3:] = (points[:, 3:] - np.array(self.mean_color)) / 256.0
         if self.use_height:
-            floor_height = np.percentile(point_cloud[:, 2], 0.99)
-            height = point_cloud[:, 2] - floor_height
-            point_cloud = np.concatenate(
-                [point_cloud, np.expand_dims(height, 1)], 1)
-        results['point_cloud'] = point_cloud
+            floor_height = np.percentile(points[:, 2], 0.99)
+            height = points[:, 2] - floor_height
+            points = np.concatenate([points, np.expand_dims(height, 1)], 1)
+        results['points'] = points
         if self.name == 'scannet':
             results['pcl_color'] = pcl_color
             results['instance_labels'] = instance_labels

tests/test_indoor_augment.py

 import numpy as np
 
-from mmdet3d.datasets.pipelines.indoor_augment import IndoorFlipData
+from mmdet3d.datasets.pipelines.indoor_augment import (IndoorFlipData,
+                                                       IndoorRotateData,
+                                                       IndoorShuffleData)
 
 
 def test_indoor_flip_data():
     sunrgbd_flip_data = IndoorFlipData('sunrgbd')
     sunrgbd_results = dict()
-    sunrgbd_results['point_cloud'] = np.array(
+    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_boxes'] = np.array([[
@@ -18,13 +20,13 @@ def test_indoor_flip_data():
                                                 -1.58242359
                                             ]])
     sunrgbd_results = sunrgbd_flip_data(sunrgbd_results)
-    sunrgbd_point_cloud = sunrgbd_results.get('point_cloud', None)
+    sunrgbd_points = sunrgbd_results.get('points', None)
     sunrgbd_gt_boxes = sunrgbd_results.get('gt_boxes', None)
-    assert sunrgbd_point_cloud.shape == (2, 4)
+    assert sunrgbd_points.shape == (2, 4)
     assert sunrgbd_gt_boxes.shape == (2, 7)
     scannet_flip_data = IndoorFlipData('scannet')
     scannet_results = dict()
-    scannet_results['point_cloud'] = np.array(
+    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_boxes'] = np.array([[
@@ -33,7 +35,54 @@ def test_indoor_flip_data():
         -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_points = scannet_results.get('points', None)
     scannet_gt_boxes = scannet_results.get('gt_boxes', None)
-    assert scannet_point_cloud.shape == (2, 4)
+    assert scannet_points.shape == (2, 4)
     assert scannet_gt_boxes.shape == (2, 6)
+
+
+def test_indoor_rotate_data():
+    sunrgbd_indoor_rotate_data = IndoorRotateData('sunrgbd')
+    scannet_indoor_rotate_data = IndoorRotateData('scannet')
+    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_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_rotate_data(sunrgbd_results)
+    sunrgbd_points = sunrgbd_results.get('points', None)
+    sunrgbd_gt_boxes = sunrgbd_results.get('gt_boxes', None)
+    assert sunrgbd_points.shape == (2, 4)
+    assert sunrgbd_gt_boxes.shape == (2, 7)
+    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_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_rotate_data(scannet_results)
+    scannet_points = scannet_results.get('points', None)
+    scannet_gt_boxes = scannet_results.get('gt_boxes', None)
+    assert scannet_points.shape == (2, 4)
+    assert scannet_gt_boxes.shape == (2, 6)
+
+
+def test_indoor_shuffle_data():
+    indoor_shuffle_data = IndoorShuffleData()
+    results = dict()
+    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]])
+    results = indoor_shuffle_data(results)
+    points = results.get('points')
+    assert points.shape == (2, 4)