Spaces:
Running
Running
| import os | |
| import glob | |
| import pickle | |
| from tqdm import trange | |
| import numpy as np | |
| import h5py | |
| from numpy.core.fromnumeric import reshape | |
| from .base_dumper import BaseDumper | |
| import sys | |
| ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../")) | |
| sys.path.insert(0, ROOT_DIR) | |
| import utils | |
| class fmbench(BaseDumper): | |
| def get_seqs(self): | |
| data_dir=os.path.join(self.config['rawdata_dir']) | |
| self.split_list=[] | |
| for seq in self.config['data_seq']: | |
| cur_split_list=np.unique(np.loadtxt(os.path.join(data_dir,seq,'pairs_which_dataset.txt'),dtype=str)) | |
| self.split_list.append(cur_split_list) | |
| for split in cur_split_list: | |
| split_dir=os.path.join(data_dir,seq,split) | |
| dump_dir=os.path.join(self.config['feature_dump_dir'],seq,split) | |
| cur_img_seq=glob.glob(os.path.join(split_dir,'Images','*.jpg')) | |
| cur_dump_seq=[os.path.join(dump_dir,path.split('/')[-1])+'_'+self.config['extractor']['name']+'_'+str(self.config['extractor']['num_kpt'])\ | |
| +'.hdf5' for path in cur_img_seq] | |
| self.img_seq+=cur_img_seq | |
| self.dump_seq+=cur_dump_seq | |
| def format_dump_folder(self): | |
| if not os.path.exists(self.config['feature_dump_dir']): | |
| os.mkdir(self.config['feature_dump_dir']) | |
| for seq_index in range(len(self.config['data_seq'])): | |
| seq_dir=os.path.join(self.config['feature_dump_dir'],self.config['data_seq'][seq_index]) | |
| if not os.path.exists(seq_dir): | |
| os.mkdir(seq_dir) | |
| for split in self.split_list[seq_index]: | |
| split_dir=os.path.join(seq_dir,split) | |
| if not os.path.exists(split_dir): | |
| os.mkdir(split_dir) | |
| def format_dump_data(self): | |
| print('Formatting data...') | |
| self.data={'K1':[],'K2':[],'R':[],'T':[],'e':[],'f':[],'fea_path1':[],'fea_path2':[],'img_path1':[],'img_path2':[]} | |
| for seq_index in range(len(self.config['data_seq'])): | |
| seq=self.config['data_seq'][seq_index] | |
| print(seq) | |
| pair_list=np.loadtxt(os.path.join(self.config['rawdata_dir'],seq,'pairs_with_gt.txt'),dtype=float) | |
| which_split_list=np.loadtxt(os.path.join(self.config['rawdata_dir'],seq,'pairs_which_dataset.txt'),dtype=str) | |
| for pair_index in trange(len(pair_list)): | |
| cur_pair=pair_list[pair_index] | |
| cur_split=which_split_list[pair_index] | |
| index1,index2=int(cur_pair[0]),int(cur_pair[1]) | |
| #get intrinsic | |
| camera=np.loadtxt(os.path.join(self.config['rawdata_dir'],seq,cur_split,'Camera.txt'),dtype=float) | |
| K1,K2=camera[index1].reshape([3,3]),camera[index2].reshape([3,3]) | |
| #get pose | |
| pose=np.loadtxt(os.path.join(self.config['rawdata_dir'],seq,cur_split,'Poses.txt'),dtype=float) | |
| pose1,pose2=pose[index1].reshape([3,4]),pose[index2].reshape([3,4]) | |
| R1,R2,t1,t2=pose1[:3,:3],pose2[:3,:3],pose1[:3,3][:,np.newaxis],pose2[:3,3][:,np.newaxis] | |
| dR = np.dot(R2, R1.T) | |
| dt = t2 - np.dot(dR, t1) | |
| dt /= np.sqrt(np.sum(dt**2)) | |
| e_gt_unnorm = np.reshape(np.matmul( | |
| np.reshape(utils.evaluation_utils.np_skew_symmetric(dt.astype('float64').reshape(1, 3)), (3, 3)), | |
| np.reshape(dR.astype('float64'), (3, 3))), (3, 3)) | |
| e_gt = e_gt_unnorm / np.linalg.norm(e_gt_unnorm) | |
| f=cur_pair[2:].reshape([3,3]) | |
| f_gt=f / np.linalg.norm(f) | |
| self.data['K1'].append(K1),self.data['K2'].append(K2) | |
| self.data['R'].append(dR),self.data['T'].append(dt) | |
| self.data['e'].append(e_gt),self.data['f'].append(f_gt) | |
| img_path1,img_path2=os.path.join(seq,cur_split,'Images',str(index1).zfill(8)+'.jpg'),\ | |
| os.path.join(seq,cur_split,'Images',str(index1).zfill(8)+'.jpg') | |
| fea_path1,fea_path2=os.path.join(self.config['feature_dump_dir'],seq,cur_split,str(index1).zfill(8)+'.jpg'+'_'+self.config['extractor']['name'] | |
| +'_'+str(self.config['extractor']['num_kpt'])+'.hdf5'),\ | |
| os.path.join(self.config['feature_dump_dir'],seq,cur_split,str(index2).zfill(8)+'.jpg'+'_'+self.config['extractor']['name'] | |
| +'_'+str(self.config['extractor']['num_kpt'])+'.hdf5') | |
| self.data['img_path1'].append(img_path1),self.data['img_path2'].append(img_path2) | |
| self.data['fea_path1'].append(fea_path1),self.data['fea_path2'].append(fea_path2) | |
| self.form_standard_dataset() | |