File size: 3,925 Bytes
663494c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
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) |