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 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)}")