Spaces:
No application file
No application file
| # --------------------------------------------- | |
| # Copyright (c) OpenMMLab. All rights reserved. | |
| # --------------------------------------------- | |
| # Modified by Zhiqi Li | |
| # --------------------------------------------- | |
| from data_converter.create_gt_database import create_groundtruth_database | |
| from data_converter import nuscenes_converter as nuscenes_converter | |
| from data_converter import lyft_converter as lyft_converter | |
| from data_converter import kitti_converter as kitti | |
| from data_converter import indoor_converter as indoor | |
| import argparse | |
| from os import path as osp | |
| import sys | |
| sys.path.append('.') | |
| def kitti_data_prep(root_path, info_prefix, version, out_dir): | |
| """Prepare data related to Kitti 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. | |
| out_dir (str): Output directory of the groundtruth database info. | |
| """ | |
| kitti.create_kitti_info_file(root_path, info_prefix) | |
| kitti.create_reduced_point_cloud(root_path, info_prefix) | |
| info_train_path = osp.join(root_path, f'{info_prefix}_infos_train.pkl') | |
| info_val_path = osp.join(root_path, f'{info_prefix}_infos_val.pkl') | |
| info_trainval_path = osp.join(root_path, | |
| f'{info_prefix}_infos_trainval.pkl') | |
| info_test_path = osp.join(root_path, f'{info_prefix}_infos_test.pkl') | |
| kitti.export_2d_annotation(root_path, info_train_path) | |
| kitti.export_2d_annotation(root_path, info_val_path) | |
| kitti.export_2d_annotation(root_path, info_trainval_path) | |
| kitti.export_2d_annotation(root_path, info_test_path) | |
| create_groundtruth_database( | |
| 'KittiDataset', | |
| root_path, | |
| info_prefix, | |
| f'{out_dir}/{info_prefix}_infos_train.pkl', | |
| relative_path=False, | |
| mask_anno_path='instances_train.json', | |
| with_mask=(version == 'mask')) | |
| 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) | |
| # create_groundtruth_database(dataset_name, root_path, info_prefix, | |
| # f'{out_dir}/{info_prefix}_infos_train.pkl') | |
| def lyft_data_prep(root_path, info_prefix, version, max_sweeps=10): | |
| """Prepare data related to Lyft dataset. | |
| Related data consists of '.pkl' files recording basic infos. | |
| Although the ground truth database and 2D annotations are not used in | |
| Lyft, it can also be generated like nuScenes. | |
| Args: | |
| root_path (str): Path of dataset root. | |
| info_prefix (str): The prefix of info filenames. | |
| version (str): Dataset version. | |
| max_sweeps (int, optional): Number of input consecutive frames. | |
| Defaults to 10. | |
| """ | |
| lyft_converter.create_lyft_infos( | |
| root_path, info_prefix, version=version, max_sweeps=max_sweeps) | |
| def scannet_data_prep(root_path, info_prefix, out_dir, workers): | |
| """Prepare the info file for scannet dataset. | |
| Args: | |
| root_path (str): Path of dataset root. | |
| info_prefix (str): The prefix of info filenames. | |
| out_dir (str): Output directory of the generated info file. | |
| workers (int): Number of threads to be used. | |
| """ | |
| indoor.create_indoor_info_file( | |
| root_path, info_prefix, out_dir, workers=workers) | |
| def s3dis_data_prep(root_path, info_prefix, out_dir, workers): | |
| """Prepare the info file for s3dis dataset. | |
| Args: | |
| root_path (str): Path of dataset root. | |
| info_prefix (str): The prefix of info filenames. | |
| out_dir (str): Output directory of the generated info file. | |
| workers (int): Number of threads to be used. | |
| """ | |
| indoor.create_indoor_info_file( | |
| root_path, info_prefix, out_dir, workers=workers) | |
| def sunrgbd_data_prep(root_path, info_prefix, out_dir, workers): | |
| """Prepare the info file for sunrgbd dataset. | |
| Args: | |
| root_path (str): Path of dataset root. | |
| info_prefix (str): The prefix of info filenames. | |
| out_dir (str): Output directory of the generated info file. | |
| workers (int): Number of threads to be used. | |
| """ | |
| indoor.create_indoor_info_file( | |
| root_path, info_prefix, out_dir, workers=workers) | |
| def waymo_data_prep(root_path, | |
| info_prefix, | |
| version, | |
| out_dir, | |
| workers, | |
| max_sweeps=5): | |
| """Prepare the info file for waymo dataset. | |
| Args: | |
| root_path (str): Path of dataset root. | |
| info_prefix (str): The prefix of info filenames. | |
| out_dir (str): Output directory of the generated info file. | |
| workers (int): Number of threads to be used. | |
| max_sweeps (int): Number of input consecutive frames. Default: 5 \ | |
| Here we store pose information of these frames for later use. | |
| """ | |
| from tools.data_converter import waymo_converter as waymo | |
| splits = ['training', 'validation', 'testing'] | |
| for i, split in enumerate(splits): | |
| load_dir = osp.join(root_path, 'waymo_format', split) | |
| if split == 'validation': | |
| save_dir = osp.join(out_dir, 'kitti_format', 'training') | |
| else: | |
| save_dir = osp.join(out_dir, 'kitti_format', split) | |
| converter = waymo.Waymo2KITTI( | |
| load_dir, | |
| save_dir, | |
| prefix=str(i), | |
| workers=workers, | |
| test_mode=(split == 'test')) | |
| converter.convert() | |
| # Generate waymo infos | |
| out_dir = osp.join(out_dir, 'kitti_format') | |
| kitti.create_waymo_info_file(out_dir, info_prefix, max_sweeps=max_sweeps) | |
| create_groundtruth_database( | |
| 'WaymoDataset', | |
| out_dir, | |
| info_prefix, | |
| f'{out_dir}/{info_prefix}_infos_train.pkl', | |
| relative_path=False, | |
| with_mask=False) | |
| 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 == 'kitti': | |
| kitti_data_prep( | |
| root_path=args.root_path, | |
| info_prefix=args.extra_tag, | |
| version=args.version, | |
| out_dir=args.out_dir) | |
| elif 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) | |
| elif args.dataset == 'lyft': | |
| train_version = f'{args.version}-train' | |
| lyft_data_prep( | |
| root_path=args.root_path, | |
| info_prefix=args.extra_tag, | |
| version=train_version, | |
| max_sweeps=args.max_sweeps) | |
| test_version = f'{args.version}-test' | |
| lyft_data_prep( | |
| root_path=args.root_path, | |
| info_prefix=args.extra_tag, | |
| version=test_version, | |
| max_sweeps=args.max_sweeps) | |
| elif args.dataset == 'waymo': | |
| waymo_data_prep( | |
| root_path=args.root_path, | |
| info_prefix=args.extra_tag, | |
| version=args.version, | |
| out_dir=args.out_dir, | |
| workers=args.workers, | |
| max_sweeps=args.max_sweeps) | |
| elif args.dataset == 'scannet': | |
| scannet_data_prep( | |
| root_path=args.root_path, | |
| info_prefix=args.extra_tag, | |
| out_dir=args.out_dir, | |
| workers=args.workers) | |
| elif args.dataset == 's3dis': | |
| s3dis_data_prep( | |
| root_path=args.root_path, | |
| info_prefix=args.extra_tag, | |
| out_dir=args.out_dir, | |
| workers=args.workers) | |
| elif args.dataset == 'sunrgbd': | |
| sunrgbd_data_prep( | |
| root_path=args.root_path, | |
| info_prefix=args.extra_tag, | |
| out_dir=args.out_dir, | |
| workers=args.workers) | |