change names, mv indoor_eval outside

mmdet3d/core/evaluation/__init__.py

 from .class_names import dataset_aliases, get_classes, kitti_classes
-from .indoor_utils import indoor_eval
+from .indoor_eval import indoor_eval
 from .kitti_utils import kitti_eval, kitti_eval_coco_style
 
 __all__ = [

mmdet3d/core/evaluation/indoor_utils/__init__.py

-from .eval import indoor_eval
-
-__all__ = ['indoor_eval']

mmdet3d/datasets/__init__.py

 from mmdet.datasets.builder import DATASETS
 from .builder import build_dataset
 from .dataset_wrappers import RepeatFactorDataset
-from .indoor_dataset import IndoorDataset
+from .indoor_base_dataset import IndoorBaseDataset
 from .kitti2d_dataset import Kitti2DDataset
 from .kitti_dataset import KittiDataset
 from .loader import DistributedGroupSampler, GroupSampler, build_dataloader
@@ -12,8 +12,8 @@ from .pipelines import (GlobalRotScale, IndoorFlipData, IndoorGlobalRotScale,
                         IndoorPointsColorNormalize, ObjectNoise,
                         ObjectRangeFilter, ObjectSample, PointShuffle,
                         PointsRangeFilter, RandomFlip3D)
-from .scannet_dataset import ScannetDataset
-from .sunrgbd_dataset import SunrgbdDataset
+from .scannet_dataset import ScannetBaseDataset
+from .sunrgbd_dataset import SunrgbdBaseDataset
 
 __all__ = [
     'KittiDataset', 'GroupSampler', 'DistributedGroupSampler',
@@ -23,6 +23,6 @@ __all__ = [
     'ObjectRangeFilter', 'PointsRangeFilter', 'Collect3D',
     'IndoorLoadPointsFromFile', 'IndoorPointsColorNormalize',
     'IndoorPointSample', 'IndoorLoadAnnotations3D', 'IndoorPointsColorJitter',
-    'IndoorGlobalRotScale', 'IndoorFlipData', 'SunrgbdDataset',
-    'ScannetDataset', 'IndoorDataset'
+    'IndoorGlobalRotScale', 'IndoorFlipData', 'SunrgbdBaseDataset',
+    'ScannetBaseDataset', 'IndoorBaseDataset'
 ]

mmdet3d/datasets/indoor_dataset.py → mmdet3d/datasets/indoor_base_dataset.py

@@ -9,28 +9,28 @@ from .pipelines import Compose
 
 
 @DATASETS.register_module()
-class IndoorDataset(torch_data.Dataset):
+class IndoorBaseDataset(torch_data.Dataset):
 
     def __init__(self,
                  root_path,
                  ann_file,
                  pipeline=None,
                  training=False,
-                 class_names=None,
+                 cat_ids=None,
                  test_mode=False,
                  with_label=True):
         super().__init__()
         self.root_path = root_path
-        self.class_names = class_names if class_names else self.CLASSES
+        self.cat_ids = cat_ids if cat_ids else self.CLASSES
         self.test_mode = test_mode
         self.training = training
         self.mode = 'TRAIN' if self.training else 'TEST'
-
+        self.label2cat = {i: cat_id for i, cat_id in enumerate(self.cat_ids)}
         mmcv.check_file_exist(ann_file)
         self.infos = mmcv.load(ann_file)
 
         # dataset config
-        self.num_class = len(self.class_names)
+        self.num_class = len(self.cat_ids)
         self.pcd_limit_range = [0, -40, -3.0, 70.4, 40, 3.0]
         if pipeline is not None:
             self.pipeline = Compose(pipeline)
@@ -134,7 +134,7 @@ class IndoorDataset(torch_data.Dataset):
         assert len(metric) > 0
         gt_annos = [copy.deepcopy(info['annos']) for info in self.infos]
         ap_result_str, ap_dict = indoor_eval(gt_annos, results, metric,
-                                             self.class2type)
+                                             self.label2cat)
         return ap_dict
 
     def __len__(self):

mmdet3d/datasets/scannet_dataset.py

@@ -4,31 +4,12 @@ import mmcv
 import numpy as np
 
 from mmdet.datasets import DATASETS
-from .indoor_dataset import IndoorDataset
+from .indoor_base_dataset import IndoorBaseDataset
 
 
 @DATASETS.register_module()
-class ScannetDataset(IndoorDataset):
-    class2type = {
-        0: 'cabinet',
-        1: 'bed',
-        2: 'chair',
-        3: 'sofa',
-        4: 'table',
-        5: 'door',
-        6: 'window',
-        7: 'bookshelf',
-        8: 'picture',
-        9: 'counter',
-        10: 'desk',
-        11: 'curtain',
-        12: 'refrigerator',
-        13: 'showercurtrain',
-        14: 'toilet',
-        15: 'sink',
-        16: 'bathtub',
-        17: 'garbagebin'
-    }
+class ScannetBaseDataset(IndoorBaseDataset):
+
     CLASSES = ('cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
                'bookshelf', 'picture', 'counter', 'desk', 'curtain',
                'refrigerator', 'showercurtrain', 'toilet', 'sink', 'bathtub',
@@ -39,10 +20,10 @@ class ScannetDataset(IndoorDataset):
                  ann_file,
                  pipeline=None,
                  training=False,
-                 class_names=None,
+                 cat_ids=None,
                  test_mode=False,
                  with_label=True):
-        super().__init__(root_path, ann_file, pipeline, training, class_names,
+        super().__init__(root_path, ann_file, pipeline, training, cat_ids,
                          test_mode, with_label)
 
         self.data_path = osp.join(root_path, 'scannet_train_instance_data')

mmdet3d/datasets/sunrgbd_dataset.py

@@ -4,24 +4,12 @@ import mmcv
 import numpy as np
 
 from mmdet.datasets import DATASETS
-from .indoor_dataset import IndoorDataset
+from .indoor_base_dataset import IndoorBaseDataset
 
 
 @DATASETS.register_module()
-class SunrgbdDataset(IndoorDataset):
-
-    class2type = {
-        0: 'bed',
-        1: 'table',
-        2: 'sofa',
-        3: 'chair',
-        4: 'toilet',
-        5: 'desk',
-        6: 'dresser',
-        7: 'night_stand',
-        8: 'bookshelf',
-        9: 'bathtub'
-    }
+class SunrgbdBaseDataset(IndoorBaseDataset):
+
     CLASSES = ('bed', 'table', 'sofa', 'chair', 'toilet', 'desk', 'dresser',
                'night_stand', 'bookshelf', 'bathtub')
 
@@ -30,10 +18,10 @@ class SunrgbdDataset(IndoorDataset):
                  ann_file,
                  pipeline=None,
                  training=False,
-                 class_names=None,
+                 cat_ids=None,
                  test_mode=False,
                  with_label=True):
-        super().__init__(root_path, ann_file, pipeline, training, class_names,
+        super().__init__(root_path, ann_file, pipeline, training, cat_ids,
                          test_mode, with_label)
         self.data_path = osp.join(root_path, 'sunrgbd_trainval')
 

tests/test_scannet_dataset.py

@@ -2,7 +2,7 @@ import numpy as np
 import pytest
 import torch
 
-from mmdet3d.datasets import ScannetDataset
+from mmdet3d.datasets import ScannetBaseDataset
 
 
 def test_getitem():
@@ -36,7 +36,7 @@ def test_getitem():
             ]),
     ]
 
-    scannet_dataset = ScannetDataset(root_path, ann_file, pipelines, True)
+    scannet_dataset = ScannetBaseDataset(root_path, ann_file, pipelines, True)
     data = scannet_dataset[0]
     points = data['points']._data
     gt_bboxes_3d = data['gt_bboxes_3d']._data
@@ -77,7 +77,7 @@ def test_evaluate():
         pytest.skip()
     root_path = './tests/data/scannet'
     ann_file = './tests/data/scannet/scannet_infos.pkl'
-    scannet_dataset = ScannetDataset(root_path, ann_file)
+    scannet_dataset = ScannetBaseDataset(root_path, ann_file)
     results = []
     pred_boxes = dict()
     pred_boxes['box3d_lidar'] = np.array([[

tests/test_sunrgbd_dataset.py

@@ -2,7 +2,7 @@ import numpy as np
 import pytest
 import torch
 
-from mmdet3d.datasets import SunrgbdDataset
+from mmdet3d.datasets import SunrgbdBaseDataset
 
 
 def test_getitem():
@@ -28,7 +28,7 @@ def test_getitem():
         dict(type='Collect3D', keys=['points', 'gt_bboxes_3d', 'gt_labels']),
     ]
 
-    sunrgbd_dataset = SunrgbdDataset(root_path, ann_file, pipelines, True)
+    sunrgbd_dataset = SunrgbdBaseDataset(root_path, ann_file, pipelines, True)
     data = sunrgbd_dataset[0]
     points = data['points']._data
     gt_bboxes_3d = data['gt_bboxes_3d']._data
@@ -67,7 +67,7 @@ def test_evaluate():
         pytest.skip()
     root_path = './tests/data/sunrgbd'
     ann_file = './tests/data/sunrgbd/sunrgbd_infos.pkl'
-    sunrgbd_dataset = SunrgbdDataset(root_path, ann_file)
+    sunrgbd_dataset = SunrgbdBaseDataset(root_path, ann_file)
     results = []
     pred_boxes = dict()
     pred_boxes['box3d_lidar'] = np.array([[