bugfixed: velocity in infos with official nuscenes-devikit, add...

bugfixed: velocity in infos with official nuscenes-devikit, add FILTER_MIN_POINTS_IN_GT config for nuscenes dataset

pcdet/datasets/nuscenes/nuscenes_dataset.py

@@ -38,6 +38,9 @@ class NuScenesDataset(DatasetTemplate):
         """
         Class-balanced sampling of nuScenes dataset from https://arxiv.org/abs/1908.09492
         """
+        if self.class_names is None:
+            return infos
+
         cls_infos = {name: [] for name in self.class_names}
         for info in infos:
             for name in set(info['gt_names']):
@@ -117,9 +120,14 @@ class NuScenesDataset(DatasetTemplate):
         }
 
         if 'gt_boxes' in info:
+            if self.dataset_cfg.get('FILTER_MIN_POINTS_IN_GT', False):
+                mask = (info['num_lidar_pts'] > self.dataset_cfg.FILTER_MIN_POINTS_IN_GT - 1)
+            else:
+                mask = None
+
             input_dict.update({
-                'gt_names': info['gt_names'],
-                'gt_boxes': info['gt_boxes']
+                'gt_names': info['gt_names'] if mask is None else info['gt_names'][mask],
+                'gt_boxes': info['gt_boxes'] if mask is None else info['gt_boxes'][mask]
             })
 
         data_dict = self.prepare_data(data_dict=input_dict)
@@ -326,12 +334,13 @@ if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='arg parser')
     parser.add_argument('--cfg_file', type=str, default=None, help='specify the config of dataset')
     parser.add_argument('--func', type=str, default='create_nuscenes_infos', help='')
+    parser.add_argument('--version', type=str, default='v1.0-trainval', help='')
     args = parser.parse_args()
 
     if args.func == 'create_nuscenes_infos':
         dataset_cfg = EasyDict(yaml.load(open(args.cfg_file)))
         ROOT_DIR = (Path(__file__).resolve().parent / '../../../').resolve()
-
+        dataset_cfg.VERSION = args.version
         create_nuscenes_info(
             version=dataset_cfg.VERSION,
             data_path=ROOT_DIR / 'data' / 'nuscenes',
@@ -342,6 +351,6 @@ if __name__ == '__main__':
         nuscenes_dataset = NuScenesDataset(
             dataset_cfg=dataset_cfg, class_names=None,
             root_path=ROOT_DIR / 'data' / 'nuscenes',
-            logger=common_utils.create_logger()
+            logger=common_utils.create_logger(), training=True
         )
         nuscenes_dataset.create_groundtruth_database(max_sweeps=dataset_cfg.MAX_SWEEPS)

pcdet/datasets/nuscenes/nuscenes_utils.py

@@ -206,8 +206,7 @@ def get_sample_data(nusc, sample_data_token, selected_anntokens=None):
         cam_intrinsic = np.array(cs_record['camera_intrinsic'])
         imsize = (sd_record['width'], sd_record['height'])
     else:
-        cam_intrinsic = None
-        imsize = None
+        cam_intrinsic = imsize = None
 
     # Retrieve all sample annotations and map to sensor coordinate system.
     if selected_anntokens is not None:
@@ -218,6 +217,7 @@ def get_sample_data(nusc, sample_data_token, selected_anntokens=None):
     # Make list of Box objects including coord system transforms.
     box_list = []
     for box in boxes:
+        box.velocity = nusc.box_velocity(box.token)
         # Move box to ego vehicle coord system
         box.translate(-np.array(pose_record['translation']))
         box.rotate(Quaternion(pose_record['rotation']).inverse)
@@ -350,8 +350,9 @@ def fill_trainval_infos(data_path, nusc, train_scenes, val_scenes, test=False, m
             annotations = [nusc.get('sample_annotation', token) for token in sample['anns']]
 
             # the filtering gives 0.5~1 map improvement
-            mask = np.array([(anno['num_lidar_pts'] + anno['num_radar_pts']) > 0
-                             for anno in annotations], dtype=bool).reshape(-1)
+            num_lidar_pts = np.array([anno['num_lidar_pts'] for anno in annotations])
+            num_radar_pts = np.array([anno['num_radar_pts'] for anno in annotations])
+            mask = (num_lidar_pts + num_radar_pts > 0)
 
             locs = np.array([b.center for b in ref_boxes]).reshape(-1, 3)
             dims = np.array([b.wlh for b in ref_boxes]).reshape(-1, 3)[:, [1, 0, 2]]  # wlh == > dxdydz (lwh)
@@ -367,6 +368,8 @@ def fill_trainval_infos(data_path, nusc, train_scenes, val_scenes, test=False, m
             info['gt_boxes_velocity'] = velocity[mask, :]
             info['gt_names'] = np.array([map_name_from_general_to_detection[name] for name in names])[mask]
             info['gt_boxes_token'] = tokens[mask]
+            info['num_lidar_pts'] = num_lidar_pts[mask]
+            info['num_radar_pts'] = num_radar_pts[mask]
 
         if sample['scene_token'] in train_scenes:
             train_nusc_infos.append(info)

tools/cfgs/dataset_configs/nuscenes_dataset.yaml

@@ -4,6 +4,7 @@ DATA_PATH: '../data/nuscenes'
 VERSION: 'v1.0-trainval'
 MAX_SWEEPS: 10
 PRED_VELOCITY: True
+FILTER_MIN_POINTS_IN_GT: 1
 
 DATA_SPLIT: {
     'train': train,