Commit 04a73eda authored by Shaoshuai Shi's avatar Shaoshuai Shi
Browse files

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
parent c23c4208
...@@ -38,6 +38,9 @@ class NuScenesDataset(DatasetTemplate): ...@@ -38,6 +38,9 @@ class NuScenesDataset(DatasetTemplate):
""" """
Class-balanced sampling of nuScenes dataset from https://arxiv.org/abs/1908.09492 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} cls_infos = {name: [] for name in self.class_names}
for info in infos: for info in infos:
for name in set(info['gt_names']): for name in set(info['gt_names']):
...@@ -117,9 +120,14 @@ class NuScenesDataset(DatasetTemplate): ...@@ -117,9 +120,14 @@ class NuScenesDataset(DatasetTemplate):
} }
if 'gt_boxes' in info: 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({ input_dict.update({
'gt_names': info['gt_names'], 'gt_names': info['gt_names'] if mask is None else info['gt_names'][mask],
'gt_boxes': info['gt_boxes'] 'gt_boxes': info['gt_boxes'] if mask is None else info['gt_boxes'][mask]
}) })
data_dict = self.prepare_data(data_dict=input_dict) data_dict = self.prepare_data(data_dict=input_dict)
...@@ -326,12 +334,13 @@ if __name__ == '__main__': ...@@ -326,12 +334,13 @@ if __name__ == '__main__':
parser = argparse.ArgumentParser(description='arg parser') parser = argparse.ArgumentParser(description='arg parser')
parser.add_argument('--cfg_file', type=str, default=None, help='specify the config of dataset') 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('--func', type=str, default='create_nuscenes_infos', help='')
parser.add_argument('--version', type=str, default='v1.0-trainval', help='')
args = parser.parse_args() args = parser.parse_args()
if args.func == 'create_nuscenes_infos': if args.func == 'create_nuscenes_infos':
dataset_cfg = EasyDict(yaml.load(open(args.cfg_file))) dataset_cfg = EasyDict(yaml.load(open(args.cfg_file)))
ROOT_DIR = (Path(__file__).resolve().parent / '../../../').resolve() ROOT_DIR = (Path(__file__).resolve().parent / '../../../').resolve()
dataset_cfg.VERSION = args.version
create_nuscenes_info( create_nuscenes_info(
version=dataset_cfg.VERSION, version=dataset_cfg.VERSION,
data_path=ROOT_DIR / 'data' / 'nuscenes', data_path=ROOT_DIR / 'data' / 'nuscenes',
...@@ -342,6 +351,6 @@ if __name__ == '__main__': ...@@ -342,6 +351,6 @@ if __name__ == '__main__':
nuscenes_dataset = NuScenesDataset( nuscenes_dataset = NuScenesDataset(
dataset_cfg=dataset_cfg, class_names=None, dataset_cfg=dataset_cfg, class_names=None,
root_path=ROOT_DIR / 'data' / 'nuscenes', 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) nuscenes_dataset.create_groundtruth_database(max_sweeps=dataset_cfg.MAX_SWEEPS)
...@@ -206,8 +206,7 @@ def get_sample_data(nusc, sample_data_token, selected_anntokens=None): ...@@ -206,8 +206,7 @@ def get_sample_data(nusc, sample_data_token, selected_anntokens=None):
cam_intrinsic = np.array(cs_record['camera_intrinsic']) cam_intrinsic = np.array(cs_record['camera_intrinsic'])
imsize = (sd_record['width'], sd_record['height']) imsize = (sd_record['width'], sd_record['height'])
else: else:
cam_intrinsic = None cam_intrinsic = imsize = None
imsize = None
# Retrieve all sample annotations and map to sensor coordinate system. # Retrieve all sample annotations and map to sensor coordinate system.
if selected_anntokens is not None: if selected_anntokens is not None:
...@@ -218,6 +217,7 @@ def get_sample_data(nusc, sample_data_token, selected_anntokens=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. # Make list of Box objects including coord system transforms.
box_list = [] box_list = []
for box in boxes: for box in boxes:
box.velocity = nusc.box_velocity(box.token)
# Move box to ego vehicle coord system # Move box to ego vehicle coord system
box.translate(-np.array(pose_record['translation'])) box.translate(-np.array(pose_record['translation']))
box.rotate(Quaternion(pose_record['rotation']).inverse) 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 ...@@ -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']] annotations = [nusc.get('sample_annotation', token) for token in sample['anns']]
# the filtering gives 0.5~1 map improvement # the filtering gives 0.5~1 map improvement
mask = np.array([(anno['num_lidar_pts'] + anno['num_radar_pts']) > 0 num_lidar_pts = np.array([anno['num_lidar_pts'] for anno in annotations])
for anno in annotations], dtype=bool).reshape(-1) 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) 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) 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 ...@@ -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_boxes_velocity'] = velocity[mask, :]
info['gt_names'] = np.array([map_name_from_general_to_detection[name] for name in names])[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['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: if sample['scene_token'] in train_scenes:
train_nusc_infos.append(info) train_nusc_infos.append(info)
......
...@@ -4,6 +4,7 @@ DATA_PATH: '../data/nuscenes' ...@@ -4,6 +4,7 @@ DATA_PATH: '../data/nuscenes'
VERSION: 'v1.0-trainval' VERSION: 'v1.0-trainval'
MAX_SWEEPS: 10 MAX_SWEEPS: 10
PRED_VELOCITY: True PRED_VELOCITY: True
FILTER_MIN_POINTS_IN_GT: 1
DATA_SPLIT: { DATA_SPLIT: {
'train': train, 'train': train,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment