|
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) |
|
) |
|
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) |
|
|
|
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), |
|
} |
|
|
