From 9157ee4436432424232624c4efc37e101dec2011 Mon Sep 17 00:00:00 2001 From: Shaoshuai Shi Date: Thu, 9 Jul 2020 15:15:32 +0800 Subject: [PATCH] add config SET_NAN_VELOCITY_TO_ZEROS=True for NuScenes --- pcdet/datasets/nuscenes/nuscenes_dataset.py | 5 +++++ tools/cfgs/dataset_configs/nuscenes_dataset.yaml | 1 + 2 files changed, 6 insertions(+) diff --git a/pcdet/datasets/nuscenes/nuscenes_dataset.py b/pcdet/datasets/nuscenes/nuscenes_dataset.py index 206ac9710..d365c2463 100644 --- a/pcdet/datasets/nuscenes/nuscenes_dataset.py +++ b/pcdet/datasets/nuscenes/nuscenes_dataset.py @@ -132,6 +132,11 @@ def __getitem__(self, index): data_dict = self.prepare_data(data_dict=input_dict) + if self.dataset_cfg.get('SET_NAN_VELOCITY_TO_ZEROS', False): + gt_boxes = data_dict['gt_boxes'] + gt_boxes[np.isnan(gt_boxes)] = 0 + data_dict['gt_boxes'] = gt_boxes + if not self.dataset_cfg.PRED_VELOCITY and 'gt_boxes' in data_dict: data_dict['gt_boxes'] = data_dict['gt_boxes'][:, [0, 1, 2, 3, 4, 5, 6, -1]] diff --git a/tools/cfgs/dataset_configs/nuscenes_dataset.yaml b/tools/cfgs/dataset_configs/nuscenes_dataset.yaml index 8d9993285..5aa89d802 100644 --- a/tools/cfgs/dataset_configs/nuscenes_dataset.yaml +++ b/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 +SET_NAN_VELOCITY_TO_ZEROS: True FILTER_MIN_POINTS_IN_GT: 1 DATA_SPLIT: {