Update nuscenes_dataset.py (open-mmlab#1139)

It will the keyerror 'gt_boxes' bug when evaluate on the test split, if we do not include this condition.
xq198109 · Oct 8, 2022 · f221374 · f221374
1 parent 8f765b6
commit f221374
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pcdet/datasets/nuscenes/nuscenes_dataset.py b/pcdet/datasets/nuscenes/nuscenes_dataset.py
@@ -140,7 +140,7 @@ def __getitem__(self, index):
 
         data_dict = self.prepare_data(data_dict=input_dict)
 
-        if self.dataset_cfg.get('SET_NAN_VELOCITY_TO_ZEROS', False):
+        if self.dataset_cfg.get('SET_NAN_VELOCITY_TO_ZEROS', False) and 'gt_boxes' in info:
             gt_boxes = data_dict['gt_boxes']
             gt_boxes[np.isnan(gt_boxes)] = 0
             data_dict['gt_boxes'] = gt_boxes