Spaces:
Running
Running
| import numpy as np | |
| import torch | |
| from kornia.geometry.epipolar import numeric | |
| from kornia.geometry.conversions import convert_points_to_homogeneous | |
| import cv2 | |
| def pose_auc(errors, thresholds): | |
| sort_idx = np.argsort(errors) | |
| errors = np.array(errors.copy())[sort_idx] | |
| recall = (np.arange(len(errors)) + 1) / len(errors) | |
| errors = np.r_[0.0, errors] | |
| recall = np.r_[0.0, recall] | |
| aucs = [] | |
| for t in thresholds: | |
| last_index = np.searchsorted(errors, t) | |
| r = np.r_[recall[:last_index], recall[last_index - 1]] | |
| e = np.r_[errors[:last_index], t] | |
| aucs.append(np.trapz(r, x=e) / t) | |
| return aucs | |
| def angle_error_vec(v1, v2): | |
| n = np.linalg.norm(v1) * np.linalg.norm(v2) | |
| return np.rad2deg(np.arccos(np.clip(np.dot(v1, v2) / n, -1.0, 1.0))) | |
| def angle_error_mat(R1, R2): | |
| cos = (np.trace(np.dot(R1.T, R2)) - 1) / 2 | |
| cos = np.clip(cos, -1.0, 1.0) # numercial errors can make it out of bounds | |
| return np.rad2deg(np.abs(np.arccos(cos))) | |
| def symmetric_epipolar_distance(pts0, pts1, E, K0, K1): | |
| """Squared symmetric epipolar distance. | |
| This can be seen as a biased estimation of the reprojection error. | |
| Args: | |
| pts0 (torch.Tensor): [N, 2] | |
| E (torch.Tensor): [3, 3] | |
| """ | |
| pts0 = (pts0 - K0[[0, 1], [2, 2]][None]) / K0[[0, 1], [0, 1]][None] | |
| pts1 = (pts1 - K1[[0, 1], [2, 2]][None]) / K1[[0, 1], [0, 1]][None] | |
| pts0 = convert_points_to_homogeneous(pts0) | |
| pts1 = convert_points_to_homogeneous(pts1) | |
| Ep0 = pts0 @ E.T # [N, 3] | |
| p1Ep0 = torch.sum(pts1 * Ep0, -1) # [N,] | |
| Etp1 = pts1 @ E # [N, 3] | |
| d = p1Ep0**2 * (1.0 / (Ep0[:, 0]**2 + Ep0[:, 1]**2) + 1.0 / (Etp1[:, 0]**2 + Etp1[:, 1]**2)) # N | |
| return d | |
| def compute_symmetrical_epipolar_errors(T_0to1, pts0, pts1, K0, K1, device='cuda'): | |
| """ | |
| Update: | |
| data (dict):{"epi_errs": [M]} | |
| """ | |
| pts0 = torch.tensor(pts0, device=device) | |
| pts1 = torch.tensor(pts1, device=device) | |
| K0 = torch.tensor(K0, device=device) | |
| K1 = torch.tensor(K1, device=device) | |
| T_0to1 = torch.tensor(T_0to1, device=device) | |
| Tx = numeric.cross_product_matrix(T_0to1[:3, 3]) | |
| E_mat = Tx @ T_0to1[:3, :3] | |
| epi_err = symmetric_epipolar_distance(pts0, pts1, E_mat, K0, K1) | |
| return epi_err | |
| def compute_pose_error(T_0to1, R, t): | |
| R_gt = T_0to1[:3, :3] | |
| t_gt = T_0to1[:3, 3] | |
| error_t = angle_error_vec(t.squeeze(), t_gt) | |
| error_t = np.minimum(error_t, 180 - error_t) # ambiguity of E estimation | |
| error_R = angle_error_mat(R, R_gt) | |
| return error_t, error_R | |
| def compute_relative_pose(R1, t1, R2, t2): | |
| rots = R2 @ (R1.T) | |
| trans = -rots @ t1 + t2 | |
| return rots, trans | |
| def estimate_pose(kpts0, kpts1, K0, K1, norm_thresh, conf=0.99999): | |
| if len(kpts0) < 5: | |
| return None | |
| K0inv = np.linalg.inv(K0[:2,:2]) | |
| K1inv = np.linalg.inv(K1[:2,:2]) | |
| kpts0 = (K0inv @ (kpts0-K0[None,:2,2]).T).T | |
| kpts1 = (K1inv @ (kpts1-K1[None,:2,2]).T).T | |
| E, mask = cv2.findEssentialMat( | |
| kpts0, kpts1, np.eye(3), threshold=norm_thresh, prob=conf | |
| ) | |
| ret = None | |
| if E is not None: | |
| best_num_inliers = 0 | |
| for _E in np.split(E, len(E) / 3): | |
| n, R, t, _ = cv2.recoverPose(_E, kpts0, kpts1, np.eye(3), 1e9, mask=mask) | |
| if n > best_num_inliers: | |
| best_num_inliers = n | |
| ret = (R, t, mask.ravel() > 0) | |
| return ret | |
| def dynamic_alpha(n_matches, | |
| milestones=[0, 300, 1000, 2000], | |
| alphas=[1.0, 0.8, 0.4, 0.2]): | |
| if n_matches == 0: | |
| return 1.0 | |
| ranges = list(zip(alphas, alphas[1:] + [None])) | |
| loc = bisect.bisect_right(milestones, n_matches) - 1 | |
| _range = ranges[loc] | |
| if _range[1] is None: | |
| return _range[0] | |
| return _range[1] + (milestones[loc + 1] - n_matches) / ( | |
| milestones[loc + 1] - milestones[loc]) * (_range[0] - _range[1]) |