import torch import numpy as np import tqdm from dkm.datasets import MegadepthBuilder from dkm.utils import warp_kpts from torch.utils.data import ConcatDataset class MegadepthDenseBenchmark: def __init__(self, data_root="data/megadepth", h = 384, w = 512, num_samples = 2000, device=None) -> None: mega = MegadepthBuilder(data_root=data_root) self.dataset = ConcatDataset( mega.build_scenes(split="test_loftr", ht=h, wt=w) ) # fixed resolution of 384,512 self.num_samples = num_samples if device is None: device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.device = device def geometric_dist(self, depth1, depth2, T_1to2, K1, K2, dense_matches): b, h1, w1, d = dense_matches.shape with torch.no_grad(): x1 = dense_matches[..., :2].reshape(b, h1 * w1, 2) # x1 = torch.stack((2*x1[...,0]/w1-1,2*x1[...,1]/h1-1),dim=-1) mask, x2 = warp_kpts( x1.double(), depth1.double(), depth2.double(), T_1to2.double(), K1.double(), K2.double(), ) x2 = torch.stack( (w1 * (x2[..., 0] + 1) / 2, h1 * (x2[..., 1] + 1) / 2), dim=-1 ) prob = mask.float().reshape(b, h1, w1) x2_hat = dense_matches[..., 2:] x2_hat = torch.stack( (w1 * (x2_hat[..., 0] + 1) / 2, h1 * (x2_hat[..., 1] + 1) / 2), dim=-1 ) gd = (x2_hat - x2.reshape(b, h1, w1, 2)).norm(dim=-1) gd = gd[prob == 1] pck_1 = (gd < 1.0).float().mean() pck_3 = (gd < 3.0).float().mean() pck_5 = (gd < 5.0).float().mean() gd = gd.mean() return gd, pck_1, pck_3, pck_5 def benchmark(self, model, batch_size=8): model.train(False) with torch.no_grad(): gd_tot = 0.0 pck_1_tot = 0.0 pck_3_tot = 0.0 pck_5_tot = 0.0 sampler = torch.utils.data.WeightedRandomSampler( torch.ones(len(self.dataset)), replacement=False, num_samples=self.num_samples ) dataloader = torch.utils.data.DataLoader( self.dataset, batch_size=8, num_workers=batch_size, sampler=sampler ) for data in tqdm.tqdm(dataloader): im1, im2, depth1, depth2, T_1to2, K1, K2 = ( data["query"], data["support"], data["query_depth"].to(self.device), data["support_depth"].to(self.device), data["T_1to2"].to(self.device), data["K1"].to(self.device), data["K2"].to(self.device), ) matches, certainty = model.match(im1, im2, batched=True) gd, pck_1, pck_3, pck_5 = self.geometric_dist( depth1, depth2, T_1to2, K1, K2, matches ) gd_tot, pck_1_tot, pck_3_tot, pck_5_tot = ( gd_tot + gd, pck_1_tot + pck_1, pck_3_tot + pck_3, pck_5_tot + pck_5, ) return { "mega_pck_1": pck_1_tot.item() / len(dataloader), "mega_pck_3": pck_3_tot.item() / len(dataloader), "mega_pck_5": pck_5_tot.item() / len(dataloader), }