fix load data comments

mmdet3d/datasets/pipelines/indoor_loading.py

@@ -6,10 +6,10 @@ from mmdet.datasets.registry import PIPELINES
 
 
 @PIPELINES.register_module
-class IndoorLoadData(object):
-    """Load Indoor Data
+class LoadPointsFromFile(object):
+    """Load Points From File.
 
-    Load sunrgbd and scannet data.
+    Load sunrgbd and scannet points from file.
 
     Args:
         name (str): scannet or sunrgbd.
@@ -18,9 +18,7 @@ class IndoorLoadData(object):
         mean_color (List[float]): Mean color of the point cloud.
     """
 
-    def __init__(self, name, use_color, use_height, mean_color):
-        assert name in ['scannet', 'sunrgbd']
-        self.name = name
+    def __init__(self, use_color, use_height, mean_color):
         self.use_color = use_color
         self.use_height = use_height
         self.mean_color = mean_color
@@ -28,32 +26,21 @@ class IndoorLoadData(object):
     def __call__(self, results):
         data_path = results.get('data_path', None)
         info = results.get('info', None)
-
-        if info['annos']['gt_num'] != 0:
-            gt_boxes = info['annos']['gt_boxes_upright_depth']
-            gt_classes = info['annos']['class'].reshape(-1, 1)
-            gt_boxes_mask = np.ones_like(gt_classes)
-        else:
-            gt_boxes = np.zeros((1, 6), dtype=np.float32)
-            gt_classes = np.zeros((1, 1))
-            gt_boxes_mask = np.zeros((1, 1))
-
-        if self.name == 'scannet':
+        name = 'scannet' if info.get('image', None) is None else 'sunrgbd'
+        if name == 'scannet':
             scan_name = info['point_cloud']['lidar_idx']
             point_cloud = 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(
                 osp.join(data_path, 'lidar',
                          '%06d.npz' % info['point_cloud']['lidar_idx']))['pc']
 
         if not self.use_color:
-            if self.name == 'scannet':
+            if name == 'scannet':
                 pcl_color = point_cloud[:, 3:6]
-            point_cloud = point_cloud[:, 0:3]  # do not use color for now
+            point_cloud = point_cloud[:, 0:3]
         else:
-            if self.name == 'scannet':
+            if name == 'scannet':
                 pcl_color = point_cloud[:, 3:6]
             point_cloud = point_cloud[:, 0:6]
             point_cloud[:, 3:] = (point_cloud[:, 3:] -
@@ -65,19 +52,12 @@ class IndoorLoadData(object):
             point_cloud = np.concatenate(
                 [point_cloud, np.expand_dims(height, 1)], 1)
         results['point_cloud'] = point_cloud
-
-        if self.name == 'scannet':
+        if name == 'scannet':
             results['pcl_color'] = pcl_color
-            results['instance_labels'] = instance_labels
-            results['semantic_labels'] = semantic_labels
-        results['gt_boxes'] = gt_boxes
-        results['gt_classes'] = gt_classes
-        results['gt_boxes_mask'] = gt_boxes_mask
         return results
 
     def __repr__(self):
         repr_str = self.__class__.__name__
-        repr_str += '(dataset_name={})'.format(self.name)
         repr_str += '(use_height={})'.format(self.use_height)
         repr_str += '(use_color={}'.format(self.use_color)
         repr_str += '(mean_color={})'.format(self.mean_color)
@@ -88,6 +68,49 @@ class IndoorLoadData(object):
         assert osp.exists(lidar_file)
         return np.load(lidar_file)
 
+
+@PIPELINES.register_module
+class LoadAnnotations3D(object):
+    """Load Annotations3D.
+
+    Load sunrgbd and scannet annotations.
+
+    Args:
+        name (str): scannet or sunrgbd.
+    """
+
+    def __init__(self):
+        pass
+
+    def __call__(self, results):
+        data_path = results.get('data_path', None)
+        info = results.get('info', None)
+        if info['annos']['gt_num'] != 0:
+            gt_boxes = info['annos']['gt_boxes_upright_depth']
+            gt_classes = info['annos']['class'].reshape(-1, 1)
+            gt_boxes_mask = np.ones_like(gt_classes)
+        else:
+            gt_boxes = np.zeros((1, 6), dtype=np.float32)
+            gt_classes = np.zeros((1, 1))
+            gt_boxes_mask = np.zeros((1, 1))
+        name = 'scannet' if info.get('image', None) is None else 'sunrgbd'
+
+        if name == 'scannet':
+            scan_name = info['point_cloud']['lidar_idx']
+            instance_labels = self._get_instance_label(scan_name, data_path)
+            semantic_labels = self._get_semantic_label(scan_name, data_path)
+            results['instance_labels'] = instance_labels
+            results['semantic_labels'] = semantic_labels
+
+        results['gt_boxes'] = gt_boxes
+        results['gt_classes'] = gt_classes
+        results['gt_boxes_mask'] = gt_boxes_mask
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        return repr_str
+
     def _get_instance_label(self, scan_name, data_path):
         ins_file = osp.join(data_path, scan_name + '_ins_label.npy')
         assert osp.exists(ins_file)

tests/test_indoor_loading.py

 import mmcv
 
-from mmdet3d.datasets.pipelines.indoor_loading import IndoorLoadData
+from mmdet3d.datasets.pipelines.indoor_loading import (LoadAnnotations3D,
+                                                       LoadPointsFromFile)
 
 
-def test_indoor_load_data():
+def test_load_points_from_file():
     sunrgbd_info = mmcv.load('./tests/data/sunrgbd/sunrgbd_infos.pkl')
-    sunrgbd_load_data = IndoorLoadData('sunrgbd', False, True, [0.5, 0.5, 0.5])
+    sunrgbd_load_points_from_file = LoadPointsFromFile(False, True,
+                                                       [0.5, 0.5, 0.5])
     sunrgbd_results = dict()
     sunrgbd_results['data_path'] = './tests/data/sunrgbd/sunrgbd_trainval'
     sunrgbd_results['info'] = sunrgbd_info[0]
-    sunrgbd_results = sunrgbd_load_data(sunrgbd_results)
+    sunrgbd_results = sunrgbd_load_points_from_file(sunrgbd_results)
     sunrgbd_point_cloud = sunrgbd_results.get('point_cloud', None)
+    assert sunrgbd_point_cloud.shape == (1000, 4)
+
+    scannet_info = mmcv.load('./tests/data/scannet/scannet_infos.pkl')
+    scannet_load_data = LoadPointsFromFile(False, True, [0.5, 0.5, 0.5])
+    scannet_results = dict()
+    scannet_results[
+        'data_path'] = './tests/data/scannet/scannet_train_instance_data'
+    scannet_results['info'] = scannet_info[0]
+    scannet_results = scannet_load_data(scannet_results)
+    scannet_point_cloud = scannet_results.get('point_cloud', None)
+    scannet_pcl_color = scannet_results.get('pcl_color', None)
+    assert scannet_point_cloud.shape == (1000, 4)
+    assert scannet_pcl_color.shape == (1000, 3)
+
+
+def test_load_annotations3D():
+    sunrgbd_info = mmcv.load('./tests/data/sunrgbd/sunrgbd_infos.pkl')
+    sunrgbd_load_annotations3D = LoadAnnotations3D()
+    sunrgbd_results = dict()
+    sunrgbd_results['data_path'] = './tests/data/sunrgbd/sunrgbd_trainval'
+    sunrgbd_results['info'] = sunrgbd_info[0]
+    sunrgbd_results = sunrgbd_load_annotations3D(sunrgbd_results)
     sunrgbd_gt_boxes = sunrgbd_results.get('gt_boxes', None)
     sunrgbd_gt_classes = sunrgbd_results.get('gt_classes', None)
     sunrgbd_gt_boxes_mask = sunrgbd_results.get('gt_boxes_mask', None)
-    assert sunrgbd_point_cloud.shape == (1000, 4)
     assert sunrgbd_gt_boxes.shape == (3, 7)
     assert sunrgbd_gt_classes.shape == (3, 1)
     assert sunrgbd_gt_boxes_mask.shape == (3, 1)
 
     scannet_info = mmcv.load('./tests/data/scannet/scannet_infos.pkl')
-    scannet_load_data = IndoorLoadData('scannet', False, True, [0.5, 0.5, 0.5])
+    scannet_load_annotations3D = LoadAnnotations3D()
     scannet_results = dict()
     scannet_results[
         'data_path'] = './tests/data/scannet/scannet_train_instance_data'
     scannet_results['info'] = scannet_info[0]
-    scannet_results = scannet_load_data(scannet_results)
-    scannet_point_cloud = scannet_results.get('point_cloud', None)
+    scannet_results = scannet_load_annotations3D(scannet_results)
     scannet_gt_boxes = scannet_results.get('gt_boxes', None)
     scannet_gt_classes = scannet_results.get('gt_classes', None)
     scannet_gt_boxes_mask = scannet_results.get('gt_boxes_mask', None)
-    scannet_pcl_color = scannet_results.get('pcl_color', None)
     scannet_instance_labels = scannet_results.get('instance_labels', None)
     scannet_semantic_labels = scannet_results.get('semantic_labels', None)
-    assert scannet_point_cloud.shape == (1000, 4)
     assert scannet_gt_boxes.shape == (27, 6)
     assert scannet_gt_classes.shape == (27, 1)
     assert scannet_gt_boxes_mask.shape == (27, 1)
-    assert scannet_pcl_color.shape == (1000, 3)
     assert scannet_instance_labels.shape == (1000, )
     assert scannet_semantic_labels.shape == (1000, )

tools/data_converter/sunrgbd_data_utils.py

@@ -112,7 +112,7 @@ class SUNRGBDData(object):
         def process_single_scene(sample_idx):
             print('%s sample_idx: %s' % (self.split, sample_idx))
             # convert depth to points
-            SAMPLE_NUM = 1000
+            SAMPLE_NUM = 50000
             pc_upright_depth = self.get_depth(sample_idx)
             # TODO : sample points in loading process and test
             pc_upright_depth_subsampled = random_sampling(