-
Notifications
You must be signed in to change notification settings - Fork 9
/
main.py
52 lines (37 loc) · 1.41 KB
/
main.py
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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import mayavi.mlab as mlab
from model.geometry import *
import os
import torch
import argparse
import cv2
from model.registration import Registration
import yaml
from easydict import EasyDict as edict
def join(loader, node):
seq = loader.construct_sequence(node)
return '_'.join([str(i) for i in seq])
yaml.add_constructor('!join', join)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('config', type=str, help= 'Path to the config file.')
args = parser.parse_args()
with open(args.config,'r') as f:
config = yaml.load(f, Loader=yaml.Loader)
config = edict(config)
if config.gpu_mode:
config.device = torch.device("cuda:0")
else:
config.device = torch.device('cpu')
"""demo data"""
intrinsics = np.loadtxt(config.intrinsics)
"""load lepard predicted matches as landmarks"""
data = np.load(config.correspondence)
ldmk_src = data['src_pcd'][0][data['match'][:,1] ]
ldmk_tgt = data['tgt_pcd'][0][data['match'][:,2] ]
uv_src = xyz_2_uv(ldmk_src, intrinsics)
uv_tgt = xyz_2_uv(ldmk_tgt, intrinsics)
landmarks = ( torch.from_numpy(uv_src).to(config.device),
torch.from_numpy(uv_tgt).to(config.device))
"""init model with source frame"""
model = Registration(config.src_depth, K=intrinsics, config=config)
model.register_a_depth_frame( config.tgt_depth, landmarks=landmarks)