Spaces:
Sleeping
Sleeping
| import torch | |
| import numpy as np | |
| import os | |
| from collections import OrderedDict, namedtuple | |
| import sys | |
| ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) | |
| sys.path.insert(0, ROOT_DIR) | |
| from sgmnet import matcher as SGM_Model | |
| from superglue import matcher as SG_Model | |
| from utils import evaluation_utils | |
| class GNN_Matcher(object): | |
| def __init__(self, config, model_name): | |
| assert model_name == "SGM" or model_name == "SG" | |
| config = namedtuple("config", config.keys())(*config.values()) | |
| self.p_th = config.p_th | |
| self.model = SGM_Model(config) if model_name == "SGM" else SG_Model(config) | |
| self.model.cuda(), self.model.eval() | |
| checkpoint = torch.load(os.path.join(config.model_dir, "model_best.pth")) | |
| # for ddp model | |
| if list(checkpoint["state_dict"].items())[0][0].split(".")[0] == "module": | |
| new_stat_dict = OrderedDict() | |
| for key, value in checkpoint["state_dict"].items(): | |
| new_stat_dict[key[7:]] = value | |
| checkpoint["state_dict"] = new_stat_dict | |
| self.model.load_state_dict(checkpoint["state_dict"]) | |
| def run(self, test_data): | |
| norm_x1, norm_x2 = evaluation_utils.normalize_size( | |
| test_data["x1"][:, :2], test_data["size1"] | |
| ), evaluation_utils.normalize_size(test_data["x2"][:, :2], test_data["size2"]) | |
| x1, x2 = np.concatenate( | |
| [norm_x1, test_data["x1"][:, 2, np.newaxis]], axis=-1 | |
| ), np.concatenate([norm_x2, test_data["x2"][:, 2, np.newaxis]], axis=-1) | |
| feed_data = { | |
| "x1": torch.from_numpy(x1[np.newaxis]).cuda().float(), | |
| "x2": torch.from_numpy(x2[np.newaxis]).cuda().float(), | |
| "desc1": torch.from_numpy(test_data["desc1"][np.newaxis]).cuda().float(), | |
| "desc2": torch.from_numpy(test_data["desc2"][np.newaxis]).cuda().float(), | |
| } | |
| with torch.no_grad(): | |
| res = self.model(feed_data, test_mode=True) | |
| p = res["p"] | |
| index1, index2 = self.match_p(p[0, :-1, :-1]) | |
| corr1, corr2 = ( | |
| test_data["x1"][:, :2][index1.cpu()], | |
| test_data["x2"][:, :2][index2.cpu()], | |
| ) | |
| if len(corr1.shape) == 1: | |
| corr1, corr2 = corr1[np.newaxis], corr2[np.newaxis] | |
| return corr1, corr2 | |
| def match_p(self, p): # p N*M | |
| score, index = torch.topk(p, k=1, dim=-1) | |
| _, index2 = torch.topk(p, k=1, dim=-2) | |
| mask_th, index, index2 = score[:, 0] > self.p_th, index[:, 0], index2.squeeze(0) | |
| mask_mc = index2[index] == torch.arange(len(p)).cuda() | |
| mask = mask_th & mask_mc | |
| index1, index2 = torch.nonzero(mask).squeeze(1), index[mask] | |
| return index1, index2 | |
| class NN_Matcher(object): | |
| def __init__(self, config): | |
| config = namedtuple("config", config.keys())(*config.values()) | |
| self.mutual_check = config.mutual_check | |
| self.ratio_th = config.ratio_th | |
| def run(self, test_data): | |
| desc1, desc2, x1, x2 = ( | |
| test_data["desc1"], | |
| test_data["desc2"], | |
| test_data["x1"], | |
| test_data["x2"], | |
| ) | |
| desc_mat = np.sqrt( | |
| abs( | |
| (desc1**2).sum(-1)[:, np.newaxis] | |
| + (desc2**2).sum(-1)[np.newaxis] | |
| - 2 * desc1 @ desc2.T | |
| ) | |
| ) | |
| nn_index = np.argpartition(desc_mat, kth=(1, 2), axis=-1) | |
| dis_value12 = np.take_along_axis(desc_mat, nn_index, axis=-1) | |
| ratio_score = dis_value12[:, 0] / dis_value12[:, 1] | |
| nn_index1 = nn_index[:, 0] | |
| nn_index2 = np.argmin(desc_mat, axis=0) | |
| mask_ratio, mask_mutual = ( | |
| ratio_score < self.ratio_th, | |
| np.arange(len(x1)) == nn_index2[nn_index1], | |
| ) | |
| corr1, corr2 = x1[:, :2], x2[:, :2][nn_index1] | |
| if self.mutual_check: | |
| mask = mask_ratio & mask_mutual | |
| else: | |
| mask = mask_ratio | |
| corr1, corr2 = corr1[mask], corr2[mask] | |
| return corr1, corr2 | |