import numpy as np import torch from dkm.utils import * from PIL import Image from tqdm import tqdm import torch.nn.functional as F class Megadepth1500Benchmark: def __init__(self, data_root="data/megadepth", scene_names = None) -> None: if scene_names is None: self.scene_names = [ "0015_0.1_0.3.npz", "0015_0.3_0.5.npz", "0022_0.1_0.3.npz", "0022_0.3_0.5.npz", "0022_0.5_0.7.npz", ] else: self.scene_names = scene_names self.scenes = [ np.load(f"{data_root}/{scene}", allow_pickle=True) for scene in self.scene_names ] self.data_root = data_root def benchmark(self, model): with torch.no_grad(): data_root = self.data_root tot_e_t, tot_e_R, tot_e_pose = [], [], [] for scene_ind in range(len(self.scenes)): scene = self.scenes[scene_ind] pairs = scene["pair_infos"] intrinsics = scene["intrinsics"] poses = scene["poses"] im_paths = scene["image_paths"] pair_inds = range(len(pairs)) for pairind in tqdm(pair_inds): idx1, idx2 = pairs[pairind][0] K1 = intrinsics[idx1].copy() T1 = poses[idx1].copy() R1, t1 = T1[:3, :3], T1[:3, 3] K2 = intrinsics[idx2].copy() T2 = poses[idx2].copy() R2, t2 = T2[:3, :3], T2[:3, 3] R, t = compute_relative_pose(R1, t1, R2, t2) im1_path = f"{data_root}/{im_paths[idx1]}" im2_path = f"{data_root}/{im_paths[idx2]}" im1 = Image.open(im1_path) w1, h1 = im1.size im2 = Image.open(im2_path) w2, h2 = im2.size scale1 = 1200 / max(w1, h1) scale2 = 1200 / max(w2, h2) w1, h1 = scale1 * w1, scale1 * h1 w2, h2 = scale2 * w2, scale2 * h2 K1[:2] = K1[:2] * scale1 K2[:2] = K2[:2] * scale2 dense_matches, dense_certainty = model.match(im1_path, im2_path) sparse_matches,_ = model.sample( dense_matches, dense_certainty, 5000 ) kpts1 = sparse_matches[:, :2] kpts1 = ( torch.stack( ( w1 * (kpts1[:, 0] + 1) / 2, h1 * (kpts1[:, 1] + 1) / 2, ), axis=-1, ) ) kpts2 = sparse_matches[:, 2:] kpts2 = ( torch.stack( ( w2 * (kpts2[:, 0] + 1) / 2, h2 * (kpts2[:, 1] + 1) / 2, ), axis=-1, ) ) for _ in range(5): shuffling = np.random.permutation(np.arange(len(kpts1))) kpts1 = kpts1[shuffling] kpts2 = kpts2[shuffling] try: norm_threshold = 0.5 / ( np.mean(np.abs(K1[:2, :2])) + np.mean(np.abs(K2[:2, :2]))) R_est, t_est, mask = estimate_pose( kpts1.cpu().numpy(), kpts2.cpu().numpy(), K1, K2, norm_threshold, conf=0.99999, ) T1_to_2_est = np.concatenate((R_est, t_est), axis=-1) # e_t, e_R = compute_pose_error(T1_to_2_est, R, t) e_pose = max(e_t, e_R) except Exception as e: print(repr(e)) e_t, e_R = 90, 90 e_pose = max(e_t, e_R) tot_e_t.append(e_t) tot_e_R.append(e_R) tot_e_pose.append(e_pose) tot_e_pose = np.array(tot_e_pose) thresholds = [5, 10, 20] auc = pose_auc(tot_e_pose, thresholds) acc_5 = (tot_e_pose < 5).mean() acc_10 = (tot_e_pose < 10).mean() acc_15 = (tot_e_pose < 15).mean() acc_20 = (tot_e_pose < 20).mean() map_5 = acc_5 map_10 = np.mean([acc_5, acc_10]) map_20 = np.mean([acc_5, acc_10, acc_15, acc_20]) return { "auc_5": auc[0], "auc_10": auc[1], "auc_20": auc[2], "map_5": map_5, "map_10": map_10, "map_20": map_20, }