|
from abc import ABCMeta, abstractmethod |
|
import os |
|
import h5py |
|
import numpy as np |
|
from tqdm import trange |
|
from torch.multiprocessing import Pool,set_start_method |
|
set_start_method('spawn',force=True) |
|
|
|
import sys |
|
ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../")) |
|
sys.path.insert(0, ROOT_DIR) |
|
from components import load_component |
|
|
|
|
|
class BaseDumper(metaclass=ABCMeta): |
|
def __init__(self,config): |
|
self.config=config |
|
self.img_seq=[] |
|
self.dump_seq=[] |
|
|
|
@abstractmethod |
|
def get_seqs(self): |
|
raise NotImplementedError |
|
|
|
@abstractmethod |
|
def format_dump_folder(self): |
|
raise NotImplementedError |
|
|
|
@abstractmethod |
|
def format_dump_data(self): |
|
raise NotImplementedError |
|
|
|
def initialize(self): |
|
self.extractor=load_component('extractor',self.config['extractor']['name'],self.config['extractor']) |
|
self.get_seqs() |
|
self.format_dump_folder() |
|
|
|
|
|
def extract(self,index): |
|
img_path,dump_path=self.img_seq[index],self.dump_seq[index] |
|
if not self.config['extractor']['overwrite'] and os.path.exists(dump_path): |
|
return |
|
kp, desc = self.extractor.run(img_path) |
|
self.write_feature(kp,desc,dump_path) |
|
|
|
def dump_feature(self): |
|
print('Extrating features...') |
|
self.num_img=len(self.dump_seq) |
|
pool=Pool(self.config['extractor']['num_process']) |
|
iteration_num=self.num_img//self.config['extractor']['num_process'] |
|
if self.num_img%self.config['extractor']['num_process']!=0: |
|
iteration_num+=1 |
|
for index in trange(iteration_num): |
|
indicies_list=range(index*self.config['extractor']['num_process'],min((index+1)*self.config['extractor']['num_process'],self.num_img)) |
|
pool.map(self.extract,indicies_list) |
|
pool.close() |
|
pool.join() |
|
|
|
def write_feature(self,pts, desc, filename): |
|
with h5py.File(filename, "w") as ifp: |
|
ifp.create_dataset('keypoints', pts.shape, dtype=np.float32) |
|
ifp.create_dataset('descriptors', desc.shape, dtype=np.float32) |
|
ifp["keypoints"][:] = pts |
|
ifp["descriptors"][:] = desc |
|
|
|
def form_standard_dataset(self): |
|
dataset_path=os.path.join(self.config['dataset_dump_dir'],self.config['data_name']+\ |
|
'_'+self.config['extractor']['name']+'_'+str(self.config['extractor']['num_kpt'])+'.hdf5') |
|
|
|
pair_data_type=['K1','K2','R','T','e','f'] |
|
num_pairs=len(self.data['K1']) |
|
with h5py.File(dataset_path, 'w') as f: |
|
print('collecting pair info...') |
|
for type in pair_data_type: |
|
dg=f.create_group(type) |
|
for idx in range(num_pairs): |
|
data_item=np.asarray(self.data[type][idx]) |
|
dg.create_dataset(str(idx),data_item.shape,data_item.dtype,data=data_item) |
|
|
|
for type in ['img_path1','img_path2']: |
|
dg=f.create_group(type) |
|
for idx in range(num_pairs): |
|
dg.create_dataset(str(idx),[1],h5py.string_dtype(encoding='ascii'),data=self.data[type][idx].encode('ascii')) |
|
|
|
|
|
print('collecting desc and kpt...') |
|
desc1_g,desc2_g,kpt1_g,kpt2_g=f.create_group('desc1'),f.create_group('desc2'),f.create_group('kpt1'),f.create_group('kpt2') |
|
for idx in trange(num_pairs): |
|
desc_file1,desc_file2=h5py.File(self.data['fea_path1'][idx],'r'),h5py.File(self.data['fea_path2'][idx],'r') |
|
desc1,desc2,kpt1,kpt2=desc_file1['descriptors'][()],desc_file2['descriptors'][()],desc_file1['keypoints'][()],desc_file2['keypoints'][()] |
|
desc1_g.create_dataset(str(idx),desc1.shape,desc1.dtype,data=desc1) |
|
desc2_g.create_dataset(str(idx),desc2.shape,desc2.dtype,data=desc2) |
|
kpt1_g.create_dataset(str(idx),kpt1.shape,kpt1.dtype,data=kpt1) |
|
kpt2_g.create_dataset(str(idx),kpt2.shape,kpt2.dtype,data=kpt2) |
|
|
|
|
|
|
|
|
