Spaces:
Running
Running
| from PIL import Image | |
| import numpy as np | |
| import os | |
| from tqdm import tqdm | |
| from roma.utils import pose_auc | |
| import cv2 | |
| class HpatchesHomogBenchmark: | |
| """Hpatches grid goes from [0,n-1] instead of [0.5,n-0.5]""" | |
| def __init__(self, dataset_path) -> None: | |
| seqs_dir = "hpatches-sequences-release" | |
| self.seqs_path = os.path.join(dataset_path, seqs_dir) | |
| self.seq_names = sorted(os.listdir(self.seqs_path)) | |
| # Ignore seqs is same as LoFTR. | |
| self.ignore_seqs = set( | |
| [ | |
| "i_contruction", | |
| "i_crownnight", | |
| "i_dc", | |
| "i_pencils", | |
| "i_whitebuilding", | |
| "v_artisans", | |
| "v_astronautis", | |
| "v_talent", | |
| ] | |
| ) | |
| def convert_coordinates(self, im_A_coords, im_A_to_im_B, wq, hq, wsup, hsup): | |
| offset = 0.5 # Hpatches assumes that the center of the top-left pixel is at [0,0] (I think) | |
| im_A_coords = ( | |
| np.stack( | |
| ( | |
| wq * (im_A_coords[..., 0] + 1) / 2, | |
| hq * (im_A_coords[..., 1] + 1) / 2, | |
| ), | |
| axis=-1, | |
| ) | |
| - offset | |
| ) | |
| im_A_to_im_B = ( | |
| np.stack( | |
| ( | |
| wsup * (im_A_to_im_B[..., 0] + 1) / 2, | |
| hsup * (im_A_to_im_B[..., 1] + 1) / 2, | |
| ), | |
| axis=-1, | |
| ) | |
| - offset | |
| ) | |
| return im_A_coords, im_A_to_im_B | |
| def benchmark(self, model, model_name=None): | |
| n_matches = [] | |
| homog_dists = [] | |
| for seq_idx, seq_name in tqdm( | |
| enumerate(self.seq_names), total=len(self.seq_names) | |
| ): | |
| im_A_path = os.path.join(self.seqs_path, seq_name, "1.ppm") | |
| im_A = Image.open(im_A_path) | |
| w1, h1 = im_A.size | |
| for im_idx in range(2, 7): | |
| im_B_path = os.path.join(self.seqs_path, seq_name, f"{im_idx}.ppm") | |
| im_B = Image.open(im_B_path) | |
| w2, h2 = im_B.size | |
| H = np.loadtxt( | |
| os.path.join(self.seqs_path, seq_name, "H_1_" + str(im_idx)) | |
| ) | |
| dense_matches, dense_certainty = model.match(im_A_path, im_B_path) | |
| good_matches, _ = model.sample(dense_matches, dense_certainty, 5000) | |
| pos_a, pos_b = self.convert_coordinates( | |
| good_matches[:, :2], good_matches[:, 2:], w1, h1, w2, h2 | |
| ) | |
| try: | |
| H_pred, inliers = cv2.findHomography( | |
| pos_a, | |
| pos_b, | |
| method=cv2.RANSAC, | |
| confidence=0.99999, | |
| ransacReprojThreshold=3 * min(w2, h2) / 480, | |
| ) | |
| except: | |
| H_pred = None | |
| if H_pred is None: | |
| H_pred = np.zeros((3, 3)) | |
| H_pred[2, 2] = 1.0 | |
| corners = np.array( | |
| [[0, 0, 1], [0, h1 - 1, 1], [w1 - 1, 0, 1], [w1 - 1, h1 - 1, 1]] | |
| ) | |
| real_warped_corners = np.dot(corners, np.transpose(H)) | |
| real_warped_corners = ( | |
| real_warped_corners[:, :2] / real_warped_corners[:, 2:] | |
| ) | |
| warped_corners = np.dot(corners, np.transpose(H_pred)) | |
| warped_corners = warped_corners[:, :2] / warped_corners[:, 2:] | |
| mean_dist = np.mean( | |
| np.linalg.norm(real_warped_corners - warped_corners, axis=1) | |
| ) / (min(w2, h2) / 480.0) | |
| homog_dists.append(mean_dist) | |
| n_matches = np.array(n_matches) | |
| thresholds = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] | |
| auc = pose_auc(np.array(homog_dists), thresholds) | |
| return { | |
| "hpatches_homog_auc_3": auc[2], | |
| "hpatches_homog_auc_5": auc[4], | |
| "hpatches_homog_auc_10": auc[9], | |
| } | |