[Refactor]Support classes balance dataset

mmdet3d/datasets/__init__.py

 # Copyright (c) OpenMMLab. All rights reserved.
 from .builder import DATASETS, PIPELINES, build_dataset
+from .dataset_wrappers import CBGSDataset
 from .det3d_dataset import Det3DDataset
 from .kitti_dataset import KittiDataset
 from .kitti_mono_dataset import KittiMonoDataset
@@ -28,17 +29,17 @@ from .utils import get_loading_pipeline
 from .waymo_dataset import WaymoDataset
 
 __all__ = [
-    'KittiDataset', 'KittiMonoDataset', 'DATASETS', 'build_dataset',
-    'NuScenesDataset', 'NuScenesMonoDataset', 'LyftDataset', 'ObjectSample',
-    'RandomFlip3D', 'ObjectNoise', 'GlobalRotScaleTrans', 'PointShuffle',
-    'ObjectRangeFilter', 'PointsRangeFilter', 'LoadPointsFromFile',
-    'S3DISSegDataset', 'S3DISDataset', 'NormalizePointsColor',
-    'IndoorPatchPointSample', 'IndoorPointSample', 'PointSample',
-    'LoadAnnotations3D', 'GlobalAlignment', 'SUNRGBDDataset', 'ScanNetDataset',
-    'ScanNetSegDataset', 'ScanNetInstanceSegDataset', 'SemanticKITTIDataset',
-    'Det3DDataset', 'Seg3DDataset', 'LoadPointsFromMultiSweeps',
-    'WaymoDataset', 'BackgroundPointsFilter', 'VoxelBasedPointSampler',
-    'get_loading_pipeline', 'RandomDropPointsColor', 'RandomJitterPoints',
-    'ObjectNameFilter', 'AffineResize', 'RandomShiftScale',
-    'LoadPointsFromDict', 'PIPELINES'
+    'KittiDataset', 'KittiMonoDataset', 'DATASETS', 'CBGSDataset',
+    'build_dataset', 'NuScenesDataset', 'NuScenesMonoDataset', 'LyftDataset',
+    'ObjectSample', 'RandomFlip3D', 'ObjectNoise', 'GlobalRotScaleTrans',
+    'PointShuffle', 'ObjectRangeFilter', 'PointsRangeFilter',
+    'LoadPointsFromFile', 'S3DISSegDataset', 'S3DISDataset',
+    'NormalizePointsColor', 'IndoorPatchPointSample', 'IndoorPointSample',
+    'PointSample', 'LoadAnnotations3D', 'GlobalAlignment', 'SUNRGBDDataset',
+    'ScanNetDataset', 'ScanNetSegDataset', 'ScanNetInstanceSegDataset',
+    'SemanticKITTIDataset', 'Det3DDataset', 'Seg3DDataset',
+    'LoadPointsFromMultiSweeps', 'WaymoDataset', 'BackgroundPointsFilter',
+    'VoxelBasedPointSampler', 'get_loading_pipeline', 'RandomDropPointsColor',
+    'RandomJitterPoints', 'ObjectNameFilter', 'AffineResize',
+    'RandomShiftScale', 'LoadPointsFromDict', 'PIPELINES'
 ]

mmdet3d/datasets/dataset_wrappers.py

@@ -17,8 +17,8 @@ class CBGSDataset(object):
     """
 
     def __init__(self, dataset):
-        self.dataset = dataset
-        self.CLASSES = dataset.CLASSES
+        self.dataset = DATASETS.build(dataset)
+        self.CLASSES = self.dataset.metainfo['CLASSES']
         self.cat2id = {name: i for i, name in enumerate(self.CLASSES)}
         self.sample_indices = self._get_sample_indices()
         # self.dataset.data_infos = self.data_infos
@@ -40,7 +40,10 @@ class CBGSDataset(object):
         for idx in range(len(self.dataset)):
             sample_cat_ids = self.dataset.get_cat_ids(idx)
             for cat_id in sample_cat_ids:
-                class_sample_idxs[cat_id].append(idx)
+                if cat_id != -1:
+                    # Filter categories that do not need to care.
+                    # -1 indicate dontcare in MMDet3d.
+                    class_sample_idxs[cat_id].append(idx)
         duplicated_samples = sum(
             [len(v) for _, v in class_sample_idxs.items()])
         class_distribution = {

mmdet3d/datasets/det3d_dataset.py

@@ -294,3 +294,20 @@ class Det3DDataset(BaseDataset):
                     example['data_sample'].gt_instances_3d.labels_3d) == 0:
                 return None
         return example
+
+    def get_cat_ids(self, idx: int) -> List[int]:
+        """Get category ids by index. Dataset wrapped by ClassBalancedDataset
+        must implement this method.
+
+        The ``CBGSDataset`` or ``ClassBalancedDataset``requires a subclass
+        which implements this method.
+
+        Args:
+            idx (int): The index of data.
+
+        Returns:
+            set[int]: All categories in the sample of specified index.
+        """
+        info = self.get_data_info(idx)
+        gt_labels = info['ann_info']['gt_labels_3d'].tolist()
+        return set(gt_labels)