Merge branch 'indoor_loading' into 'master'

Indoor loading

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

mmdet3d/datasets/__init__.py

@@ -5,8 +5,10 @@ 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, ObjectNoise, ObjectRangeFilter,
-                        ObjectSample, PointShuffle, PointsRangeFilter,
+from .pipelines import (GlobalRotScale, IndoorLoadAnnotations3D,
+                        IndoorLoadPointsFromFile, IndoorPointsColorNormalize,
+                        ObjectNoise, ObjectRangeFilter, ObjectSample,
+                        PointSample, PointShuffle, PointsRangeFilter,
                         RandomFlip3D)
 
 __all__ = [
@@ -14,5 +16,7 @@ __all__ = [
     'build_dataloader', 'RepeatFactorDataset', 'DATASETS', 'build_dataset',
     'CocoDataset', 'Kitti2DDataset', 'NuScenesDataset', 'ObjectSample',
     'RandomFlip3D', 'ObjectNoise', 'GlobalRotScale', 'PointShuffle',
-    'ObjectRangeFilter', 'PointsRangeFilter', 'Collect3D'
+    'ObjectRangeFilter', 'PointsRangeFilter', 'Collect3D',
+    'IndoorLoadPointsFromFile', 'IndoorPointsColorNormalize', 'PointSample',
+    'IndoorLoadAnnotations3D'
 ]

mmdet3d/datasets/pipelines/__init__.py

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

mmdet3d/datasets/pipelines/indoor_loading.py

+import mmcv
+import numpy as np
+
+from mmdet.datasets.builder import PIPELINES
+
+
+@PIPELINES.register_module()
+class IndoorPointsColorNormalize(object):
+    """Indoor Points Color Normalize
+
+    Normalize color of the points.
+
+    Args:
+        color_mean (List[float]): Mean color of the point cloud.
+    """
+
+    def __init__(self, color_mean):
+        self.color_mean = color_mean
+
+    def __call__(self, results):
+        points = results['points']
+        assert points.shape[1] >= 6,\
+            f'Expect points have channel >=6, got {points.shape[1]}'
+        points[:, 3:6] = points[:, 3:6] - np.array(self.color_mean) / 256.0
+        results['points'] = points
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += '(color_mean={})'.format(self.color_mean)
+        return repr_str
+
+
+@PIPELINES.register_module()
+class IndoorLoadPointsFromFile(object):
+    """Indoor Load Points From File.
+
+    Load sunrgbd and scannet points from file.
+
+    Args:
+        use_height (bool): Whether to use height.
+        load_dim (int): The dimension of the loaded points.
+            Default: 6.
+        use_dim (List[int]): Which dimensions of the points to be used.
+            Default: [0, 1, 2].
+    """
+
+    def __init__(self, use_height, load_dim=6, use_dim=[0, 1, 2]):
+        self.use_height = use_height
+        assert max(use_dim) < load_dim, \
+            f'Expect all used dimensions < {load_dim}, got {use_dim}'
+
+        self.load_dim = load_dim
+        self.use_dim = use_dim
+
+    def __call__(self, results):
+        pts_filename = results['pts_filename']
+        mmcv.check_file_exist(pts_filename)
+        points = np.load(pts_filename)
+        points = points.reshape(-1, self.load_dim)
+        points = points[:, self.use_dim]
+
+        if self.use_height:
+            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
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += '(use_height={})'.format(self.use_height)
+        repr_str += '(mean_color={})'.format(self.color_mean)
+        repr_str += '(load_dim={})'.format(self.load_dim)
+        repr_str += '(use_dim={})'.format(self.use_dim)
+        return repr_str
+
+
+@PIPELINES.register_module
+class IndoorLoadAnnotations3D(object):
+    """Indoor Load Annotations3D.
+
+    Load instance mask and semantic mask of points.
+    """
+
+    def __init__(self):
+        pass
+
+    def __call__(self, results):
+        pts_instance_mask_path = results['pts_instance_mask_path']
+        pts_semantic_mask_path = results['pts_semantic_mask_path']
+
+        mmcv.check_file_exist(pts_instance_mask_path)
+        mmcv.check_file_exist(pts_semantic_mask_path)
+        pts_instance_mask = np.load(pts_instance_mask_path)
+        pts_semantic_mask = np.load(pts_semantic_mask_path)
+        results['pts_instance_mask'] = pts_instance_mask
+        results['pts_semantic_mask'] = pts_semantic_mask
+
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        return repr_str

tests/test_indoor_loading.py

+import os.path as osp
+
+import mmcv
+import numpy as np
+
+from mmdet3d.datasets.pipelines import (IndoorLoadAnnotations3D,
+                                        IndoorLoadPointsFromFile)
+
+
+def test_indoor_load_points_from_file():
+    sunrgbd_info = mmcv.load('./tests/data/sunrgbd/sunrgbd_infos.pkl')
+    sunrgbd_load_points_from_file = IndoorLoadPointsFromFile(True, 6)
+    sunrgbd_results = dict()
+    data_path = './tests/data/sunrgbd/sunrgbd_trainval'
+    sunrgbd_info = sunrgbd_info[0]
+    scan_name = sunrgbd_info['point_cloud']['lidar_idx']
+    sunrgbd_results['pts_filename'] = osp.join(data_path, 'lidar',
+                                               f'{scan_name:06d}.npy')
+    sunrgbd_results = sunrgbd_load_points_from_file(sunrgbd_results)
+    sunrgbd_point_cloud = sunrgbd_results['points']
+    assert sunrgbd_point_cloud.shape == (100, 4)
+
+    scannet_info = mmcv.load('./tests/data/scannet/scannet_infos.pkl')
+    scannet_load_data = IndoorLoadPointsFromFile(True)
+    scannet_results = dict()
+    data_path = './tests/data/scannet/scannet_train_instance_data'
+    scannet_results['data_path'] = data_path
+    scannet_info = scannet_info[0]
+    scan_name = scannet_info['point_cloud']['lidar_idx']
+
+    scannet_results['pts_filename'] = osp.join(data_path,
+                                               f'{scan_name}_vert.npy')
+    scannet_results = scannet_load_data(scannet_results)
+    scannet_point_cloud = scannet_results['points']
+    assert scannet_point_cloud.shape == (100, 4)
+
+
+def test_load_annotations3D():
+    sunrgbd_info = mmcv.load('./tests/data/sunrgbd/sunrgbd_infos.pkl')[0]
+    if sunrgbd_info['annos']['gt_num'] != 0:
+        sunrgbd_gt_bboxes_3d = sunrgbd_info['annos']['gt_boxes_upright_depth']
+        sunrgbd_gt_labels = sunrgbd_info['annos']['class'].reshape(-1, 1)
+        sunrgbd_gt_bboxes_3d_mask = np.ones_like(sunrgbd_gt_labels)
+    else:
+        sunrgbd_gt_bboxes_3d = np.zeros((1, 6), dtype=np.float32)
+        sunrgbd_gt_labels = np.zeros((1, 1))
+        sunrgbd_gt_bboxes_3d_mask = np.zeros((1, 1))
+    assert sunrgbd_gt_bboxes_3d.shape == (3, 7)
+    assert sunrgbd_gt_labels.shape == (3, 1)
+    assert sunrgbd_gt_bboxes_3d_mask.shape == (3, 1)
+
+    scannet_info = mmcv.load('./tests/data/scannet/scannet_infos.pkl')[0]
+    scannet_load_annotations3D = IndoorLoadAnnotations3D()
+    scannet_results = dict()
+    data_path = './tests/data/scannet/scannet_train_instance_data'
+    if scannet_info['annos']['gt_num'] != 0:
+        scannet_gt_bboxes_3d = scannet_info['annos']['gt_boxes_upright_depth']
+        scannet_gt_labels = scannet_info['annos']['class'].reshape(-1, 1)
+        scannet_gt_bboxes_3d_mask = np.ones_like(scannet_gt_labels)
+    else:
+        scannet_gt_bboxes_3d = np.zeros((1, 6), dtype=np.float32)
+        scannet_gt_labels = np.zeros((1, 1))
+        scannet_gt_bboxes_3d_mask = np.zeros((1, 1))
+    scan_name = scannet_info['point_cloud']['lidar_idx']
+    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')
+    scannet_results['info'] = scannet_info
+    scannet_results['gt_bboxes_3d'] = scannet_gt_bboxes_3d
+    scannet_results['gt_labels'] = scannet_gt_labels
+    scannet_results['gt_bboxes_3d_mask'] = scannet_gt_bboxes_3d_mask
+    scannet_results = scannet_load_annotations3D(scannet_results)
+    scannet_gt_boxes = scannet_results['gt_bboxes_3d']
+    scannet_gt_lbaels = scannet_results['gt_labels']
+    scannet_gt_boxes_mask = scannet_results['gt_bboxes_3d_mask']
+    scannet_pts_instance_mask = scannet_results['pts_instance_mask']
+    scannet_pts_semantic_mask = scannet_results['pts_semantic_mask']
+    assert scannet_gt_boxes.shape == (27, 6)
+    assert scannet_gt_lbaels.shape == (27, 1)
+    assert scannet_gt_boxes_mask.shape == (27, 1)
+    assert scannet_pts_instance_mask.shape == (100, )
+    assert scannet_pts_semantic_mask.shape == (100, )

tests/test_indoor_sample.py

 import numpy as np
 
-from mmdet3d.datasets.pipelines.indoor_sample import PointSample
+from mmdet3d.datasets.pipelines import PointSample
 
 
 def test_indoor_sample():

tools/data_converter/sunrgbd_data_utils.py

@@ -146,9 +146,9 @@ class SUNRGBDData(object):
             pc_upright_depth = self.get_depth(sample_idx)
             pc_upright_depth_subsampled = random_sampling(
                 pc_upright_depth, SAMPLE_NUM)
-            np.savez_compressed(
-                os.path.join(self.root_dir, 'lidar', f'{sample_idx:06d}.npz'),
-                pc=pc_upright_depth_subsampled)
+            np.save(
+                os.path.join(self.root_dir, 'lidar', f'{sample_idx:06d}.npy'),
+                pc_upright_depth_subsampled)
 
             info = dict()
             pc_info = {'num_features': 6, 'lidar_idx': sample_idx}