|
import argparse |
|
import os |
|
import numpy as np |
|
import h5py |
|
import cv2 |
|
from numpy.core.numeric import indices |
|
import pyxis as px |
|
from tqdm import trange |
|
|
|
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 |
|
|
|
parser = argparse.ArgumentParser(description='checking training data.') |
|
parser.add_argument('--meta_dir', type=str, default='dataset/valid') |
|
parser.add_argument('--dataset_dir', type=str, default='dataset') |
|
parser.add_argument('--desc_dir', type=str, default='desc') |
|
parser.add_argument('--raw_dir', type=str, default='raw_data') |
|
parser.add_argument('--desc_suffix', type=str, default='_root_1000.hdf5') |
|
parser.add_argument('--vis_folder',type=str,default=None) |
|
args=parser.parse_args() |
|
|
|
|
|
|
|
if __name__=='__main__': |
|
if args.vis_folder is not None and not os.path.exists(args.vis_folder): |
|
os.mkdir(args.vis_folder) |
|
|
|
pair_num_list=np.loadtxt(os.path.join(args.meta_dir,'pair_num.txt'),dtype=str) |
|
pair_seq_list,accu_pair_list=train_utils.parse_pair_seq(pair_num_list) |
|
total_pair=int(pair_num_list[0,1]) |
|
total_inlier_rate,total_corr_num,total_incorr_num=[],[],[] |
|
pair_num_list=pair_num_list[1:] |
|
|
|
for index in trange(total_pair): |
|
seq=pair_seq_list[index] |
|
index_within_seq=index-accu_pair_list[seq] |
|
with h5py.File(os.path.join(args.dataset_dir,seq,'info.h5py'),'r') as data: |
|
corr=data['corr'][str(index_within_seq)][()] |
|
corr1,corr2=corr[:,0],corr[:,1] |
|
incorr1,incorr2=data['incorr1'][str(index_within_seq)][()],data['incorr2'][str(index_within_seq)][()] |
|
img_path1,img_path2=data['img_path1'][str(index_within_seq)][()][0].decode(),data['img_path2'][str(index_within_seq)][()][0].decode() |
|
img_name1,img_name2=img_path1.split('/')[-1],img_path2.split('/')[-1] |
|
fea_path1,fea_path2=os.path.join(args.desc_dir,seq,img_name1+args.desc_suffix),os.path.join(args.desc_dir,seq,img_name2+args.desc_suffix) |
|
with h5py.File(fea_path1,'r') as fea1, h5py.File(fea_path2,'r') as fea2: |
|
desc1,kpt1=fea1['descriptors'][()],fea1['keypoints'][()][:,:2] |
|
desc2,kpt2=fea2['descriptors'][()],fea2['keypoints'][()][:,:2] |
|
sim_mat=desc1@desc2.T |
|
nn_index1,nn_index2=np.argmax(sim_mat,axis=1),np.argmax(sim_mat,axis=0) |
|
mask_mutual=(nn_index2[nn_index1]==np.arange(len(nn_index1)))[corr1] |
|
mask_inlier=nn_index1[corr1]==corr2 |
|
mask_nn_correct=np.logical_and(mask_mutual,mask_inlier) |
|
|
|
total_inlier_rate.append(mask_nn_correct.mean()) |
|
total_corr_num.append(len(corr1)) |
|
total_incorr_num.append((len(incorr1)+len(incorr2))/2) |
|
|
|
if args.vis_folder is not None: |
|
|
|
img1,img2=cv2.imread(os.path.join(args.raw_dir,img_path1)),cv2.imread(os.path.join(args.raw_dir,img_path2)) |
|
corr1_pos,corr2_pos=np.take_along_axis(kpt1,corr1[:,np.newaxis],axis=0),np.take_along_axis(kpt2,corr2[:,np.newaxis],axis=0) |
|
dis_corr=evaluation_utils.draw_match(img1,img2,corr1_pos,corr2_pos) |
|
cv2.imwrite(os.path.join(args.vis_folder,str(index)+'.png'),dis_corr) |
|
|
|
incorr1_pos,incorr2_pos=np.take_along_axis(kpt1,incorr1[:,np.newaxis],axis=0),np.take_along_axis(kpt2,incorr2[:,np.newaxis],axis=0) |
|
dis_incorr1,dis_incorr2=evaluation_utils.draw_points(img1,incorr1_pos),evaluation_utils.draw_points(img2,incorr2_pos) |
|
cv2.imwrite(os.path.join(args.vis_folder,str(index)+'_incorr1.png'),dis_incorr1) |
|
cv2.imwrite(os.path.join(args.vis_folder,str(index)+'_incorr2.png'),dis_incorr2) |
|
|
|
print('NN matching accuracy: ',np.asarray(total_inlier_rate).mean()) |
|
print('mean corr number: ',np.asarray(total_corr_num).mean()) |
|
print('mean incorr number: ',np.asarray(total_incorr_num).mean()) |
|
|
