Change nus related (#75)

* Add bottom2gravity and balance_class option.

* Add valid_flag.

* Add use_valid_flag.

* Change balance_class option.

* Add wrapper, remove gravity and balance_class option.

* Change flag, change load_infos to repeat_indices.

* Remove pdb, change docstring.

* Change name, remove unnecessary args.

* Remove _concat_dataset.

* Change name of ClassSampledDataset.

* Change docstring.

* Add get_cat_ids in nus_dataset, add unittest for wrapper.

* Change to len(self.dataset)

* Update nus data.

mmdet3d/datasets/__init__.py

-from mmdet.datasets.builder import DATASETS, build_dataloader, build_dataset
+from mmdet.datasets.builder import build_dataloader
+from .builder import DATASETS, build_dataset
 from .custom_3d import Custom3DDataset
 from .kitti_dataset import KittiDataset
 from .lyft_dataset import LyftDataset

mmdet3d/datasets/builder.py

+import platform
+from mmcv.utils import build_from_cfg
+
+from mmdet.datasets import DATASETS
+from mmdet.datasets.builder import _concat_dataset
+
+if platform.system() != 'Windows':
+    # https://github.com/pytorch/pytorch/issues/973
+    import resource
+    rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
+    hard_limit = rlimit[1]
+    soft_limit = min(4096, hard_limit)
+    resource.setrlimit(resource.RLIMIT_NOFILE, (soft_limit, hard_limit))
+
+
+def build_dataset(cfg, default_args=None):
+    from mmdet3d.datasets.dataset_wrappers import CBGSDataset
+    from mmdet.datasets.dataset_wrappers import (ClassBalancedDataset,
+                                                 ConcatDataset, RepeatDataset)
+    if isinstance(cfg, (list, tuple)):
+        dataset = ConcatDataset([build_dataset(c, default_args) for c in cfg])
+    elif cfg['type'] == 'ConcatDataset':
+        dataset = ConcatDataset(
+            [build_dataset(c, default_args) for c in cfg['datasets']],
+            cfg.get('separate_eval', True))
+    elif cfg['type'] == 'RepeatDataset':
+        dataset = RepeatDataset(
+            build_dataset(cfg['dataset'], default_args), cfg['times'])
+    elif cfg['type'] == 'ClassBalancedDataset':
+        dataset = ClassBalancedDataset(
+            build_dataset(cfg['dataset'], default_args), cfg['oversample_thr'])
+    elif cfg['type'] == 'CBGSDataset':
+        dataset = CBGSDataset(build_dataset(cfg['dataset'], default_args))
+    elif isinstance(cfg.get('ann_file'), (list, tuple)):
+        dataset = _concat_dataset(cfg, default_args)
+    else:
+        dataset = build_from_cfg(cfg, DATASETS, default_args)
+
+    return dataset

mmdet3d/datasets/dataset_wrappers.py

+import numpy as np
+
+from .builder import DATASETS
+
+
+@DATASETS.register_module()
+class CBGSDataset(object):
+    """A wrapper of class sampled dataset with ann_file path. Implementation of
+    paper `Class-balanced Grouping and Sampling for Point Cloud 3D Object
+    Detection <https://arxiv.org/abs/1908.09492.>`_.
+
+    Balance the number of scenes under different classes.
+
+    Args:
+        dataset (:obj:`CustomDataset`): The dataset to be class sampled.
+    """
+
+    def __init__(self, dataset):
+        self.dataset = dataset
+        self.CLASSES = dataset.CLASSES
+        self.sample_indices = self._get_sample_indices()
+        # self.dataset.data_infos = self.data_infos
+        if hasattr(self.dataset, 'flag'):
+            self.flag = np.array(
+                [self.dataset.flag[ind] for ind in self.sample_indices],
+                dtype=np.uint8)
+
+    def _get_sample_indices(self):
+        """Load annotations from ann_file.
+
+        Args:
+            ann_file (str): Path of the annotation file.
+
+        Returns:
+            list[dict]: List of annotations after class sampling.
+        """
+        class_sample_idxs = {name: [] for name in self.CLASSES}
+        for idx in range(len(self.dataset)):
+            class_sample_idx = self.dataset.get_cat_ids(idx)
+            for key in class_sample_idxs.keys():
+                class_sample_idxs[key] += class_sample_idx[key]
+        duplicated_samples = sum([len(v) for _, v in class_sample_idx.items()])
+        class_distribution = {
+            k: len(v) / duplicated_samples
+            for k, v in class_sample_idx.items()
+        }
+
+        sample_indices = []
+
+        frac = 1.0 / len(self.CLASSES)
+        ratios = [frac / v for v in class_distribution.values()]
+        for cls_inds, ratio in zip(list(class_sample_idx.values()), ratios):
+            sample_indices += np.random.choice(cls_inds,
+                                               int(len(cls_inds) *
+                                                   ratio)).tolist()
+        return sample_indices
+
+    def __getitem__(self, idx):
+        """Get item from infos according to the given index.
+
+        Returns:
+            dict: Data dictionary of the corresponding index.
+        """
+        ori_idx = self.sample_indices[idx]
+        return self.dataset[ori_idx]
+
+    def __len__(self):
+        """Return the length of data infos.
+
+        Returns:
+            int: Length of data infos.
+        """
+        return len(self.sample_indices)

mmdet3d/datasets/nuscenes_dataset.py

@@ -36,8 +36,7 @@ class NuScenesDataset(Custom3DDataset):
         box_type_3d (str, optional): Type of 3D box of this dataset.
             Based on the `box_type_3d`, the dataset will encapsulate the box
             to its original format then converted them to `box_type_3d`.
-            Defaults to 'LiDAR' in this dataset. Available options includes
-
+            Defaults to 'LiDAR' in this dataset. Available options includes.
             - 'LiDAR': Box in LiDAR coordinates.
             - 'Depth': Box in depth coordinates, usually for indoor dataset.
             - 'Camera': Box in camera coordinates.
@@ -47,6 +46,8 @@ class NuScenesDataset(Custom3DDataset):
             Defaults to False.
         eval_version (bool, optional): Configuration version of evaluation.
             Defaults to  'detection_cvpr_2019'.
+        use_valid_flag (bool): Whether to use `use_valid_flag` key in the info
+            file as mask to filter gt_boxes and gt_names. Defaults to False.
     """
     NameMapping = {
         'movable_object.barrier': 'barrier',
@@ -111,8 +112,10 @@ class NuScenesDataset(Custom3DDataset):
                  box_type_3d='LiDAR',
                  filter_empty_gt=True,
                  test_mode=False,
-                 eval_version='detection_cvpr_2019'):
+                 eval_version='detection_cvpr_2019',
+                 use_valid_flag=False):
         self.load_interval = load_interval
+        self.use_valid_flag = use_valid_flag
         super().__init__(
             data_root=data_root,
             ann_file=ann_file,
@@ -127,7 +130,6 @@ class NuScenesDataset(Custom3DDataset):
         self.eval_version = eval_version
         from nuscenes.eval.detection.config import config_factory
         self.eval_detection_configs = config_factory(self.eval_version)
-
         if self.modality is None:
             self.modality = dict(
                 use_camera=False,
@@ -137,6 +139,29 @@ class NuScenesDataset(Custom3DDataset):
                 use_external=False,
             )
 
+    def get_cat_ids(self, idx):
+        """Get category distribution of single scene.
+
+        Args:
+            idx (int): Index of the data_info.
+
+        Returns:
+            dict[list]: for each category, if the current scene
+                contains such boxes, store a list containing idx,
+                otherwise, store empty list.
+        """
+        class_sample_idx = {name: [] for name in self.CLASSES}
+        info = self.data_infos[idx]
+        if self.use_valid_flag:
+            mask = info['valid_flag']
+            gt_names = set(info['gt_names'][mask])
+        else:
+            gt_names = set(info['gt_names'])
+        for name in gt_names:
+            if name in self.CLASSES:
+                class_sample_idx[name].append(idx)
+        return class_sample_idx
+
     def load_annotations(self, ann_file):
         """Load annotations from ann_file.
 
@@ -173,7 +198,6 @@ class NuScenesDataset(Custom3DDataset):
                 - ann_info (dict): Annotation info.
         """
         info = self.data_infos[index]
-
         # standard protocal modified from SECOND.Pytorch
         input_dict = dict(
             sample_idx=info['token'],
@@ -228,6 +252,9 @@ class NuScenesDataset(Custom3DDataset):
         """
         info = self.data_infos[index]
         # filter out bbox containing no points
+        if self.use_valid_flag:
+            mask = info['valid_flag']
+        else:
             mask = info['num_lidar_pts'] > 0
         gt_bboxes_3d = info['gt_boxes'][mask]
         gt_names_3d = info['gt_names'][mask]

tests/test_dataset/test_dataset_wrappers.py

+import numpy as np
+
+from mmdet3d.datasets.builder import build_dataset
+
+
+def test_getitem():
+    np.random.seed(0)
+    point_cloud_range = [-50, -50, -5, 50, 50, 3]
+    file_client_args = dict(backend='disk')
+    class_names = [
+        'car', 'truck', 'trailer', 'bus', 'construction_vehicle', 'bicycle',
+        'motorcycle', 'pedestrian', 'traffic_cone', 'barrier'
+    ]
+    pipeline = [
+        dict(
+            type='LoadPointsFromFile',
+            load_dim=5,
+            use_dim=5,
+            file_client_args=file_client_args),
+        dict(
+            type='LoadPointsFromMultiSweeps',
+            sweeps_num=9,
+            use_dim=[0, 1, 2, 3, 4],
+            file_client_args=file_client_args,
+            pad_empty_sweeps=True,
+            remove_close=True,
+            test_mode=True),
+        dict(type='LoadAnnotations3D', with_bbox_3d=True, with_label_3d=True),
+        # dict(type='ObjectSample', db_sampler=db_sampler),
+        dict(
+            type='GlobalRotScaleTrans',
+            rot_range=[-0.3925, 0.3925],
+            scale_ratio_range=[0.95, 1.05],
+            translation_std=[0, 0, 0]),
+        dict(type='RandomFlip3D', flip_ratio_bev_horizontal=0.5),
+        dict(type='PointsRangeFilter', point_cloud_range=point_cloud_range),
+        dict(type='ObjectRangeFilter', point_cloud_range=point_cloud_range),
+        dict(type='ObjectNameFilter', classes=class_names),
+        dict(type='PointShuffle'),
+        dict(type='DefaultFormatBundle3D', class_names=class_names),
+        dict(
+            type='Collect3D', keys=['points', 'gt_bboxes_3d', 'gt_labels_3d'])
+    ]
+    input_modality = dict(
+        use_lidar=True,
+        use_camera=False,
+        use_radar=False,
+        use_map=False,
+        use_external=False)
+    dataset_cfg = dict(
+        type='CBGSDataset',
+        dataset=dict(
+            type='NuScenesDataset',
+            data_root='tests/data/nuscenes',
+            ann_file='tests/data/nuscenes/nus_info.pkl',
+            pipeline=pipeline,
+            classes=class_names,
+            modality=input_modality,
+            test_mode=False,
+            use_valid_flag=True,
+            # we use box_type_3d='LiDAR' in kitti and nuscenes dataset
+            # and box_type_3d='Depth' in sunrgbd and scannet dataset.
+            box_type_3d='LiDAR'))
+    nus_dataset = build_dataset(dataset_cfg)
+    assert len(nus_dataset) == 10
+    data = nus_dataset[0]
+    assert data['img_metas'].data['flip'] is True
+    assert data['img_metas'].data['pcd_horizontal_flip'] is True
+    assert data['points']._data.shape == (537, 5)
+
+    data = nus_dataset[1]
+    assert data['img_metas'].data['flip'] is False
+    assert data['img_metas'].data['pcd_horizontal_flip'] is False
+    assert data['points']._data.shape == (537, 5)

tools/data_converter/nuscenes_converter.py

@@ -230,6 +230,10 @@ def _fill_trainval_infos(nusc,
                              for b in boxes]).reshape(-1, 1)
             velocity = np.array(
                 [nusc.box_velocity(token)[:2] for token in sample['anns']])
+            valid_flag = np.array(
+                [(anno['num_lidar_pts'] + anno['num_radar_pts']) > 0
+                 for anno in annotations],
+                dtype=bool).reshape(-1)
             # convert velo from global to lidar
             for i in range(len(boxes)):
                 velo = np.array([*velocity[i], 0.0])
@@ -253,6 +257,7 @@ def _fill_trainval_infos(nusc,
                 [a['num_lidar_pts'] for a in annotations])
             info['num_radar_pts'] = np.array(
                 [a['num_radar_pts'] for a in annotations])
+            info['valid_flag'] = valid_flag
 
         if sample['scene_token'] in train_scenes:
             train_nusc_infos.append(info)