| import argparse | |
| from os import path as osp | |
| import sys | |
| from data_converter import e2e_nusc as nuscenes_converter | |
| sys.path.append('.') | |
| def nuscenes_data_prep(root_path, | |
| can_bus_root_path, | |
| info_prefix, | |
| version, | |
| dataset_name, | |
| out_dir, | |
| max_sweeps=10): | |
| """Prepare data related to nuScenes dataset. | |
| Related data consists of '.pkl' files recording basic infos, | |
| 2D annotations and groundtruth database. | |
| Args: | |
| root_path (str): Path of dataset root. | |
| info_prefix (str): The prefix of info filenames. | |
| version (str): Dataset version. | |
| dataset_name (str): The dataset class name. | |
| out_dir (str): Output directory of the groundtruth database info. | |
| max_sweeps (int): Number of input consecutive frames. Default: 10 | |
| """ | |
| nuscenes_converter.create_nuscenes_infos( | |
| root_path, out_dir, can_bus_root_path, info_prefix, version=version, max_sweeps=max_sweeps) | |
| if version == 'v1.0-test': | |
| info_test_path = osp.join( | |
| out_dir, f'{info_prefix}_infos_temporal_test.pkl') | |
| nuscenes_converter.export_2d_annotation( | |
| root_path, info_test_path, version=version) | |
| else: | |
| info_train_path = osp.join( | |
| out_dir, f'{info_prefix}_infos_temporal_train.pkl') | |
| info_val_path = osp.join( | |
| out_dir, f'{info_prefix}_infos_temporal_val.pkl') | |
| nuscenes_converter.export_2d_annotation( | |
| root_path, info_train_path, version=version) | |
| nuscenes_converter.export_2d_annotation( | |
| root_path, info_val_path, version=version) | |
| parser = argparse.ArgumentParser(description='Data converter arg parser') | |
| parser.add_argument('dataset', metavar='kitti', help='name of the dataset') | |
| parser.add_argument( | |
| '--root-path', | |
| type=str, | |
| default='./data/kitti', | |
| help='specify the root path of dataset') | |
| parser.add_argument( | |
| '--canbus', | |
| type=str, | |
| default='./data', | |
| help='specify the root path of nuScenes canbus') | |
| parser.add_argument( | |
| '--version', | |
| type=str, | |
| default='v1.0', | |
| required=False, | |
| help='specify the dataset version, no need for kitti') | |
| parser.add_argument( | |
| '--max-sweeps', | |
| type=int, | |
| default=10, | |
| required=False, | |
| help='specify sweeps of lidar per example') | |
| parser.add_argument( | |
| '--out-dir', | |
| type=str, | |
| default='./data/kitti', | |
| required='False', | |
| help='name of info pkl') | |
| parser.add_argument('--extra-tag', type=str, default='kitti') | |
| parser.add_argument( | |
| '--workers', type=int, default=4, help='number of threads to be used') | |
| args = parser.parse_args() | |
| if __name__ == '__main__': | |
| if args.dataset == 'nuscenes' and args.version != 'v1.0-mini': | |
| train_version = f'{args.version}-trainval' | |
| nuscenes_data_prep( | |
| root_path=args.root_path, | |
| can_bus_root_path=args.canbus, | |
| info_prefix=args.extra_tag, | |
| version=train_version, | |
| dataset_name='NuScenesDataset', | |
| out_dir=args.out_dir, | |
| max_sweeps=args.max_sweeps) | |
| test_version = f'{args.version}-test' | |
| nuscenes_data_prep( | |
| root_path=args.root_path, | |
| can_bus_root_path=args.canbus, | |
| info_prefix=args.extra_tag, | |
| version=test_version, | |
| dataset_name='NuScenesDataset', | |
| out_dir=args.out_dir, | |
| max_sweeps=args.max_sweeps) | |
| elif args.dataset == 'nuscenes' and args.version == 'v1.0-mini': | |
| train_version = f'{args.version}' | |
| nuscenes_data_prep( | |
| root_path=args.root_path, | |
| can_bus_root_path=args.canbus, | |
| info_prefix=args.extra_tag, | |
| version=train_version, | |
| dataset_name='NuScenesDataset', | |
| out_dir=args.out_dir, | |
| max_sweeps=args.max_sweeps) |