Refactor pipeline

mmdet3d/datasets/scannet_dataset.py

@@ -3,11 +3,11 @@ import os.path as osp
 import numpy as np
 
 from mmdet.datasets import DATASETS
-from .indoor_base_dataset import IndoorBaseDataset
+from .custom_3d import Custom3DDataset
 
 
 @DATASETS.register_module()
-class ScanNetDataset(IndoorBaseDataset):
+class ScanNetDataset(Custom3DDataset):
 
     CLASSES = ('cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
                'bookshelf', 'picture', 'counter', 'desk', 'curtain',
@@ -15,38 +15,36 @@ class ScanNetDataset(IndoorBaseDataset):
                'garbagebin')
 
     def __init__(self,
-                 root_path,
+                 data_root,
                  ann_file,
                  pipeline=None,
                  classes=None,
-                 test_mode=False,
-                 with_label=True):
-        super().__init__(root_path, ann_file, pipeline, classes, test_mode,
-                         with_label)
+                 test_mode=False):
+        super().__init__(data_root, ann_file, pipeline, classes, test_mode)
 
     def _get_pts_filename(self, sample_idx):
-        pts_filename = osp.join(self.root_path, f'{sample_idx}_vert.npy')
+        pts_filename = osp.join(self.data_root, f'{sample_idx}_vert.npy')
         return pts_filename
 
-    def _get_ann_info(self, index, sample_idx):
+    def get_ann_info(self, index, sample_idx):
         # Use index to get the annos, thus the evalhook could also use this api
         info = self.data_infos[index]
         if info['annos']['gt_num'] != 0:
             gt_bboxes_3d = info['annos']['gt_boxes_upright_depth']  # k, 6
-            gt_labels = info['annos']['class']
-            gt_bboxes_3d_mask = np.ones_like(gt_labels, dtype=np.bool)
+            gt_labels_3d = info['annos']['class']
+            gt_bboxes_3d_mask = np.ones_like(gt_labels_3d, dtype=np.bool)
         else:
             gt_bboxes_3d = np.zeros((1, 6), dtype=np.float32)
-            gt_labels = np.zeros(1, dtype=np.bool)
+            gt_labels_3d = np.zeros(1, dtype=np.bool)
             gt_bboxes_3d_mask = np.zeros(1, dtype=np.bool)
-        pts_instance_mask_path = osp.join(self.root_path,
+        pts_instance_mask_path = osp.join(self.data_root,
                                           f'{sample_idx}_ins_label.npy')
-        pts_semantic_mask_path = osp.join(self.root_path,
+        pts_semantic_mask_path = osp.join(self.data_root,
                                           f'{sample_idx}_sem_label.npy')
 
         anns_results = dict(
             gt_bboxes_3d=gt_bboxes_3d,
-            gt_labels=gt_labels,
+            gt_labels_3d=gt_labels_3d,
             gt_bboxes_3d_mask=gt_bboxes_3d_mask,
             pts_instance_mask_path=pts_instance_mask_path,
             pts_semantic_mask_path=pts_semantic_mask_path)

mmdet3d/datasets/sunrgbd_dataset.py

@@ -3,44 +3,42 @@ import os.path as osp
 import numpy as np
 
 from mmdet.datasets import DATASETS
-from .indoor_base_dataset import IndoorBaseDataset
+from .custom_3d import Custom3DDataset
 
 
 @DATASETS.register_module()
-class SUNRGBDDataset(IndoorBaseDataset):
+class SUNRGBDDataset(Custom3DDataset):
 
     CLASSES = ('bed', 'table', 'sofa', 'chair', 'toilet', 'desk', 'dresser',
                'night_stand', 'bookshelf', 'bathtub')
 
     def __init__(self,
-                 root_path,
+                 data_root,
                  ann_file,
                  pipeline=None,
                  classes=None,
-                 test_mode=False,
-                 with_label=True):
-        super().__init__(root_path, ann_file, pipeline, classes, test_mode,
-                         with_label)
+                 test_mode=False):
+        super().__init__(data_root, ann_file, pipeline, classes, test_mode)
 
     def _get_pts_filename(self, sample_idx):
-        pts_filename = osp.join(self.root_path, 'lidar',
+        pts_filename = osp.join(self.data_root, 'lidar',
                                 f'{sample_idx:06d}.npy')
         return pts_filename
 
-    def _get_ann_info(self, index, sample_idx):
+    def get_ann_info(self, index, sample_idx):
         # Use index to get the annos, thus the evalhook could also use this api
         info = self.data_infos[index]
         if info['annos']['gt_num'] != 0:
             gt_bboxes_3d = info['annos']['gt_boxes_upright_depth']  # k, 6
-            gt_labels = info['annos']['class']
-            gt_bboxes_3d_mask = np.ones_like(gt_labels, dtype=np.bool)
+            gt_labels_3d = info['annos']['class']
+            gt_bboxes_3d_mask = np.ones_like(gt_labels_3d, dtype=np.bool)
         else:
             gt_bboxes_3d = np.zeros((1, 6), dtype=np.float32)
-            gt_labels = np.zeros(1, dtype=np.bool)
+            gt_labels_3d = np.zeros(1, dtype=np.bool)
             gt_bboxes_3d_mask = np.zeros(1, dtype=np.bool)
 
         anns_results = dict(
             gt_bboxes_3d=gt_bboxes_3d,
-            gt_labels=gt_labels,
+            gt_labels_3d=gt_labels_3d,
             gt_bboxes_3d_mask=gt_bboxes_3d_mask)
         return anns_results

mmdet3d/models/detectors/parta2.py

@@ -119,12 +119,7 @@ class PartA2(TwoStageDetector):
         else:
             return self.forward_test(**kwargs)
 
-    def simple_test(self,
-                    points,
-                    img_meta,
-                    gt_bboxes_3d,
-                    proposals=None,
-                    rescale=False):
+    def simple_test(self, points, img_meta, proposals=None, rescale=False):
         feats_dict, voxels_dict = self.extract_feat(points, img_meta)
 
         if self.with_rpn:

mmdet3d/models/roi_heads/bbox_heads/parta2_bbox_head.py

@@ -545,7 +545,6 @@ class PartA2BboxHead(nn.Module):
                         score_thr,
                         nms_thr,
                         use_rotate_nms=True):
-        normalized_scores = torch.sigmoid(box_probs)
         if use_rotate_nms:
             nms_func = nms_gpu
         else:
@@ -562,7 +561,7 @@ class PartA2BboxHead(nn.Module):
         nms_thresh = nms_thr if isinstance(
             nms_thr, list) else [nms_thr for x in range(self.num_classes)]
         for k in range(0, self.num_classes):
-            class_scores_keep = normalized_scores[:, k] >= score_thresh[k]
+            class_scores_keep = box_probs[:, k] >= score_thresh[k]
 
             if class_scores_keep.int().sum() > 0:
                 original_idxs = class_scores_keep.nonzero().view(-1)

tests/test_anchor.py

@@ -17,7 +17,7 @@ def test_aligned_anchor_generator():
     anchor_generator_cfg = dict(
         type='AlignedAnchor3DRangeGenerator',
         ranges=[[-51.2, -51.2, -1.80, 51.2, 51.2, -1.80]],
-        strides=[1, 2, 4],
+        scales=[1, 2, 4],
         sizes=[
             [0.8660, 2.5981, 1.],  # 1.5/sqrt(3)
             [0.5774, 1.7321, 1.],  # 1/sqrt(3)

tests/test_dataset/test_scannet_dataset.py

@@ -15,23 +15,28 @@ def test_getitem():
                    'sink', 'bathtub', 'garbagebin')
     pipelines = [
         dict(
-            type='IndoorLoadPointsFromFile',
-            use_height=True,
+            type='LoadPointsFromFile',
+            shift_height=True,
             load_dim=6,
             use_dim=[0, 1, 2]),
-        dict(type='IndoorLoadAnnotations3D'),
+        dict(
+            type='LoadAnnotations3D',
+            with_bbox_3d=True,
+            with_label_3d=True,
+            with_mask_3d=True,
+            with_seg_3d=True),
         dict(type='IndoorPointSample', num_points=5),
         dict(type='IndoorFlipData', flip_ratio_yz=1.0, flip_ratio_xz=1.0),
         dict(
             type='IndoorGlobalRotScale',
-            use_height=True,
+            shift_height=True,
             rot_range=[-1 / 36, 1 / 36],
             scale_range=None),
         dict(type='DefaultFormatBundle3D', class_names=class_names),
         dict(
             type='Collect3D',
             keys=[
-                'points', 'gt_bboxes_3d', 'gt_labels', 'pts_semantic_mask',
+                'points', 'gt_bboxes_3d', 'gt_labels_3d', 'pts_semantic_mask',
                 'pts_instance_mask'
             ]),
     ]
@@ -40,7 +45,7 @@ def test_getitem():
     data = scannet_dataset[0]
     points = data['points']._data
     gt_bboxes_3d = data['gt_bboxes_3d']._data
-    gt_labels = data['gt_labels']._data
+    gt_labels = data['gt_labels_3d']._data
     pts_semantic_mask = data['pts_semantic_mask']._data
     pts_instance_mask = data['pts_instance_mask']._data
 
@@ -84,6 +89,7 @@ def test_getitem():
     assert scannet_dataset.CLASSES != original_classes
     assert scannet_dataset.CLASSES == ('cabinet', 'bed')
 
+    # Test load classes from file
     import tempfile
     tmp_file = tempfile.NamedTemporaryFile()
     with open(tmp_file.name, 'w') as f:

tests/test_dataset/test_sunrgbd_dataset.py

@@ -13,26 +13,28 @@ def test_getitem():
                    'dresser', 'night_stand', 'bookshelf', 'bathtub')
     pipelines = [
         dict(
-            type='IndoorLoadPointsFromFile',
-            use_height=True,
+            type='LoadPointsFromFile',
+            shift_height=True,
             load_dim=6,
             use_dim=[0, 1, 2]),
+        dict(type='LoadAnnotations3D'),
         dict(type='IndoorFlipData', flip_ratio_yz=1.0),
         dict(
             type='IndoorGlobalRotScale',
-            use_height=True,
+            shift_height=True,
             rot_range=[-1 / 6, 1 / 6],
             scale_range=[0.85, 1.15]),
         dict(type='IndoorPointSample', num_points=5),
         dict(type='DefaultFormatBundle3D', class_names=class_names),
-        dict(type='Collect3D', keys=['points', 'gt_bboxes_3d', 'gt_labels']),
+        dict(
+            type='Collect3D', keys=['points', 'gt_bboxes_3d', 'gt_labels_3d']),
     ]
 
     sunrgbd_dataset = SUNRGBDDataset(root_path, ann_file, pipelines)
     data = sunrgbd_dataset[0]
     points = data['points']._data
     gt_bboxes_3d = data['gt_bboxes_3d']._data
-    gt_labels = data['gt_labels']._data
+    gt_labels_3d = data['gt_labels_3d']._data
 
     expected_points = np.array(
         [[0.6570105, 1.5538014, 0.24514851, 1.0165423],
@@ -59,7 +61,7 @@ def test_getitem():
 
     assert np.allclose(points, expected_points)
     assert np.allclose(gt_bboxes_3d, expected_gt_bboxes_3d)
-    assert np.all(gt_labels.numpy() == expected_gt_labels)
+    assert np.all(gt_labels_3d.numpy() == expected_gt_labels)
     assert original_classes == class_names
 
     SUNRGBD_dataset = SUNRGBDDataset(

tests/test_pipeline/test_indoor_pipeline.py

@@ -15,16 +15,21 @@ def test_scannet_pipeline():
     np.random.seed(0)
     pipelines = [
         dict(
-            type='IndoorLoadPointsFromFile',
-            use_height=True,
+            type='LoadPointsFromFile',
+            shift_height=True,
             load_dim=6,
             use_dim=[0, 1, 2]),
-        dict(type='IndoorLoadAnnotations3D'),
+        dict(
+            type='LoadAnnotations3D',
+            with_bbox_3d=True,
+            with_label_3d=True,
+            with_mask_3d=True,
+            with_seg_3d=True),
         dict(type='IndoorPointSample', num_points=5),
         dict(type='IndoorFlipData', flip_ratio_yz=1.0, flip_ratio_xz=1.0),
         dict(
             type='IndoorGlobalRotScale',
-            use_height=True,
+            shift_height=True,
             rot_range=[-1 / 36, 1 / 36],
             scale_range=None),
         dict(type='DefaultFormatBundle3D', class_names=class_names),
@@ -45,21 +50,21 @@ def test_scannet_pipeline():
     if info['annos']['gt_num'] != 0:
         scannet_gt_bboxes_3d = info['annos']['gt_boxes_upright_depth']
         scannet_gt_labels_3d = info['annos']['class']
-        scannet_gt_bboxes_3d_mask = np.ones_like(
-            scannet_gt_labels_3d, dtype=np.bool)
     else:
         scannet_gt_bboxes_3d = np.zeros((1, 6), dtype=np.float32)
         scannet_gt_labels_3d = np.zeros((1, ))
-        scannet_gt_bboxes_3d_mask = np.zeros((1, ), dtype=np.bool)
     scan_name = info['point_cloud']['lidar_idx']
+    results['ann_info'] = dict()
+    results['ann_info']['pts_instance_mask_path'] = osp.join(
+        data_path, f'{scan_name}_ins_label.npy')
+    results['ann_info']['pts_semantic_mask_path'] = osp.join(
+        data_path, f'{scan_name}_sem_label.npy')
+    results['ann_info']['gt_bboxes_3d'] = scannet_gt_bboxes_3d
+    results['ann_info']['gt_labels_3d'] = scannet_gt_labels_3d
 
-    results['pts_instance_mask_path'] = osp.join(data_path,
-                                                 f'{scan_name}_ins_label.npy')
-    results['pts_semantic_mask_path'] = osp.join(data_path,
-                                                 f'{scan_name}_sem_label.npy')
-    results['gt_bboxes_3d'] = scannet_gt_bboxes_3d
-    results['gt_labels_3d'] = scannet_gt_labels_3d
-    results['gt_bboxes_3d_mask'] = scannet_gt_bboxes_3d_mask
+    results['bbox3d_fields'] = []
+    results['pts_mask_fields'] = []
+    results['pts_seg_fields'] = []
 
     results = pipeline(results)
 
@@ -100,14 +105,15 @@ def test_sunrgbd_pipeline():
     np.random.seed(0)
     pipelines = [
         dict(
-            type='IndoorLoadPointsFromFile',
-            use_height=True,
+            type='LoadPointsFromFile',
+            shift_height=True,
             load_dim=6,
             use_dim=[0, 1, 2]),
+        dict(type='LoadAnnotations3D'),
         dict(type='IndoorFlipData', flip_ratio_yz=1.0),
         dict(
             type='IndoorGlobalRotScale',
-            use_height=True,
+            shift_height=True,
             rot_range=[-1 / 6, 1 / 6],
             scale_range=[0.85, 1.15]),
         dict(type='IndoorPointSample', num_points=5),
@@ -126,15 +132,18 @@ def test_sunrgbd_pipeline():
     if info['annos']['gt_num'] != 0:
         gt_bboxes_3d = info['annos']['gt_boxes_upright_depth']
         gt_labels_3d = info['annos']['class']
-        gt_bboxes_3d_mask = np.ones_like(gt_labels_3d, dtype=np.bool)
     else:
         gt_bboxes_3d = np.zeros((1, 6), dtype=np.float32)
         gt_labels_3d = np.zeros((1, ))
-        gt_bboxes_3d_mask = np.zeros((1, ), dtype=np.bool)
 
-    results['gt_bboxes_3d'] = gt_bboxes_3d
-    results['gt_labels_3d'] = gt_labels_3d
-    results['gt_bboxes_3d_mask'] = gt_bboxes_3d_mask
+    # prepare input of pipeline
+    results['ann_info'] = dict()
+    results['ann_info']['gt_bboxes_3d'] = gt_bboxes_3d
+    results['ann_info']['gt_labels_3d'] = gt_labels_3d
+    results['bbox3d_fields'] = []
+    results['pts_mask_fields'] = []
+    results['pts_seg_fields'] = []
+
     results = pipeline(results)
     points = results['points']._data
     gt_bboxes_3d = results['gt_bboxes_3d']._data

tests/test_pipeline/test_indoor_loading.py → tests/test_pipeline/test_loading.py

@@ -2,14 +2,14 @@ import os.path as osp
 
 import mmcv
 import numpy as np
+import pytest
 
-from mmdet3d.datasets.pipelines import (IndoorLoadAnnotations3D,
-                                        IndoorLoadPointsFromFile)
+from mmdet3d.datasets.pipelines import LoadAnnotations3D, LoadPointsFromFile
 
 
-def test_indoor_load_points_from_file():
+def test_load_points_from_indoor_file():
     sunrgbd_info = mmcv.load('./tests/data/sunrgbd/sunrgbd_infos.pkl')
-    sunrgbd_load_points_from_file = IndoorLoadPointsFromFile(True, 6)
+    sunrgbd_load_points_from_file = LoadPointsFromFile(6, shift_height=True)
     sunrgbd_results = dict()
     data_path = './tests/data/sunrgbd/sunrgbd_trainval'
     sunrgbd_info = sunrgbd_info[0]
@@ -21,7 +21,7 @@ def test_indoor_load_points_from_file():
     assert sunrgbd_point_cloud.shape == (100, 4)
 
     scannet_info = mmcv.load('./tests/data/scannet/scannet_infos.pkl')
-    scannet_load_data = IndoorLoadPointsFromFile(True)
+    scannet_load_data = LoadPointsFromFile(shift_height=True)
     scannet_results = dict()
     data_path = './tests/data/scannet/scannet_train_instance_data'
     scannet_results['data_path'] = data_path
@@ -35,50 +35,68 @@ def test_indoor_load_points_from_file():
     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_3d = sunrgbd_info['annos']['class']
-        sunrgbd_gt_bboxes_3d_mask = np.ones_like(
-            sunrgbd_gt_labels_3d, dtype=np.bool)
-    else:
-        sunrgbd_gt_bboxes_3d = np.zeros((1, 6), dtype=np.float32)
-        sunrgbd_gt_labels_3d = np.zeros((1, ))
-        sunrgbd_gt_bboxes_3d_mask = np.zeros((1, ), dtype=np.bool)
-    assert sunrgbd_gt_bboxes_3d.shape == (3, 7)
-    assert sunrgbd_gt_labels_3d.shape == (3, )
-    assert sunrgbd_gt_bboxes_3d_mask.shape == (3, )
+def test_load_points_from_outdoor_file():
+    data_path = 'tests/data/kitti/a.bin'
+    load_points_from_file = LoadPointsFromFile(4, 4)
+    results = dict()
+    results['pts_filename'] = data_path
+    results = load_points_from_file(results)
+    points = results['points']
+    assert points.shape == (50, 4)
+    assert np.allclose(points.sum(), 2637.479)
+
+    load_points_from_file = LoadPointsFromFile(4, [0, 1, 2, 3])
+    results = dict()
+    results['pts_filename'] = data_path
+    results = load_points_from_file(results)
+    new_points = results['points']
+    assert new_points.shape == (50, 4)
+    assert np.allclose(points.sum(), 2637.479)
+    np.equal(points, new_points)
 
+    with pytest.raises(AssertionError):
+        LoadPointsFromFile(4, 5)
+
+
+def test_load_annotations3D():
+    # Test scannet LoadAnnotations3D
     scannet_info = mmcv.load('./tests/data/scannet/scannet_infos.pkl')[0]
-    scannet_load_annotations3D = IndoorLoadAnnotations3D()
+    scannet_load_annotations3D = LoadAnnotations3D(
+        with_bbox_3d=True,
+        with_label_3d=True,
+        with_mask_3d=True,
+        with_seg_3d=True)
     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_3d = scannet_info['annos']['class']
-        scannet_gt_bboxes_3d_mask = np.ones_like(
-            scannet_gt_labels_3d, dtype=np.bool)
     else:
         scannet_gt_bboxes_3d = np.zeros((1, 6), dtype=np.float32)
         scannet_gt_labels_3d = np.zeros((1, ))
-        scannet_gt_bboxes_3d_mask = np.zeros((1, ), dtype=np.bool)
+
+    # prepare input of loading pipeline
     scan_name = scannet_info['point_cloud']['lidar_idx']
-    scannet_results['pts_instance_mask_path'] = osp.join(
+    scannet_results['ann_info'] = dict()
+    scannet_results['ann_info']['pts_instance_mask_path'] = osp.join(
         data_path, f'{scan_name}_ins_label.npy')
-    scannet_results['pts_semantic_mask_path'] = osp.join(
+    scannet_results['ann_info']['pts_semantic_mask_path'] = osp.join(
         data_path, f'{scan_name}_sem_label.npy')
-    scannet_results['gt_bboxes_3d'] = scannet_gt_bboxes_3d
-    scannet_results['gt_labels_3d'] = scannet_gt_labels_3d
-    scannet_results['gt_bboxes_3d_mask'] = scannet_gt_bboxes_3d_mask
+    scannet_results['ann_info']['gt_bboxes_3d'] = scannet_gt_bboxes_3d
+    scannet_results['ann_info']['gt_labels_3d'] = scannet_gt_labels_3d
+
+    scannet_results['bbox3d_fields'] = []
+    scannet_results['pts_mask_fields'] = []
+    scannet_results['pts_seg_fields'] = []
+
     scannet_results = scannet_load_annotations3D(scannet_results)
     scannet_gt_boxes = scannet_results['gt_bboxes_3d']
     scannet_gt_lbaels = scannet_results['gt_labels_3d']
-    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, )
-    assert scannet_gt_boxes_mask.shape == (27, )
     assert scannet_pts_instance_mask.shape == (100, )
     assert scannet_pts_semantic_mask.shape == (100, )