change color and name

mmdet3d/datasets/pipelines/indoor_loading.py

@@ -12,109 +12,87 @@ class LoadPointsFromFile(object):
     Load sunrgbd and scannet points from file.
 
     Args:
-        name (str): scannet or sunrgbd.
-        use_color (bool): Whether to use color.
         use_height (bool): Whether to use height.
         color_mean (List[float]): Mean color of the point cloud.
+        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_color, use_height, color_mean):
-        self.use_color = use_color
+    def __init__(self, use_height, color_mean, load_dim=6, use_dim=[0, 1, 2]):
         self.use_height = use_height
         self.color_mean = color_mean
+        assert max(use_dim) < load_dim
+        self.load_dim = load_dim
+        self.use_dim = use_dim
 
     def __call__(self, results):
-        data_path = results.get('data_path', None)
+        pts_filename = results.get('pts_filename', None)
         info = results.get('info', None)
         name = 'scannet' if info.get('image', None) is None else 'sunrgbd'
+        assert osp.exists(pts_filename)
         if name == 'scannet':
-            pts_filename = info['point_cloud']['lidar_idx']
-            points = self._get_lidar(pts_filename, data_path)
+            points = np.load(pts_filename)
         else:
-            points = np.load(
-                osp.join(data_path, 'lidar',
-                         '%06d.npz' % info['point_cloud']['lidar_idx']))['pc']
-
-        if not self.use_color:
-            if name == 'scannet':
-                pts_color = points[:, 3:6]
-            points = points[:, 0:3]
-        else:
-            if name == 'scannet':
-                pts_color = points[:, 3:6]
-            points = points[:, 0:6]
-            points[:, 3:] = (points[:, 3:] - np.array(self.color_mean)) / 256.0
+            points = np.load(pts_filename)['pc']
+        points = points.reshape(-1, self.load_dim)
+        points[:, 3:6] = points[:, 3:6] - np.array(self.color_mean) / 256.0
+        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
-        if name == 'scannet':
-            results['pts_color'] = pts_color
         return results
 
     def __repr__(self):
         repr_str = self.__class__.__name__
         repr_str += '(use_height={})'.format(self.use_height)
-        repr_str += '(use_color={}'.format(self.use_color)
         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
 
-    def _get_lidar(self, scan_name, data_path):
-        lidar_file = osp.join(data_path, scan_name + '_vert.npy')
-        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)
+        ins_labelname = results.get('ins_labelname', None)
+        sem_labelname = results.get('sem_labelname', None)
         info = results.get('info', None)
         if info['annos']['gt_num'] != 0:
             gt_bboxes_3d = info['annos']['gt_boxes_upright_depth']
-            gt_classes = info['annos']['class'].reshape(-1, 1)
-            gt_bboxes_3d_mask = np.ones_like(gt_classes)
+            gt_labels = info['annos']['class'].reshape(-1, 1)
+            gt_bboxes_3d_mask = np.ones_like(gt_labels)
         else:
             gt_bboxes_3d = np.zeros((1, 6), dtype=np.float32)
-            gt_classes = np.zeros((1, 1))
+            gt_labels = np.zeros((1, 1))
             gt_bboxes_3d_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
+            assert osp.exists(ins_labelname)
+            assert osp.exists(sem_labelname)
+            pts_instance_mask = np.load(ins_labelname)
+            pts_semantic_mask = np.load(sem_labelname)
+            results['pts_instance_mask'] = pts_instance_mask
+            results['pts_semantic_mask'] = pts_semantic_mask
 
         results['gt_bboxes_3d'] = gt_bboxes_3d
-        results['gt_classes'] = gt_classes
+        results['gt_labels'] = gt_labels
         results['gt_bboxes_3d_mask'] = gt_bboxes_3d_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)
-        return np.load(ins_file)
-
-    def _get_semantic_label(self, scan_name, data_path):
-        sem_file = osp.join(data_path, scan_name + '_sem_label.npy')
-        assert osp.exists(sem_file)
-        return np.load(sem_file)

tests/test_indoor_loading.py

+import os.path as osp
+
 import mmcv
 
 from mmdet3d.datasets.pipelines.indoor_loading import (LoadAnnotations3D,
@@ -6,26 +8,33 @@ from mmdet3d.datasets.pipelines.indoor_loading import (LoadAnnotations3D,
 
 def test_load_points_from_file():
     sunrgbd_info = mmcv.load('./tests/data/sunrgbd/sunrgbd_infos.pkl')
-    sunrgbd_load_points_from_file = LoadPointsFromFile(False, True,
-                                                       [0.5, 0.5, 0.5])
+    sunrgbd_load_points_from_file = LoadPointsFromFile(True, [0.5, 0.5, 0.5],
+                                                       6)
     sunrgbd_results = dict()
-    sunrgbd_results['data_path'] = './tests/data/sunrgbd/sunrgbd_trainval'
-    sunrgbd_results['info'] = sunrgbd_info[0]
+    data_path = './tests/data/sunrgbd/sunrgbd_trainval'
+    sunrgbd_results['data_path'] = data_path
+    sunrgbd_info = sunrgbd_info[0]
+    scan_name = sunrgbd_info['point_cloud']['lidar_idx']
+    sunrgbd_results['info'] = sunrgbd_info
+    sunrgbd_results['pts_filename'] = osp.join(data_path, 'lidar',
+                                               '%06d.npz' % scan_name)
     sunrgbd_results = sunrgbd_load_points_from_file(sunrgbd_results)
     sunrgbd_point_cloud = sunrgbd_results.get('points', 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_load_data = LoadPointsFromFile(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]
+    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['info'] = scannet_info
+    scannet_results['pts_filename'] = osp.join(data_path,
+                                               scan_name + '_vert.npy')
     scannet_results = scannet_load_data(scannet_results)
     scannet_point_cloud = scannet_results.get('points', None)
-    scannet_pcl_color = scannet_results.get('pts_color', None)
     assert scannet_point_cloud.shape == (1000, 4)
-    assert scannet_pcl_color.shape == (1000, 3)
 
 
 def test_load_annotations3D():
@@ -36,26 +45,32 @@ def test_load_annotations3D():
     sunrgbd_results['info'] = sunrgbd_info[0]
     sunrgbd_results = sunrgbd_load_annotations3D(sunrgbd_results)
     sunrgbd_gt_boxes = sunrgbd_results.get('gt_bboxes_3d', None)
-    sunrgbd_gt_classes = sunrgbd_results.get('gt_classes', None)
+    sunrgbd_gt_lbaels = sunrgbd_results.get('gt_labels', None)
     sunrgbd_gt_boxes_mask = sunrgbd_results.get('gt_bboxes_3d_mask', None)
     assert sunrgbd_gt_boxes.shape == (3, 7)
-    assert sunrgbd_gt_classes.shape == (3, 1)
+    assert sunrgbd_gt_lbaels.shape == (3, 1)
     assert sunrgbd_gt_boxes_mask.shape == (3, 1)
 
     scannet_info = mmcv.load('./tests/data/scannet/scannet_infos.pkl')
     scannet_load_annotations3D = LoadAnnotations3D()
     scannet_results = dict()
-    scannet_results[
-        'data_path'] = './tests/data/scannet/scannet_train_instance_data'
-    scannet_results['info'] = scannet_info[0]
+    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['ins_labelname'] = osp.join(data_path,
+                                                scan_name + '_ins_label.npy')
+    scannet_results['sem_labelname'] = osp.join(data_path,
+                                                scan_name + '_sem_label.npy')
+    scannet_results['info'] = scannet_info
     scannet_results = scannet_load_annotations3D(scannet_results)
     scannet_gt_boxes = scannet_results.get('gt_bboxes_3d', None)
-    scannet_gt_classes = scannet_results.get('gt_classes', None)
+    scannet_gt_lbaels = scannet_results.get('gt_labels', None)
     scannet_gt_boxes_mask = scannet_results.get('gt_bboxes_3d_mask', None)
-    scannet_instance_labels = scannet_results.get('instance_labels', None)
-    scannet_semantic_labels = scannet_results.get('semantic_labels', None)
+    scannet_pts_instance_mask = scannet_results.get('pts_instance_mask', None)
+    scannet_pts_semantic_mask = scannet_results.get('pts_semantic_mask', None)
     assert scannet_gt_boxes.shape == (27, 6)
-    assert scannet_gt_classes.shape == (27, 1)
+    assert scannet_gt_lbaels.shape == (27, 1)
     assert scannet_gt_boxes_mask.shape == (27, 1)
-    assert scannet_instance_labels.shape == (1000, )
-    assert scannet_semantic_labels.shape == (1000, )
+    assert scannet_pts_instance_mask.shape == (1000, )
+    assert scannet_pts_semantic_mask.shape == (1000, )