Spaces:
Running
Running
| 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 | |