Spaces:
Running
Running
File size: 3,780 Bytes
c608946 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
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],
}
|