add config SET_NAN_VELOCITY_TO_ZEROS=True for NuScenes

hcheng1005 · Jul 9, 2020 · 9157ee4 · 9157ee4
1 parent 67b43a3
commit 9157ee4
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git 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
@@ -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: {