Spaces:
Running
Running
File size: 3,448 Bytes
a80d6bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import torch
import numpy as np
import cv2
import os
from loss import batch_episym
from tqdm import tqdm
import sys
ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, ROOT_DIR)
from utils import evaluation_utils,train_utils
def valid(valid_loader, model,match_loss, config,model_config):
model.eval()
loader_iter = iter(valid_loader)
num_pair = 0
total_loss,total_acc_corr,total_acc_incorr=0,0,0
total_precision,total_recall=torch.zeros(model_config.layer_num ,device='cuda'),\
torch.zeros(model_config.layer_num ,device='cuda')
total_acc_mid=torch.zeros(len(model_config.seedlayer)-1,device='cuda')
with torch.no_grad():
if config.local_rank==0:
loader_iter=tqdm(loader_iter)
print('validating...')
for test_data in loader_iter:
num_pair+= 1
test_data = train_utils.tocuda(test_data)
res= model(test_data)
loss_res=match_loss.run(test_data,res)
total_acc_corr+=loss_res['acc_corr']
total_acc_incorr+=loss_res['acc_incorr']
total_loss+=loss_res['total_loss']
if config.model_name=='SGM':
total_acc_mid+=loss_res['mid_acc_corr']
total_precision,total_recall=total_precision+loss_res['pre_seed_conf'],total_recall+loss_res['recall_seed_conf']
total_acc_corr/=num_pair
total_acc_incorr /= num_pair
total_precision/=num_pair
total_recall/=num_pair
total_acc_mid/=num_pair
#apply tensor reduction
total_loss,total_acc_corr,total_acc_incorr,total_precision,total_recall,total_acc_mid=train_utils.reduce_tensor(total_loss,'sum'),\
train_utils.reduce_tensor(total_acc_corr,'mean'),train_utils.reduce_tensor(total_acc_incorr,'mean'),\
train_utils.reduce_tensor(total_precision,'mean'),train_utils.reduce_tensor(total_recall,'mean'),train_utils.reduce_tensor(total_acc_mid,'mean')
model.train()
return total_loss,total_acc_corr,total_acc_incorr,total_precision,total_recall,total_acc_mid
def dump_train_vis(res,data,step,config):
#batch matching
p=res['p'][:,:-1,:-1]
score,index1=torch.max(p,dim=-1)
_,index2=torch.max(p,dim=-2)
mask_th=score>0.2
mask_mc=index2.gather(index=index1,dim=1) == torch.arange(len(p[0])).cuda()[None]
mask_p=mask_th&mask_mc#B*N
corr1,corr2=data['x1'],data['x2'].gather(index=index1[:,:,None].expand(-1,-1,2),dim=1)
corr1_kpt,corr2_kpt=data['kpt1'],data['kpt2'].gather(index=index1[:,:,None].expand(-1,-1,2),dim=1)
epi_dis=batch_episym(corr1,corr2,data['e_gt'])
mask_inlier=epi_dis<config.inlier_th#B*N
#dump vis
for cur_mask_p,cur_mask_inlier,cur_corr1,cur_corr2,img_path1,img_path2 in zip(mask_p,mask_inlier,corr1_kpt,corr2_kpt,data['img_path1'],data['img_path2']):
img1,img2=cv2.imread(img_path1),cv2.imread(img_path2)
dis_play=evaluation_utils.draw_match(img1,img2,cur_corr1[cur_mask_p].cpu().numpy(),cur_corr2[cur_mask_p].cpu().numpy(),inlier=cur_mask_inlier)
base_name_seq=os.path.join(img_path1.split('/')[-1]+'_'+img_path2.split('/')[-1]+'_'+img_path1.split('/')[-2])
save_path=os.path.join(config.train_vis_folder,'train_vis',config.log_base,str(step),base_name_seq+'.png')
cv2.imwrite(save_path,dis_play) |