def forward(self, kp1, kp2, w_kp1, w_kp2, w_vis_kp1_mask, w_vis_kp2_mask,
shift_scale1, shift_scale2, intrinsics1, intrinsics2,
extrinsics1, extrinsics2):
mutual_gt_matches_mask, nn_kp_ids = \
get_gt_matches(kp1, kp2, w_kp1, w_kp2, w_vis_kp1_mask, w_vis_kp2_mask, self.px_thresh)
r_kp1 = revert_data_transform(kp1, shift_scale1)
r_kp2 = revert_data_transform(kp2, shift_scale2)
nn_r_kp2 = select_kp(r_kp2, nn_kp_ids)
F = compose_gt_transform(intrinsics1, intrinsics2, extrinsics1,
extrinsics2)
# TODO. Squared epipolar distance maybe?
ep_dist = epipolar_distance(r_kp1, nn_r_kp2,
F).clamp(max=self.px_thresh)
loss = ep_dist * mutual_gt_matches_mask.float()
loss = loss.sum(dim=-1) / mutual_gt_matches_mask.float().sum(
dim=-1).clamp(min=1e-8)
loss = self.loss_lambda * loss.mean()
return loss
def relative_pose_error(kp1,
kp2,
kp1_desc,
kp2_desc,
shift_scale1,
shift_scale2,
intrinsics1,
intrinsics2,
extrinsics1,
extrinsics2,
px_thresh,
detailed=False):
"""
:param kp1: B x N x 2
:param kp2: B x N x 2
:param kp1_desc: B x N x C
:param kp2_desc: B x N x C
:param shift_scale1: B x 4
:param shift_scale2: B x 4
:param intrinsics1: B x 3 x 3
:param intrinsics2: B x 3 x 3
:param extrinsics1: B x 4 x 4
:param extrinsics2: B x 4 x 4
:param px_thresh: list
:param detailed: bool
"""
mutual_desc_matches_mask, nn_desc_ids = get_mutual_desc_matches(
kp1_desc, kp2_desc, None, 0.9)
r_kp1 = revert_data_transform(kp1, shift_scale1)
r_kp2 = revert_data_transform(kp2, shift_scale2)
nn_r_kp2 = select_kp(r_kp2, nn_desc_ids)
num_thresh = len(px_thresh)
b, n = kp1.shape[:2]
gt_rel_pose = get_gt_rel_pose(extrinsics1, extrinsics2)
R_err = torch.zeros(num_thresh, b)
t_err = torch.zeros(num_thresh, b)
est_inl_mask = torch.zeros(num_thresh, b, n,
dtype=torch.bool).to(kp1.device)
for i, thresh in enumerate(px_thresh):
i_est_rel_pose, i_est_inl_mask = prepare_rel_pose(
r_kp1, nn_r_kp2, mutual_desc_matches_mask, intrinsics1,
intrinsics2, thresh)
R_err[i] = angle_mat(i_est_rel_pose[:, :3, :3], gt_rel_pose[:, :3, :3])
t_err[i] = angle_vec(i_est_rel_pose[:, :3, 3], gt_rel_pose[:, :3, 3])
if detailed:
return R_err, t_err, est_inl_mask, mutual_desc_matches_mask, nn_desc_ids
else:
return R_err, t_err, est_inl_mask
def forward(self, kp1, kp2, w_kp1, w_kp2, w_vis_kp1_mask, w_vis_kp2_mask,
shift_scale1, shift_scale2, intrinsics1, intrinsics2,
extrinsics1, extrinsics2):
mutual_gt_matches_mask, nn_kp_ids = \
get_gt_matches(kp1, kp2, w_kp1, w_kp2, w_vis_kp1_mask, w_vis_kp2_mask, self.px_thresh)
r_kp1 = revert_data_transform(kp1, shift_scale1)
r_kp2 = revert_data_transform(kp2, shift_scale2)
nn_r_kp2 = select_kp(r_kp2, nn_kp_ids)
nn_r_i1_kp2 = change_intrinsics(nn_r_kp2, intrinsics2, intrinsics1)
gt_E_param = compose_gt_transform(intrinsics1, intrinsics2,
extrinsics1, extrinsics2, E_param)
if self.loss_version == '1':
t_px_thresh = torch.tensor(self.px_thresh).to(kp1.device)
est_E_param, success_mask = ParamRelPose().apply(
r_kp1, nn_r_i1_kp2, mutual_gt_matches_mask, intrinsics1,
intrinsics2, t_px_thresh)
elif self.loss_version == '2':
t_px_thresh = torch.tensor(self.px_thresh).to(kp1.device)
est_E_param, success_mask = ParamTruncRelPose().apply(
r_kp1, nn_r_i1_kp2, nn_r_kp2, mutual_gt_matches_mask,
intrinsics1, intrinsics2, extrinsics1, extrinsics2,
t_px_thresh)
else:
raise NotImplementedError
loss = (est_E_param - gt_E_param).norm(dim=-1)
loss = loss.sum() / success_mask.float().sum().clamp(min=1e-8)
loss = self.loss_lambda * loss
return loss
def relative_param_pose_error(kp1, kp2, kp1_desc, kp2_desc, shift_scale1,
shift_scale2, intrinsics1, intrinsics2,
extrinsics1, extrinsics2, px_thresh):
mutual_desc_matches_mask, nn_desc_ids = get_mutual_desc_matches(
kp1_desc, kp2_desc, None, 0.9)
r_kp1 = revert_data_transform(kp1, shift_scale1)
r_kp2 = revert_data_transform(kp2, shift_scale2)
nn_r_kp2 = select_kp(r_kp2, nn_desc_ids)
nn_r_i1_kp2 = change_intrinsics(nn_r_kp2, intrinsics2, intrinsics1)
gt_E_param = compose_gt_transform(intrinsics1, intrinsics2, extrinsics1,
extrinsics2, E_param)
num_thresh = len(px_thresh)
b = kp1.shape[0]
R_param_err = torch.zeros(num_thresh, b)
t_param_err = torch.zeros(num_thresh, b)
success_mask = torch.zeros(num_thresh, b, dtype=torch.bool)
for i, thresh in enumerate(px_thresh):
i_est_E_param, i_success_mask = prepare_param_rel_pose(
r_kp1, nn_r_i1_kp2, mutual_desc_matches_mask, intrinsics1,
intrinsics2, thresh)
R_param_err[i] = (i_est_E_param[:, :3] -
gt_E_param[:, :3]).norm(dim=-1)
t_param_err[i] = (i_est_E_param[:, 3:] -
gt_E_param[:, 3:]).norm(dim=-1)
success_mask[i] = i_success_mask
return R_param_err, t_param_err, success_mask
def save_aachen_inference(dataset_config, output):
batch, endpoint = output
scene_name, image1_name = batch.get(du.SCENE_NAME)[0], batch.get(
du.IMAGE1_NAME)[0]
kp1 = endpoint[eu.KP1]
shift_scale1 = batch.get(du.SHIFT_SCALE1).to(kp1.device)
kp1 = revert_data_transform(kp1, shift_scale1).cpu().squeeze().numpy()
kp1_desc = endpoint[eu.KP1_DESC].cpu().squeeze().numpy()
dataset_root = dataset_config[du.AACHEN][du.DATASET_ROOT]
method_name = "NetJoint"
file_path = os.path.join(dataset_root, scene_name,
f"{image1_name}.{method_name}.npz")
r_file_path = os.path.join(dataset_root, scene_name,
f"{image1_name}.{method_name}")
np.savez(file_path, keypoints=kp1, descriptors=kp1_desc)
os.rename(file_path, r_file_path)
def epipolar_match_score(kp1,
kp2,
w_kp1,
w_kp2,
w_vis_kp1_mask,
w_vis_kp2_mask,
kp1_desc,
kp2_desc,
shift_scale1,
shift_scale2,
intrinsics1,
intrinsics2,
extrinsics1,
extrinsics2,
px_thresh,
dd_measure,
detailed=False):
mutual_desc_matches_mask, nn_desc_ids = get_mutual_desc_matches(
kp1_desc, kp2_desc, dd_measure, None)
# Verify matches by using the largest pixel threshold
v_mutual_desc_matches_mask, nn_kp_values = verify_mutual_desc_matches(
nn_desc_ids,
kp1,
kp2,
w_kp1,
w_kp2,
w_vis_kp1_mask,
w_vis_kp2_mask,
px_thresh[-1],
return_reproj=True)
# Select minimum number of visible points for each scene
num_vis_gt_matches = get_num_vis_gt_matches(w_vis_kp1_mask, w_vis_kp2_mask)
num_thresh = len(px_thresh)
b, n = kp1.shape[:2]
o_kp1 = revert_data_transform(kp1, shift_scale1)
o_kp2 = revert_data_transform(kp2, shift_scale2)
nn_o_kp2 = select_kp(o_kp2, nn_desc_ids)
F = compose_gt_transform(intrinsics1, intrinsics2, extrinsics1,
extrinsics2)
ep_dist = epipolar_distance(o_kp1, nn_o_kp2, F)
if detailed:
em_scores = torch.zeros(num_thresh, b)
num_matches = torch.zeros(num_thresh, b)
match_mask = torch.zeros(num_thresh, b, n)
else:
em_scores = torch.zeros(num_thresh)
for i, thresh in enumerate(px_thresh):
if i != num_thresh - 1:
i_mutual_matches_mask = mutual_desc_matches_mask * nn_kp_values.le(
thresh) * ep_dist.le(thresh)
else:
i_mutual_matches_mask = mutual_desc_matches_mask * v_mutual_desc_matches_mask * ep_dist.le(
thresh)
i_num_matches = i_mutual_matches_mask.sum(dim=-1).float()
if detailed:
em_scores[i] = i_num_matches / num_vis_gt_matches
num_matches[i] = i_num_matches
match_mask[i] = i_mutual_matches_mask
else:
em_scores[i] = (i_num_matches / num_vis_gt_matches).mean()
if detailed:
return em_scores, num_matches, num_vis_gt_matches, nn_desc_ids, match_mask
else:
return em_scores