Spaces:
Running
Running
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']) |