File size: 1,395 Bytes
404d2af |
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 |
import os
import numpy as np
import h5py
import cv2
class standard_reader:
def __init__(self,config):
self.raw_dir=config['rawdata_dir']
self.dataset=h5py.File(config['dataset_dir'],'r')
self.num_kpt=config['num_kpt']
def run(self,index):
K1,K2=np.asarray(self.dataset['K1'][str(index)]),np.asarray(self.dataset['K2'][str(index)])
R = np.asarray(self.dataset['R'][str(index)])
t = np.asarray(self.dataset['T'][str(index)])
t = t / np.sqrt((t ** 2).sum())
desc1,desc2=self.dataset['desc1'][str(index)][()][:self.num_kpt],self.dataset['desc2'][str(index)][()][:self.num_kpt]
x1, x2 = self.dataset['kpt1'][str(index)][()][:self.num_kpt], self.dataset['kpt2'][str(index)][()][:self.num_kpt]
e,f=self.dataset['e'][str(index)][()],self.dataset['f'][str(index)][()]
img1_path,img2_path=self.dataset['img_path1'][str(index)][()][0].decode(),self.dataset['img_path2'][str(index)][()][0].decode()
img1,img2=cv2.imread(os.path.join(self.raw_dir,img1_path)),cv2.imread(os.path.join(self.raw_dir,img2_path))
info={'index':index,'K1':K1,'K2':K2,'R':R,'t':t,'x1':x1,'x2':x2,'desc1':desc1,'desc2':desc2,'img1':img1,'img2':img2,'e':e,'f':f,'r_gt':R,'t_gt':t}
return info
def close(self):
self.dataset.close()
def __len__(self):
return len(self.dataset['K1']) |