def split_data(data_dir, idx_path, output_dir):
'''
split dataset according to the index file.
inputs:
data_dir(str): the dataset dir
idx_path(str): the idx_path
output_dir(str): the output dir
outputs:
num_data(int): the # of samples should be in the output_dir
'''
# get idx_file
idx_list = get_idx_list(idx_path)
os.mkdir(output_dir)
os.mkdir(os.path.join(output_dir, 'calib'))
os.mkdir(os.path.join(output_dir, 'image_2'))
os.mkdir(os.path.join(output_dir, 'label_2'))
os.mkdir(os.path.join(output_dir, 'velodyne'))
# symlink data
for idx in idx_list:
os.symlink(os.path.join(data_dir, 'calib', idx + '.txt'),
os.path.join(output_dir, 'calib', idx + '.txt'))
os.symlink(os.path.join(data_dir, 'image_2', idx + '.png'),
os.path.join(output_dir, 'image_2', idx + '.png'))
os.symlink(os.path.join(data_dir, 'label_2', idx + '.txt'),
os.path.join(output_dir, 'label_2', idx + '.txt'))
os.symlink(os.path.join(data_dir, 'velodyne', idx + '.bin'),
os.path.join(output_dir, 'velodyne', idx + '.bin'))
return len(idx_list)
def main():
'''
visualize data
'''
parser = argparse.ArgumentParser(description='Visulize Dataset')
parser.add_argument('--data-dir',
type=str,
metavar='INPUT PATH',
help='dataset dir')
parser.add_argument(
'--idx-file',
type=str,
metavar='INDEX FILE PATH',
help='the txt file containing the indeces of the smapled data')
parser.add_argument('--output-dir',
type=str,
metavar='OUTPUT PATH',
help='output dir')
parser.add_argument('--dataset',
type=str,
metavar='DATASET',
help='KITTI' or 'CARLA')
args = parser.parse_args()
global dataset, data_dir, output_dir
data_dir = args.data_dir
idx_path = args.idx_file
output_dir = args.output_dir
dataset = args.dataset.upper()
idx_list = get_idx_list(idx_path)
with Pool(8) as p:
r = list(tqdm(p.imap(vis_fn, idx_list), total=len(idx_list)))
def create_info_file(root_dir:str, idx_path:str, save_path:str):
'''
Create KITTI_infos_xxx.pkl
[<info>,]
info: {
tag: str (e.g. '000000'),
pc_path: str,
reduced_pc_path: str,
img_path: str,
calib: KittiCalib,
label: KittiLable,
}
'''
global g_data_dir, g_infos
# load idx
root_dir = Path(root_dir)
(root_dir/"reduced_velodyne").mkdir(exist_ok=True)
idx_list = get_idx_list(idx_path)
idx_list.sort()
infos = Manager().list()
g_data_dir = root_dir
g_infos = infos
with Pool(8) as p:
r = list(tqdm(p.imap(create_info_file_wk_fn, idx_list), total=len(idx_list)))
infos = list(infos)
save_pickle(infos, save_path)
print(f"Created {save_path}: {len(infos)} samples")
def create_info_file(root_dir: str, idx_path: str, save_path: str,
dataset: str):
'''
Create <dataset>_infos_xxx.pkl
[<info>,]
info: {
tag: str (e.g. '000000'),
pc_paths: {
velo_top: abs_path,
...
}
img_path: str,
calib: CarlaCalib/WaymoCalib/KittiCalib,
label: CarlaLabel/WaymoLabel/KittiLabel,
}
'''
global g_data_dir, g_infos
root_dir = Path(root_dir)
if dataset == "kitti":
(root_dir / "reduced_velodyne").mkdir(exist_ok=True)
idx_list = get_idx_list(idx_path)
idx_list.sort()
infos = Manager().list()
g_data_dir = root_dir
g_infos = infos
with Pool(8) as p:
if dataset == "carla":
r = list(
tqdm(p.imap(create_info_file_carla_wk_fn, idx_list),
total=len(idx_list)))
elif dataset == "waymo":
r = list(
tqdm(p.imap(create_info_file_waymo_wk_fn, idx_list),
total=len(idx_list)))
elif dataset == "kitti":
r = list(
tqdm(p.imap(create_info_file_kitti_wk_fn, idx_list),
total=len(idx_list)))
infos = list(infos)
save_pickle(infos, save_path)
print(f"Created {save_path}: {len(infos)} samples")
def split_data(data_dir, idx_path, output_dir):
'''
split dataset according to the index file.
inputs:
data_dir(str): the dataset dir
idx_path(str): the idx_path
output_dir(str): the output dir
outputs:
num_data(int): the # of samples should be in the output_dir
'''
# get idx_file
idx_list = get_idx_list(idx_path)
os.mkdir(output_dir)
os.mkdir(os.path.join(output_dir, 'calib'))
# os.mkdir(os.path.join(output_dir, 'image_2'))
os.mkdir(os.path.join(output_dir, 'label_imu'))
os.mkdir(os.path.join(output_dir, 'velo_top'))
os.mkdir(os.path.join(output_dir, 'velo_left'))
os.mkdir(os.path.join(output_dir, 'velo_right'))
# symlink data
idx_count = 0
for idx in idx_list:
if not os.path.isfile(os.path.join(data_dir, 'calib', idx + '.txt')):
return idx_count
idx_count = idx_count + 1
os.symlink(os.path.join(data_dir, 'calib', idx + '.txt'),
os.path.join(output_dir, 'calib', idx + '.txt'))
# os.symlink(os.path.join(data_dir, 'image_2', idx+'.png'), os.path.join(output_dir, 'image_2', idx+'.png'))
os.symlink(os.path.join(data_dir, 'label_imu', idx + '.txt'),
os.path.join(output_dir, 'label_imu', idx + '.txt'))
os.symlink(os.path.join(data_dir, 'velo_top', idx + '.npy'),
os.path.join(output_dir, 'velo_top', idx + '.npy'))
os.symlink(os.path.join(data_dir, 'velo_left', idx + '.npy'),
os.path.join(output_dir, 'velo_left', idx + '.npy'))
os.symlink(os.path.join(data_dir, 'velo_right', idx + '.npy'),
os.path.join(output_dir, 'velo_right', idx + '.npy'))
return len(idx_list)
def create_db_file(root_dir:str, idx_path:str, save_dir:str):
'''
Create KITTI_dbinfos_xxx.pkl and save gt_pc into gt_database/
dbinfo: {
"Car": [<car_dbinfo>, ],
"Pedestrian": [<ped_dbinfo>, ],
...
}
car_dbinfo:{
name: str,
gtpc_path: str,
tag: str,
gt_idx: int, # no. of obj
box3d_cam: KittiObj,
num_points_in_gt: int,
calib: KittiCalib
}
'''
global g_data_dir, g_dbinfos
root_dir = Path(root_dir)
save_dir = Path(save_dir)
(save_dir/"gt_database").mkdir(exist_ok=True)
idx_list = get_idx_list(idx_path)
idx_list.sort()
# get classes
cls_list = get_classes(root_dir/"label_2", idx_list)
dbinfo = {itm: [] for itm in cls_list}
g_data_dir = root_dir
g_dbinfos = Manager().list()
with Pool(8) as p:
r = list(tqdm(p.imap(create_db_file_wk_fn, idx_list), total=len(idx_list)))
for info in g_dbinfos:
dbinfo[info["name"]].append(info)
save_pickle(dbinfo, str(save_dir/"KITTI_dbinfos_train.pkl"))
print("KITTI_dbinfos_train.pkl saved.")
for k, v in dbinfo.items():
print(f"{k}: {len(v)}")
def create_db_file(root_dir: str, idx_path: str, save_dir: str, dataset: str):
'''
Create <dataset>_dbinfos_xxx.pkl and save gt_pc into gt_database/
dbinfo: {
"Car": [<car_dbinfo>, ],
"Pedestrian": [<ped_dbinfo>, ],
...
}
car_dbinfo:{
name: str,
gtpc_paths: {
velo_top: abs_path of saved gtpc (IMU Frame),
...
}
tag: str,
gt_idx: int, # no. of obj
box3d_imu: CarlaObj, WaymoObj, KittiObj
num_points_in_gt: int,
calib: CarlaCalib, WaymoCalib, KittiCalib
}
'''
global g_data_dir, g_dbinfos
root_dir = Path(root_dir)
save_dir = Path(save_dir)
(save_dir / "gt_database").mkdir(exist_ok=False)
idx_list = get_idx_list(idx_path)
idx_list.sort()
# get classes
if dataset in ["carla", "waymo"]:
cls_list = get_classes(root_dir / "label_imu",
idx_list,
dataset=dataset)
elif dataset == "kitti":
cls_list = get_classes(root_dir / "label_2", idx_list, dataset=dataset)
dbinfo = {itm: [] for itm in cls_list}
g_data_dir = root_dir
g_dbinfos = Manager().list()
if dataset == "carla":
with Pool(8) as p:
r = list(
tqdm(p.imap(create_db_file_carla_wk_fn, idx_list),
total=len(idx_list)))
for info in g_dbinfos:
dbinfo[info["name"]].append(info)
save_pickle(dbinfo, str(save_dir / "CARLA_dbinfos_train.pkl"))
print("CARLA_dbinfos_train.pkl saved.")
elif dataset == "waymo":
with Pool(8) as p:
r = list(
tqdm(p.imap(create_db_file_waymo_wk_fn, idx_list),
total=len(idx_list)))
for info in g_dbinfos:
dbinfo[info["name"]].append(info)
save_pickle(dbinfo, str(save_dir / "WAYMO_dbinfos_train.pkl"))
print("WAYMO_dbinfos_train.pkl saved.")
elif dataset == "kitti":
with Pool(8) as p:
r = list(
tqdm(p.imap(create_db_file_kitti_wk_fn, idx_list),
total=len(idx_list)))
for info in g_dbinfos:
dbinfo[info["name"]].append(info)
save_pickle(dbinfo, str(save_dir / "KITTI_dbinfos_train.pkl"))
print("KITTI_dbinfos_train.pkl saved.")
else:
raise NotImplementedError
for k, v in dbinfo.items():
print(f"{k}: {len(v)}")