Spaces:
Running
Running
image-matching-webui
/
third_party
/DKM
/dkm
/benchmarks
/deprecated
/hpatches_sequences_dense_benchmark.py
from PIL import Image | |
import numpy as np | |
import os | |
import torch | |
from tqdm import tqdm | |
from dkm.utils import * | |
class HpatchesDenseBenchmark: | |
"""WARNING: 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)) | |
def convert_coordinates(self, query_coords, query_to_support, wq, hq, wsup, hsup): | |
# Get matches in output format on the grid [0, n] where the center of the top-left coordinate is [0.5, 0.5] | |
offset = ( | |
0.5 # Hpatches assumes that the center of the top-left pixel is at [0,0] | |
) | |
query_coords = ( | |
torch.stack( | |
( | |
wq * (query_coords[..., 0] + 1) / 2, | |
hq * (query_coords[..., 1] + 1) / 2, | |
), | |
axis=-1, | |
) | |
- offset | |
) | |
query_to_support = ( | |
torch.stack( | |
( | |
wsup * (query_to_support[..., 0] + 1) / 2, | |
hsup * (query_to_support[..., 1] + 1) / 2, | |
), | |
axis=-1, | |
) | |
- offset | |
) | |
return query_coords, query_to_support | |
def inside_image(self, x, w, h): | |
return torch.logical_and( | |
x[:, 0] < (w - 1), | |
torch.logical_and(x[:, 1] < (h - 1), (x > 0).prod(dim=-1)), | |
) | |
def benchmark(self, model): | |
use_cuda = torch.cuda.is_available() | |
device = torch.device("cuda:0" if use_cuda else "cpu") | |
aepes = [] | |
pcks = [] | |
for seq_idx, seq_name in tqdm( | |
enumerate(self.seq_names), total=len(self.seq_names) | |
): | |
if seq_name[0] == "i": | |
continue | |
im1_path = os.path.join(self.seqs_path, seq_name, "1.ppm") | |
im1 = Image.open(im1_path) | |
w1, h1 = im1.size | |
for im_idx in range(2, 7): | |
im2_path = os.path.join(self.seqs_path, seq_name, f"{im_idx}.ppm") | |
im2 = Image.open(im2_path) | |
w2, h2 = im2.size | |
matches, certainty = model.match(im2, im1, do_pred_in_og_res=True) | |
matches, certainty = matches.reshape(-1, 4), certainty.reshape(-1) | |
inv_homography = torch.from_numpy( | |
np.loadtxt( | |
os.path.join(self.seqs_path, seq_name, "H_1_" + str(im_idx)) | |
) | |
).to(device) | |
homography = torch.linalg.inv(inv_homography) | |
pos_a, pos_b = self.convert_coordinates( | |
matches[:, :2], matches[:, 2:], w2, h2, w1, h1 | |
) | |
pos_a, pos_b = pos_a.double(), pos_b.double() | |
pos_a_h = torch.cat( | |
[pos_a, torch.ones([pos_a.shape[0], 1], device=device)], dim=1 | |
) | |
pos_b_proj_h = (homography @ pos_a_h.t()).t() | |
pos_b_proj = pos_b_proj_h[:, :2] / pos_b_proj_h[:, 2:] | |
mask = self.inside_image(pos_b_proj, w1, h1) | |
residual = pos_b - pos_b_proj | |
dist = (residual**2).sum(dim=1).sqrt()[mask] | |
aepes.append(torch.mean(dist).item()) | |
pck1 = (dist < 1.0).float().mean().item() | |
pck3 = (dist < 3.0).float().mean().item() | |
pck5 = (dist < 5.0).float().mean().item() | |
pcks.append([pck1, pck3, pck5]) | |
m_pcks = np.mean(np.array(pcks), axis=0) | |
return { | |
"hp_pck1": m_pcks[0], | |
"hp_pck3": m_pcks[1], | |
"hp_pck5": m_pcks[2], | |
} | |