Spaces:
Running
Running
| import torch | |
| import numpy as np | |
| import tqdm | |
| from roma.datasets import MegadepthBuilder | |
| from roma.utils import warp_kpts | |
| from torch.utils.data import ConcatDataset | |
| import roma | |
| class MegadepthDenseBenchmark: | |
| def __init__(self, data_root="data/megadepth", h = 384, w = 512, num_samples = 2000) -> 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 | |
| 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() | |
| return gd, pck_1, pck_3, pck_5, prob | |
| 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 | |
| ) | |
| B = batch_size | |
| dataloader = torch.utils.data.DataLoader( | |
| self.dataset, batch_size=B, num_workers=batch_size, sampler=sampler | |
| ) | |
| for idx, data in tqdm.tqdm(enumerate(dataloader), disable = roma.RANK > 0): | |
| im_A, im_B, depth1, depth2, T_1to2, K1, K2 = ( | |
| data["im_A"], | |
| data["im_B"], | |
| data["im_A_depth"].cuda(), | |
| data["im_B_depth"].cuda(), | |
| data["T_1to2"].cuda(), | |
| data["K1"].cuda(), | |
| data["K2"].cuda(), | |
| ) | |
| matches, certainty = model.match(im_A, im_B, batched=True) | |
| gd, pck_1, pck_3, pck_5, prob = self.geometric_dist( | |
| depth1, depth2, T_1to2, K1, K2, matches | |
| ) | |
| if roma.DEBUG_MODE: | |
| from roma.utils.utils import tensor_to_pil | |
| import torch.nn.functional as F | |
| path = "vis" | |
| H, W = model.get_output_resolution() | |
| white_im = torch.ones((B,1,H,W),device="cuda") | |
| im_B_transfer_rgb = F.grid_sample( | |
| im_B.cuda(), matches[:,:,:W, 2:], mode="bilinear", align_corners=False | |
| ) | |
| warp_im = im_B_transfer_rgb | |
| c_b = certainty[:,None]#(certainty*0.9 + 0.1*torch.ones_like(certainty))[:,None] | |
| vis_im = c_b * warp_im + (1 - c_b) * white_im | |
| for b in range(B): | |
| import os | |
| os.makedirs(f"{path}/{model.name}/{idx}_{b}_{H}_{W}",exist_ok=True) | |
| tensor_to_pil(vis_im[b], unnormalize=True).save( | |
| f"{path}/{model.name}/{idx}_{b}_{H}_{W}/warp.jpg") | |
| tensor_to_pil(im_A[b].cuda(), unnormalize=True).save( | |
| f"{path}/{model.name}/{idx}_{b}_{H}_{W}/im_A.jpg") | |
| tensor_to_pil(im_B[b].cuda(), unnormalize=True).save( | |
| f"{path}/{model.name}/{idx}_{b}_{H}_{W}/im_B.jpg") | |
| gd_tot, pck_1_tot, pck_3_tot, pck_5_tot = ( | |
| gd_tot + gd.mean(), | |
| pck_1_tot + pck_1, | |
| pck_3_tot + pck_3, | |
| pck_5_tot + pck_5, | |
| ) | |
| return { | |
| "epe": gd_tot.item() / len(dataloader), | |
| "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), | |
| } | |