Spaces:
Running
Running
| import pickle | |
| import h5py | |
| import numpy as np | |
| import torch | |
| from dkm.utils import * | |
| from PIL import Image | |
| from tqdm import tqdm | |
| class Yfcc100mBenchmark: | |
| def __init__(self, data_root="data/yfcc100m_test") -> None: | |
| self.scenes = [ | |
| "buckingham_palace", | |
| "notre_dame_front_facade", | |
| "reichstag", | |
| "sacre_coeur", | |
| ] | |
| self.data_root = data_root | |
| def benchmark(self, model, r=2): | |
| model.train(False) | |
| with torch.no_grad(): | |
| data_root = self.data_root | |
| meta_info = open( | |
| f"{data_root}/yfcc_test_pairs_with_gt.txt", "r" | |
| ).readlines() | |
| tot_e_t, tot_e_R, tot_e_pose = [], [], [] | |
| for scene_ind in range(len(self.scenes)): | |
| scene = self.scenes[scene_ind] | |
| pairs = np.array( | |
| pickle.load( | |
| open(f"{data_root}/pairs/{scene}-te-1000-pairs.pkl", "rb") | |
| ) | |
| ) | |
| scene_dir = f"{data_root}/yfcc100m/{scene}/test/" | |
| calibs = open(scene_dir + "calibration.txt", "r").read().split("\n") | |
| images = open(scene_dir + "images.txt", "r").read().split("\n") | |
| pair_inds = np.random.choice( | |
| range(len(pairs)), size=len(pairs), replace=False | |
| ) | |
| for pairind in tqdm(pair_inds): | |
| idx1, idx2 = pairs[pairind] | |
| params = meta_info[1000 * scene_ind + pairind].split() | |
| rot1, rot2 = int(params[2]), int(params[3]) | |
| calib1 = h5py.File(scene_dir + calibs[idx1], "r") | |
| K1, R1, t1, _, _ = get_pose(calib1) | |
| calib2 = h5py.File(scene_dir + calibs[idx2], "r") | |
| K2, R2, t2, _, _ = get_pose(calib2) | |
| R, t = compute_relative_pose(R1, t1, R2, t2) | |
| im1 = images[idx1] | |
| im2 = images[idx2] | |
| im1 = Image.open(scene_dir + im1).rotate(rot1 * 90, expand=True) | |
| w1, h1 = im1.size | |
| im2 = Image.open(scene_dir + im2).rotate(rot2 * 90, expand=True) | |
| w2, h2 = im2.size | |
| K1 = rotate_intrinsic(K1, rot1) | |
| K2 = rotate_intrinsic(K2, rot2) | |
| dense_matches, dense_certainty = model.match(im1, im2) | |
| dense_certainty = dense_certainty ** (1 / r) | |
| sparse_matches, sparse_confidence = model.sample( | |
| dense_matches, dense_certainty, 10000 | |
| ) | |
| scale1 = 480 / min(w1, h1) | |
| scale2 = 480 / min(w2, h2) | |
| w1, h1 = scale1 * w1, scale1 * h1 | |
| w2, h2 = scale2 * w2, scale2 * h2 | |
| K1 = K1 * scale1 | |
| K2 = K2 * scale2 | |
| kpts1 = sparse_matches[:, :2] | |
| kpts1 = np.stack( | |
| (w1 * kpts1[:, 0] / 2, h1 * kpts1[:, 1] / 2), axis=-1 | |
| ) | |
| kpts2 = sparse_matches[:, 2:] | |
| kpts2 = np.stack( | |
| (w2 * kpts2[:, 0] / 2, h2 * kpts2[:, 1] / 2), axis=-1 | |
| ) | |
| try: | |
| threshold = 1.0 | |
| norm_threshold = threshold / ( | |
| np.mean(np.abs(K1[:2, :2])) + np.mean(np.abs(K2[:2, :2])) | |
| ) | |
| R_est, t_est, mask = estimate_pose( | |
| kpts1, | |
| kpts2, | |
| K1[:2, :2], | |
| K2[:2, :2], | |
| norm_threshold, | |
| conf=0.9999999, | |
| ) | |
| T1_to_2 = np.concatenate((R_est, t_est), axis=-1) # | |
| e_t, e_R = compute_pose_error(T1_to_2, R, t) | |
| e_pose = max(e_t, e_R) | |
| except: | |
| 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, | |
| } | |