def _fill_trainval_infos(nusc,
train_scenes,
val_scenes,
test=False,
max_sweeps=10):
train_nusc_infos = []
val_nusc_infos = []
from pyquaternion import Quaternion
for sample in prog_bar(nusc.sample):
lidar_token = sample["data"]["LIDAR_TOP"]
cam_front_token = sample["data"]["CAM_FRONT"]
sd_rec = nusc.get('sample_data', sample['data']["LIDAR_TOP"])
cs_record = nusc.get('calibrated_sensor',
sd_rec['calibrated_sensor_token'])
pose_record = nusc.get('ego_pose', sd_rec['ego_pose_token'])
lidar_path, boxes, _ = nusc.get_sample_data(lidar_token)
cam_path, _, cam_intrinsic = nusc.get_sample_data(cam_front_token)
assert Path(lidar_path).exists(), (
"you must download all trainval data, key-frame only dataset performs far worse than sweeps."
)
info = {
"lidar_path": lidar_path,
"cam_front_path": cam_path,
"token": sample["token"],
"sweeps": [],
"lidar2ego_translation": cs_record['translation'],
"lidar2ego_rotation": cs_record['rotation'],
"ego2global_translation": pose_record['translation'],
"ego2global_rotation": pose_record['rotation'],
"timestamp": sample["timestamp"],
}
l2e_r = info["lidar2ego_rotation"]
l2e_t = info["lidar2ego_translation"]
e2g_r = info["ego2global_rotation"]
e2g_t = info["ego2global_translation"]
l2e_r_mat = Quaternion(l2e_r).rotation_matrix
e2g_r_mat = Quaternion(e2g_r).rotation_matrix
sd_rec = nusc.get('sample_data', sample['data']["LIDAR_TOP"])
sweeps = []
while len(sweeps) < max_sweeps:
if not sd_rec['prev'] == "":
sd_rec = nusc.get('sample_data', sd_rec['prev'])
cs_record = nusc.get('calibrated_sensor',
sd_rec['calibrated_sensor_token'])
pose_record = nusc.get('ego_pose', sd_rec['ego_pose_token'])
lidar_path = nusc.get_sample_data_path(sd_rec['token'])
sweep = {
"lidar_path": lidar_path,
"sample_data_token": sd_rec['token'],
"lidar2ego_translation": cs_record['translation'],
"lidar2ego_rotation": cs_record['rotation'],
"ego2global_translation": pose_record['translation'],
"ego2global_rotation": pose_record['rotation'],
"timestamp": sd_rec["timestamp"]
}
l2e_r_s = sweep["lidar2ego_rotation"]
l2e_t_s = sweep["lidar2ego_translation"]
e2g_r_s = sweep["ego2global_rotation"]
e2g_t_s = sweep["ego2global_translation"]
# sweep->ego->global->ego'->lidar
l2e_r_s_mat = Quaternion(l2e_r_s).rotation_matrix
e2g_r_s_mat = Quaternion(e2g_r_s).rotation_matrix
R = (l2e_r_s_mat.T @ e2g_r_s_mat.T) @ (
np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T)
T = (l2e_t_s @ e2g_r_s_mat.T + e2g_t_s) @ (
np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T)
T -= e2g_t @ (np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(
l2e_r_mat).T) + l2e_t @ np.linalg.inv(l2e_r_mat).T
sweep["sweep2lidar_rotation"] = R.T # points @ R.T + T
sweep["sweep2lidar_translation"] = T
sweeps.append(sweep)
else:
break
info["sweeps"] = sweeps
if not test:
annotations = [
nusc.get('sample_annotation', token)
for token in sample['anns']
]
locs = np.array([b.center for b in boxes]).reshape(-1, 3)
dims = np.array([b.wlh for b in boxes]).reshape(-1, 3)
rots = np.array([b.orientation.yaw_pitch_roll[0]
for b in boxes]).reshape(-1, 1)
velocity = np.array(
[nusc.box_velocity(token)[:2] for token in sample['anns']])
# convert velo from global to lidar
for i in range(len(boxes)):
velo = np.array([*velocity[i], 0.0])
velo = velo @ np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(
l2e_r_mat).T
velocity[i] = velo[:2]
names = [b.name for b in boxes]
for i in range(len(names)):
if names[i] in NuScenesDataset.NameMapping:
names[i] = NuScenesDataset.NameMapping[names[i]]
names = np.array(names)
# we need to convert rot to SECOND format.
# change the rot format will break all checkpoint, so...
gt_boxes = np.concatenate([locs, dims, -rots - np.pi / 2], axis=1)
assert len(gt_boxes) == len(
annotations), f"{len(gt_boxes)}, {len(annotations)}"
info["gt_boxes"] = gt_boxes
info["gt_names"] = names
info["gt_velocity"] = velocity.reshape(-1, 2)
info["num_lidar_pts"] = np.array(
[a["num_lidar_pts"] for a in annotations])
info["num_radar_pts"] = np.array(
[a["num_radar_pts"] for a in annotations])
if sample["scene_token"] in train_scenes:
train_nusc_infos.append(info)
else:
val_nusc_infos.append(info)
return train_nusc_infos, val_nusc_infos
def _fill_trainval_infos(lyft_ds,
train_scenes,
val_scenes,
do_test=False,
max_sweeps=10):
train_nusc_infos = []
val_nusc_infos = []
for sample in prog_bar(lyft_ds.sample):
lidar_token = sample['data']['LIDAR_TOP']
cam_front_token = sample['data']['CAM_FRONT']
sd_rec = lyft_ds.get('sample_data', lidar_token)
cs_record = lyft_ds.get('calibrated_sensor',
sd_rec['calibrated_sensor_token'])
pose_record = lyft_ds.get('ego_pose', sd_rec['ego_pose_token'])
lidar_path, boxes, _ = lyft_ds.get_sample_data(lidar_token)
cam_path, _, cam_intrinsic = lyft_ds.get_sample_data(cam_front_token)
assert Path(lidar_path).exists(), (
'you must download all trainval data, key-frame only dataset performs far worse than sweeps.'
)
info = {
'lidar_path': lidar_path,
'cam_front_path': cam_path,
'token': sample['token'],
'sweeps': [],
'lidar2ego_translation': cs_record['translation'],
'lidar2ego_rotation': cs_record['rotation'],
'ego2global_translation': pose_record['translation'],
'ego2global_rotation': pose_record['rotation'],
'timestamp': sample['timestamp'],
}
l2e_r = info['lidar2ego_rotation']
l2e_t = info['lidar2ego_translation']
e2g_r = info['ego2global_rotation']
e2g_t = info['ego2global_translation']
l2e_r_mat = Quaternion(l2e_r).rotation_matrix
e2g_r_mat = Quaternion(e2g_r).rotation_matrix
sweeps = []
while len(sweeps) < max_sweeps:
if not sd_rec['prev'] == '':
sd_rec = lyft_ds.get('sample_data', sd_rec['prev'])
cs_record = lyft_ds.get('calibrated_sensor',
sd_rec['calibrated_sensor_token'])
pose_record = lyft_ds.get('ego_pose', sd_rec['ego_pose_token'])
lidar_path = lyft_ds.get_sample_data_path(sd_rec['token'])
sweep = {
'lidar_path': lidar_path,
'sample_data_token': sd_rec['token'],
'lidar2ego_translation': cs_record['translation'],
'lidar2ego_rotation': cs_record['rotation'],
'ego2global_translation': pose_record['translation'],
'ego2global_rotation': pose_record['rotation'],
'timestamp': sd_rec['timestamp']
}
l2e_r_s = sweep['lidar2ego_rotation']
l2e_t_s = sweep['lidar2ego_translation']
e2g_r_s = sweep['ego2global_rotation']
e2g_t_s = sweep['ego2global_translation']
# sweep->ego->global->ego'->lidar
l2e_r_s_mat = Quaternion(l2e_r_s).rotation_matrix
e2g_r_s_mat = Quaternion(e2g_r_s).rotation_matrix
R = (l2e_r_s_mat.T @ e2g_r_s_mat.T) @ (
np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T)
T = (l2e_t_s @ e2g_r_s_mat.T + e2g_t_s) @ (
np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T)
T -= e2g_t @ (np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(
l2e_r_mat).T) + l2e_t @ np.linalg.inv(l2e_r_mat).T
sweep['sweep2lidar_rotation'] = R.T # points @ R.T + T
sweep['sweep2lidar_translation'] = T
sweeps.append(sweep)
else:
break
info['sweeps'] = sweeps
if not do_test:
annotations = [
lyft_ds.get('sample_annotation', token)
for token in sample['anns']
]
locs = np.array([b.center for b in boxes]).reshape(-1, 3)
dims = np.array([b.wlh for b in boxes]).reshape(-1, 3)
rots = np.array([b.orientation.yaw_pitch_roll[0]
for b in boxes]).reshape(-1, 1)
velocity = np.array(
[lyft_ds.box_velocity(token)[:2] for token in sample['anns']])
# convert velo from global to lidar
for i in range(len(boxes)):
velo = np.array([*velocity[i], 0.0])
velo = velo @ np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(
l2e_r_mat).T
velocity[i] = velo[:2]
names = np.array([b.name for b in boxes])
# we need to convert rot to SECOND format.
# change the rot format will break all checkpoint, so...
gt_boxes = np.concatenate([locs, dims, -rots - np.pi / 2], axis=1)
assert len(gt_boxes) == len(
annotations), f'{len(gt_boxes)}, {len(annotations)}'
info['gt_boxes'] = gt_boxes
info['gt_names'] = names
info['gt_velocity'] = velocity.reshape(-1, 2)
info['num_lidar_pts'] = np.array(
[a['num_lidar_pts'] for a in annotations])
info['num_radar_pts'] = np.array(
[a['num_radar_pts'] for a in annotations])
if sample['scene_token'] in train_scenes:
train_nusc_infos.append(info)
else:
val_nusc_infos.append(info)
return train_nusc_infos, val_nusc_infos
def create_groundtruth_database(dataset_class_name,
data_path,
info_path=None,
used_classes=None,
database_save_path=None,
db_info_save_path=None,
relative_path=True,
add_rgb=False,
lidar_only=False,
bev_only=False,
coors_range=None):
dataset = get_dataset_class(dataset_class_name)(
info_path=info_path,
root_path=data_path,
)
root_path = Path(data_path)
if database_save_path is None:
database_save_path = root_path / 'gt_database'
else:
database_save_path = Path(database_save_path)
if db_info_save_path is None:
db_info_save_path = root_path / "kitti_dbinfos_train.pkl"
database_save_path.mkdir(parents=True, exist_ok=True)
all_db_infos = {}
group_counter = 0
for j in prog_bar(list(range(len(dataset)))):
image_idx = j
sensor_data = dataset.get_sensor_data(j)
if "image_idx" in sensor_data["metadata"]:
image_idx = sensor_data["metadata"]["image_idx"]
points = sensor_data["lidar"]["points"]
annos = sensor_data["lidar"]["annotations"]
gt_boxes = annos["boxes"]
names = annos["names"]
group_dict = {}
group_ids = np.full([gt_boxes.shape[0]], -1, dtype=np.int64)
if "group_ids" in annos:
group_ids = annos["group_ids"]
else:
group_ids = np.arange(gt_boxes.shape[0], dtype=np.int64)
difficulty = np.zeros(gt_boxes.shape[0], dtype=np.int32)
if "difficulty" in annos:
difficulty = annos["difficulty"]
num_obj = gt_boxes.shape[0]
point_indices = box_np_ops.points_in_rbbox(points, gt_boxes)
for i in range(num_obj):
filename = f"{image_idx}_{names[i]}_{i}.bin"
filepath = database_save_path / filename
gt_points = points[point_indices[:, i]]
gt_points[:, :3] -= gt_boxes[i, :3]
with open(filepath, 'w') as f:
gt_points.tofile(f)
if (used_classes is None) or names[i] in used_classes:
if relative_path:
db_path = str(database_save_path.stem + "/" + filename)
else:
db_path = str(filepath)
db_info = {
"name": names[i],
"path": db_path,
"image_idx": image_idx,
"gt_idx": i,
"box3d_lidar": gt_boxes[i],
"num_points_in_gt": gt_points.shape[0],
"difficulty": difficulty[i],
# "group_id": -1,
# "bbox": bboxes[i],
}
local_group_id = group_ids[i]
# if local_group_id >= 0:
if local_group_id not in group_dict:
group_dict[local_group_id] = group_counter
group_counter += 1
db_info["group_id"] = group_dict[local_group_id]
if "score" in annos:
db_info["score"] = annos["score"][i]
if names[i] in all_db_infos:
all_db_infos[names[i]].append(db_info)
else:
all_db_infos[names[i]] = [db_info]
for k, v in all_db_infos.items():
print(f"load {len(v)} {k} database infos")
with open(db_info_save_path, 'wb') as f:
pickle.dump(all_db_infos, f)
def create_groundtruth_database(data_path,
info_path=None,
used_classes=None,
database_save_path=None,
db_info_save_path=None,
relative_path=True,
lidar_only=True,
bev_only=False,
coors_range=None):
root_path = pathlib.Path(data_path)
if info_path is None:
info_path = root_path / 'kitti_infos_train.pkl'
if database_save_path is None:
database_save_path = root_path / 'gt_database'
else:
database_save_path = pathlib.Path(database_save_path)
if db_info_save_path is None:
db_info_save_path = root_path / "kitti_dbinfos_train.pkl"
database_save_path.mkdir(parents=True, exist_ok=True)
with open(info_path, 'rb') as f:
kitti_infos = pickle.load(f)
all_db_infos = {}
if used_classes is None:
used_classes = list(kitti.get_classes())
used_classes.pop(used_classes.index('DontCare'))
for name in used_classes:
all_db_infos[name] = []
group_counter = 0
for info in prog_bar(kitti_infos):
velodyne_path = info['velodyne_path']
if relative_path:
# velodyne_path = str(root_path / velodyne_path) + "_reduced"
velodyne_path = str(root_path / velodyne_path)
num_features = 4
if 'pointcloud_num_features' in info:
num_features = info['pointcloud_num_features']
points = np.fromfile(velodyne_path, dtype=np.float32,
count=-1).reshape([-1, num_features])
image_idx = info["image_idx"]
rect = info['calib/R0_rect']
P2 = info['calib/P2']
Trv2c = info['calib/Tr_velo_to_cam']
if not lidar_only:
points = box_np_ops.remove_outside_points(points, rect, Trv2c, P2,
info["img_shape"])
#print(points)
annos = info["annos"]
names = annos["name"]
bboxes = annos["bbox"]
difficulty = annos["difficulty"]
#print(difficulty)
gt_idxes = annos["index"]
num_obj = np.sum(annos["index"] >= 0)
rbbox_cam = kitti.anno_to_rbboxes(annos)[:num_obj]
#print(rbbox_cam)
rbbox_lidar = box_np_ops.box_camera_to_lidar(rbbox_cam, rect, Trv2c)
#print(rbbox_lidar)
if bev_only: # set z and h to limits
assert coors_range is not None
rbbox_lidar[:, 2] = coors_range[2]
rbbox_lidar[:, 5] = coors_range[5] - coors_range[2]
#print(rbbox_lidar)
group_dict = {}
group_ids = np.full([bboxes.shape[0]], -1, dtype=np.int64)
if "group_ids" in annos:
group_ids = annos["group_ids"]
else:
group_ids = np.arange(bboxes.shape[0], dtype=np.int64)
#print(group_ids)
point_indices = box_np_ops.points_in_rbbox(points, rbbox_lidar)
#print(point_indices)
#print(num_obj)
for i in range(num_obj):
filename = f"{image_idx}_{names[i]}_{gt_idxes[i]}.bin"
filepath = database_save_path / filename
gt_points = points[point_indices[:, i]]
#print(gt_points)
gt_points[:, :3] -= rbbox_lidar[i, :3]
#print(gt_points)
with open(filepath, 'w') as f:
gt_points.tofile(f)
if names[i] in used_classes:
if relative_path:
db_path = str(database_save_path.stem + "/" + filename)
else:
db_path = str(filepath)
db_info = {
"name": names[i],
"path": db_path,
"image_idx": image_idx,
"gt_idx": gt_idxes[i],
"box3d_lidar": rbbox_lidar[i],
"num_points_in_gt": gt_points.shape[0],
"difficulty": difficulty[i],
# "group_id": -1,
# "bbox": bboxes[i],
}
local_group_id = group_ids[i]
# if local_group_id >= 0:
if local_group_id not in group_dict:
group_dict[local_group_id] = group_counter
group_counter += 1
db_info["group_id"] = group_dict[local_group_id]
if "score" in annos:
db_info["score"] = annos["score"][i]
all_db_infos[names[i]].append(db_info)
for k, v in all_db_infos.items():
print(f"load {len(v)} {k} database infos")
with open(db_info_save_path, 'wb') as f:
pickle.dump(all_db_infos, f)
def evaluation_nusc(self, detections, output_dir, clean_after=True):
version = self.version
eval_set_map = {
"v1.0-mini": "mini_train",
"v1.0-trainval": "val",
}
nusc_annos = {}
mapped_class_names = self._class_names
token2info = {}
for info in self._nusc_infos:
token2info[info["token"]] = info
for det in prog_bar(detections):
annos = []
boxes = _second_det_to_nusc_box(det)
for i, box in enumerate(boxes):
name = mapped_class_names[box.label]
velocity = box.velocity[:2].tolist()
if len(token2info[det["metadata"]["token"]]["sweeps"]) == 0:
velocity = (np.nan, np.nan)
box.velocity = np.array([*velocity, 0.0])
boxes = _lidar_nusc_box_to_global(
token2info[det["metadata"]["token"]], boxes,
mapped_class_names, "cvpr_2019")
for i, box in enumerate(boxes):
name = mapped_class_names[box.label]
velocity = box.velocity[:2].tolist()
nusc_anno = {
"sample_token": det["metadata"]["token"],
"translation": box.center.tolist(),
"size": box.wlh.tolist(),
"rotation": box.orientation.elements.tolist(),
"velocity": velocity,
"detection_name": name,
"detection_score": box.score,
"attribute_name": NuScenesDataset.DefaultAttribute[name],
}
annos.append(nusc_anno)
nusc_annos[det["metadata"]["token"]] = annos
nusc_submissions = {
"meta": {
"use_camera": False,
"use_lidar": True,
"use_radar": False,
"use_map": False,
"use_external": False,
},
"results": nusc_annos,
}
res_path = Path(output_dir) / "results_nusc.json"
with open(res_path, "w") as f:
json.dump(nusc_submissions, f)
print(f"evaluation_nusc(): dumping results to {res_path}")
gt_annos = self.ground_truth_annotations
if gt_annos is None:
return None
eval_main_file = Path(__file__).resolve().parent / "nusc_eval.py"
# why add \"{}\"? to support path with spaces.
cmd = f"python {str(eval_main_file)} --root_path=\"{str(self._root_path)}\""
cmd += f" --version={self.version} --eval_version={self.eval_version}"
cmd += f" --res_path=\"{str(res_path)}\" --eval_set={eval_set_map[self.version]}"
cmd += f" --output_dir=\"{output_dir}\""
# use subprocess can release all nusc memory after evaluation
subprocess.check_output(cmd, shell=True)
with open(Path(output_dir) / "metrics_summary.json", "r") as f:
metrics = json.load(f)
detail = {}
if clean_after: # only delete results_nusc.json when needed
res_path.unlink() # delete results_nusc.json since it's very large
result = f"Nusc {version} Evaluation\n"
for name in mapped_class_names:
detail[name] = {}
for k, v in metrics["label_aps"][name].items():
detail[name][f"dist@{k}"] = v
tp_errs = []
tp_names = []
for k, v in metrics["label_tp_errors"][name].items():
detail[name][k] = v
tp_errs.append(f"{v:.4f}")
tp_names.append(k)
threshs = ', '.join(list(metrics["label_aps"][name].keys()))
scores = list(metrics["label_aps"][name].values())
scores = ', '.join([f"{s * 100:.2f}" for s in scores])
result += f"{name} Nusc dist AP@{threshs} and TP errors\n"
result += scores
result += "\n"
result += ', '.join(tp_names) + ": " + ', '.join(tp_errs)
result += "\n"
return {
"results": {
"nusc": result
},
"detail": {
"nusc": detail
},
}
def _fill_trainval_infos(nusc,
train_scenes,
val_scenes,
test=False,
max_sweeps=9):
train_nusc_infos = []
val_nusc_infos = []
from pyquaternion import Quaternion
for sample in prog_bar(nusc.sample):
lidar_token = sample["data"]["LIDAR_TOP"]
cam_front_token = sample["data"]["CAM_FRONT"]
sd_rec = nusc.get('sample_data', sample['data']["LIDAR_TOP"])
cs_record = nusc.get('calibrated_sensor',
sd_rec['calibrated_sensor_token'])
pose_record = nusc.get('ego_pose', sd_rec['ego_pose_token'])
lidar_path, boxes, _ = nusc.get_sample_data(lidar_token)
cam_path, _, cam_intrinsic = nusc.get_sample_data(cam_front_token)
# if Path(lidar_path).exists():
info = {
"lidar_path": lidar_path,
"cam_front_path": cam_path,
"token": sample["token"],
"sweeps": [],
"lidar2ego_translation": cs_record['translation'],
"lidar2ego_rotation": cs_record['rotation'],
"ego2global_translation": pose_record['translation'],
"ego2global_rotation": pose_record['rotation'],
"timestamp": sample["timestamp"],
}
l2e_r = info["lidar2ego_rotation"]
l2e_t = info["lidar2ego_translation"]
e2g_r = info["ego2global_rotation"]
e2g_t = info["ego2global_translation"]
l2e_r_mat = Quaternion(l2e_r).rotation_matrix
e2g_r_mat = Quaternion(e2g_r).rotation_matrix
sd_rec = nusc.get('sample_data', sample['data']["LIDAR_TOP"])
sweeps = []
while len(sweeps) < max_sweeps:
if not sd_rec['prev'] == "":
sd_rec = nusc.get('sample_data', sd_rec['prev'])
cs_record = nusc.get('calibrated_sensor',
sd_rec['calibrated_sensor_token'])
pose_record = nusc.get('ego_pose', sd_rec['ego_pose_token'])
lidar_path = nusc.get_sample_data_path(sd_rec['token'])
sweep = {
"lidar_path": lidar_path,
"sample_data_token": sd_rec['token'],
"lidar2ego_translation": cs_record['translation'],
"lidar2ego_rotation": cs_record['rotation'],
"ego2global_translation": pose_record['translation'],
"ego2global_rotation": pose_record['rotation'],
"timestamp": sd_rec["timestamp"]
}
l2e_r_s = sweep["lidar2ego_rotation"]
l2e_t_s = sweep["lidar2ego_translation"]
e2g_r_s = sweep["ego2global_rotation"]
e2g_t_s = sweep["ego2global_translation"]
# sweep->ego->global->ego'->lidar
l2e_r_s_mat = Quaternion(l2e_r_s).rotation_matrix
e2g_r_s_mat = Quaternion(e2g_r_s).rotation_matrix
R = (l2e_r_s_mat.T @ e2g_r_s_mat.T) @ (
np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T)
T = (l2e_t_s @ e2g_r_s_mat.T + e2g_t_s) @ (
np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T)
T -= e2g_t @ (np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(
l2e_r_mat).T) + l2e_t @ np.linalg.inv(l2e_r_mat).T
sweep["sweep2lidar_rotation"] = R
sweep["sweep2lidar_translation"] = T
sweeps.append(sweep)
else:
break
info["sweeps"] = sweeps
if not test:
locs = np.array([b.center for b in boxes]).reshape(-1, 3)
dims = np.array([b.wlh for b in boxes]).reshape(-1, 3)
rots = np.array([b.orientation.yaw_pitch_roll[0]
for b in boxes]).reshape(-1, 1)
names = np.array([b.name for b in boxes])
gt_boxes = np.concatenate([locs, dims, -rots - np.pi / 2], axis=1)
info["gt_boxes"] = gt_boxes
info["gt_names"] = names
if sample["scene_token"] in train_scenes:
train_nusc_infos.append(info)
else:
val_nusc_infos.append(info)
return train_nusc_infos, val_nusc_infos
def _create_reduced_point_cloud(data_path,
info_path,
save_path=None,
color=True,
back=False):
bgr_path = None
with open(info_path, 'rb') as f:
kitti_infos = pickle.load(f)
for info in prog_bar(kitti_infos):
pc_info = info["point_cloud"]
image_info = info["image"]
image_path = image_info['image_path']
calib = info["calib"]
v_path = pc_info['velodyne_path']
v_path = Path(data_path) / v_path
points_v = np.fromfile(str(v_path), dtype=np.float32,
count=-1).reshape([-1, 4])
rect = calib['R0_rect']
P2 = calib['P2']
Trv2c = calib['Tr_velo_to_cam']
# first remove z < 0 points
# keep = points_v[:, -1] > 0
# points_v = points_v[keep]
# then remove outside.
if back:
points_v[:, 0] = -points_v[:, 0]
points_v = box_np_ops.remove_outside_points(points_v, rect, Trv2c, P2,
image_info["image_shape"])
points_xyz = lidar_to_camera(points_v, rect, Trv2c)
points_xy = project_to_image(points_xyz, P2)
image_path = v_path.parent.parent.parent / image_path
BGR_image = cv2.imread(str(image_path))
points_bgr = np.zeros((len(points_v), 3))
image_shape = image_info["image_shape"]
for i in range(len(points_bgr)):
new_y = int(points_xy[i][0])
if new_y >= image_shape[1]:
new_y = image_shape[1] - 1
new_x = int(points_xy[i][1])
if new_x >= image_shape[0]:
new_x = image_shape[0] - 1
points_bgr[i] = BGR_image[new_x][new_y] / 255
if bgr_path is None:
save_bgr_filename = v_path.parent.parent / (v_path.parent.stem +
"_bgr") / v_path.name
# save_filename = str(v_path) + '_reduced'
if back:
save_filename += "_bgr_back"
else:
save_bgr_filename = str(Path(bgr_path) / v_path.name)
if back:
save_bgr_filename += "_bgr_back"
with open(save_bgr_filename, 'w') as f:
points_bgr.tofile(f)
if save_path is None:
save_filename = v_path.parent.parent / (v_path.parent.stem +
"_bgr") / v_path.name
# save_filename = str(v_path) + '_reduced'
if back:
save_filename += "_back"
else:
save_filename = str(Path(save_path) / v_path.name)
if back:
save_filename += "_back"
with open(save_filename, 'w') as f:
points_v.tofile(f)
def create_my_groundtruth_database(dataset_class_name,
data_path,
info_path=None,
used_classes=None,
database_save_path=None,
db_info_save_path=None,
relative_path=True):
dataset = get_dataset_class(dataset_class_name)(
info_path=info_path,
root_path=data_path,
)
root_path = Path(data_path)
if database_save_path is None:
database_save_path = root_path / 'gt_database'
else:
database_save_path = Path(database_save_path)
if db_info_save_path is None:
db_info_save_path = root_path / "my_dbinfos_train.pkl"
database_save_path.mkdir(parents=True, exist_ok=True)
all_db_infos = {}
group_counter = 0
for j in prog_bar(list(range(len(dataset)))):
image_idx = j
sensor_data = dataset.get_sensor_data(j)
if "image_idx" in sensor_data["metadata"]:
image_idx = sensor_data["metadata"]["image_idx"]
points = sensor_data["lidar"]["points"]
annos = sensor_data["lidar"]["annotations"]
gt_boxes = annos["boxes"]
names = annos["names"]
group_dict = {}
group_ids = np.arange(gt_boxes.shape[0], dtype=np.int64)
difficulty = np.zeros(gt_boxes.shape[0], dtype=np.int32)
num_obj = gt_boxes.shape[0]
point_indices = None
if num_obj > 0:
point_indices = box_np_ops.points_in_rbbox(points, gt_boxes)
for i in range(num_obj):
filename = "{}_{}_{}.bin".format(image_idx, names[i], i)
filepath = database_save_path / filename
gt_points = points[point_indices[:, i]]
if gt_points.shape[0] >= 5:
gt_points[:, :3] -= gt_boxes[i, :3]
with open(filepath, 'w') as f:
gt_points.tofile(f)
if (used_classes is None) or names[i] in used_classes:
if relative_path:
db_path = str(database_save_path.stem + "/" + filename)
else:
db_path = str(filepath)
db_info = {
"name": names[i],
"path": db_path,
"image_idx": image_idx,
"gt_idx": i,
"box3d_lidar": gt_boxes[i],
"num_points_in_gt": gt_points.shape[0],
"difficulty": difficulty[i],
}
local_group_id = group_ids[i]
if local_group_id not in group_dict:
group_dict[local_group_id] = group_counter
group_counter += 1
db_info["group_id"] = group_dict[local_group_id]
if names[i] in all_db_infos:
all_db_infos[names[i]].append(db_info)
else:
all_db_infos[names[i]] = [db_info]
for k, v in all_db_infos.items():
print("load {} {} database infos".format(len(v), k))
with open(db_info_save_path, 'wb') as f:
pickle.dump(all_db_infos, f)
def _fill_trainval_infos(lyftdata,
use_flat_vehicle_coordinates,
use_second_format_direction,
calc_num_points,
is_test=False):
train_infos, val_infos = [], []
train_scene_tokens = [
scene['token'] for scene in lyftdata.scene
if scene['name'] in splits.train
]
for sample in prog_bar(sorted(lyftdata.sample, key=lambda s: s['timestamp'])):
lidar_token = sample['data']['LIDAR_TOP']
cam_front_token = sample['data']['CAM_FRONT']
sd_record = lyftdata.get('sample_data', lidar_token)
cs_record = lyftdata.get('calibrated_sensor', sd_record['calibrated_sensor_token'])
pose_record = lyftdata.get('ego_pose', sd_record['ego_pose_token'])
lidar_path, boxes_lidar, _ = lyftdata.get_sample_data(lidar_token)
cam_path, _, cam_intrinsic = lyftdata.get_sample_data(cam_front_token)
info = {
'lidar_path': lidar_path,
'cam_front_path': cam_path,
'token': sample['token'],
'lidar2ego_translation': cs_record['translation'],
'lidar2ego_rotation': cs_record['rotation'],
'ego2global_translation': pose_record['translation'],
'ego2global_rotation': pose_record['rotation'],
'timestamp': sample['timestamp'],
}
if not is_test:
locs = np.array([b.center for b in boxes_lidar]).reshape(-1, 3)
dims = np.array([b.wlh for b in boxes_lidar]).reshape(-1, 3)
rots = np.array([b.orientation.yaw_pitch_roll[0] for b in boxes_lidar]).reshape(-1, 1)
gt_boxes = np.concatenate([locs, dims, rots], axis=1)
gt_names = np.array([b.name for b in boxes_lidar])
if calc_num_points:
try:
pointcloud = LidarPointCloud.from_file(lidar_path)
except Exception as e:
print('ERROR:', e, lidar_path)
continue
indices = box_np_ops.points_in_rbbox(pointcloud.points.T[:, :3], gt_boxes)
num_points_in_gt = indices.sum(0)
info['num_lidar_pts'] = num_points_in_gt.astype(np.int32)
if use_flat_vehicle_coordinates:
_, boxes_flat_vehicle, _ = lyftdata.get_sample_data(
lidar_token,
use_flat_vehicle_coordinates=True
)
locs = np.array([b.center for b in boxes_flat_vehicle]).reshape(-1, 3)
dims = np.array([b.wlh for b in boxes_flat_vehicle]).reshape(-1, 3)
rots = np.array([b.orientation.yaw_pitch_roll[0] for b in boxes_flat_vehicle]).reshape(-1, 1)
gt_boxes = np.concatenate([locs, dims, rots], axis=1)
gt_names = np.array([b.name for b in boxes_flat_vehicle])
if use_second_format_direction:
gt_boxes[:, 6] = np.apply_along_axis(
lambda r: -ds_utils.wrap_to_pi(r + np.pi / 2),
axis=1,
arr=gt_boxes[:, 6:])
annotations = [
lyftdata.get('sample_annotation', token)
for token in sample['anns']
]
assert len(gt_boxes) == len(annotations), f"{len(gt_boxes)} != {len(annotations)}"
info['gt_boxes'] = gt_boxes
info['gt_names'] = gt_names
if sample['scene_token'] in train_scene_tokens:
train_infos.append(info)
else:
val_infos.append(info)
return train_infos, val_infos
def create_groundtruth_database_with_sweep_info(dataset_class_name,
data_path,
info_path=None,
used_classes=None,
database_save_path=None,
db_info_save_path=None,
relative_path=True,
add_rgb=False,
lidar_only=False,
bev_only=False,
coors_range=None):
dataset = get_dataset_class(dataset_class_name)(
info_path=info_path,
root_path=data_path,
)
root_path = Path(data_path)
if database_save_path is None:
database_save_path = root_path / 'gt_database_with_sweep_info'
else:
database_save_path = Path(database_save_path)
if db_info_save_path is None:
db_info_save_path = root_path / "dbinfos_train_with_sweep_info.pkl"
database_save_path.mkdir(parents=True, exist_ok=True)
all_db_infos = {}
group_counter = 0
for j in prog_bar(list(range(len(dataset)))):
image_idx = j
sensor_data = dataset.get_sensor_data(j)
if "image_idx" in sensor_data["metadata"]:
image_idx = sensor_data["metadata"]["image_idx"]
assert("sweep_points_list" in sensor_data["lidar"])
sweep_points_list = sensor_data["lidar"]["sweep_points_list"]
assert("sweep_annotations" in sensor_data["lidar"])
sweep_gt_boxes_list = sensor_data["lidar"]["sweep_annotations"]["boxes_list"]
sweep_gt_tokens_list = sensor_data["lidar"]["sweep_annotations"]["tokens_list"]
sweep_gt_names_list = sensor_data["lidar"]["sweep_annotations"]["names_list"]
# we focus on the objects in the first frame
# and find the bounding box index of the same object in every frame
points = sensor_data["lidar"]["points"]
annos = sensor_data["lidar"]["annotations"]
names = annos["names"]
gt_boxes = annos["boxes"]
attrs = annos["attrs"]
tokens = sweep_gt_tokens_list[0]
# sanity check with redundancy
assert(len(sweep_gt_boxes_list) == len(sweep_gt_tokens_list) == len(sweep_gt_names_list))
assert(len(gt_boxes) == len(attrs) == len(tokens))
assert(np.all(names == sweep_gt_names_list[0]))
assert(np.all(gt_boxes == sweep_gt_boxes_list[0]))
# on every frame, we mask points inside each bounding box
# but we are not looking at every bounding box
sweep_point_indices_list = []
for sweep_points, sweep_gt_boxes in zip(sweep_points_list, sweep_gt_boxes_list):
sweep_point_indices_list.append(box_np_ops.points_in_rbbox(sweep_points, sweep_gt_boxes))
# crop point cloud based on boxes in the current frame
point_indices = box_np_ops.points_in_rbbox(points, gt_boxes)
#
group_dict = {}
group_ids = np.full([gt_boxes.shape[0]], -1, dtype=np.int64)
if "group_ids" in annos:
group_ids = annos["group_ids"]
else:
group_ids = np.arange(gt_boxes.shape[0], dtype=np.int64)
#
difficulty = np.zeros(gt_boxes.shape[0], dtype=np.int32)
if "difficulty" in annos:
difficulty = annos["difficulty"]
num_obj = gt_boxes.shape[0]
for i in range(num_obj):
filename = f"{image_idx}_{names[i]}_{i}.bin"
filepath = database_save_path / filename
# only use non-key frame boxes when the object is moving
if attrs[i] in ['vehicle.moving', 'cycle.with_rider', 'pedestrian.moving']:
gt_points_list = []
for t in range(len(sweep_points_list)):
# fast pass for most frames
if i < len(sweep_gt_tokens_list[t]) and sweep_gt_tokens_list[t][i] == tokens[i]:
box_idx = i
else:
I = np.flatnonzero(tokens[i] == sweep_gt_tokens_list[t])
if len(I) == 0: continue
elif len(I) == 1: box_idx = I[0]
else: raise ValueError('Identical object tokens')
gt_points_list.append(sweep_points_list[t][sweep_point_indices_list[t][:, box_idx]])
gt_points = np.concatenate(gt_points_list, axis=0)[:, [0, 1, 2, 4]]
else:
gt_points = points[point_indices[:, i]]
# offset points based on the bounding box in the current frame
gt_points[:, :3] -= gt_boxes[i, :3]
with open(filepath, 'w') as f:
gt_points.tofile(f)
if (used_classes is None) or names[i] in used_classes:
if relative_path:
db_path = str(database_save_path.stem + "/" + filename)
else:
db_path = str(filepath)
db_info = {
"name": names[i],
"path": db_path,
"image_idx": image_idx,
"gt_idx": i,
"box3d_lidar": gt_boxes[i],
"num_points_in_gt": gt_points.shape[0],
"difficulty": difficulty[i]
}
local_group_id = group_ids[i]
if local_group_id not in group_dict:
group_dict[local_group_id] = group_counter
group_counter += 1
db_info["group_id"] = group_dict[local_group_id]
if "score" in annos:
db_info["score"] = annos["score"][i]
if names[i] in all_db_infos:
all_db_infos[names[i]].append(db_info)
else:
all_db_infos[names[i]] = [db_info]
for k, v in all_db_infos.items():
print(f"load {len(v)} {k} database infos")
with open(db_info_save_path, 'wb') as f:
pickle.dump(all_db_infos, f)
def _fill_train_infos(root_path):
train_udi_infos = []
lidar_root_path = root_path+ "/lidar"
label_root_path = root_path + "/label"
img_root_path = root_path + "/image"
filenames = os.listdir(lidar_root_path)
for filename in prog_bar(filenames):
index = filename.split(".")[0]
lidar_path = lidar_root_path + "/" + index + ".bin"
cam_path = img_root_path + "/" + index + ".jpg"
label_path = label_root_path + "/" + index + "_bin.json"
assert Path(lidar_path).exists()
assert Path(cam_path).exists()
assert Path(label_path).exists()
with open(label_path, encoding='utf-8') as f:
res = f.read()
result = json.loads(res)
boxes = result["elem"]
info = {
"lidar_path": lidar_path,
"cam_front_path": cam_path,
"filename": filename,
"token": int(index),
}
gt_locs_list = []
gt_dims_list = []
print("label file path:", label_path)
for box in boxes:
box_loc = box["position"]
box_size = box["size"]
box_loc_ = np.array([box_loc["x"],box_loc["y"], box_loc["z"]], dtype=np.float)
box_size_ = np.array([box_size["width"],box_size["depth"],box_size["height"]], dtype=np.float)
box_loc_ = box_loc_.reshape(-1, 3)
box_size_ = box_size_.reshape(-1, 3)
gt_locs_list.append(box_loc_)
gt_dims_list.append(box_size_)
locs = np.concatenate(gt_locs_list, axis=0)
dims = np.concatenate(gt_dims_list, axis=0)
rots = np.array([b["yaw"] for b in boxes]).reshape(-1, 1)
names = [b["class"] for b in boxes]
for i in range(len(names)):
if names[i] in UDIDataset.NameMapping:
names[i] = UDIDataset.NameMapping[names[i]]
names = np.array(names)
# we need to convert rot to SECOND format.
# change the rot format will break all checkpoint.
gt_boxes = np.concatenate([locs, dims, -rots - np.pi / 2], axis=1)
info["gt_boxes"] = gt_boxes
info["gt_names"] = names
train_udi_infos.append(info)
return train_udi_infos
