-
Notifications
You must be signed in to change notification settings - Fork 1
/
test.py
115 lines (90 loc) · 3.96 KB
/
test.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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
# predicted as a batch
from params import par
from model import DeepVO
import numpy as np
from PIL import Image
import glob
import os
import torch
from utils.data import get_data_info
from utils.util import pose_to_SE3
import torchvision.transforms.functional as TF
def test(path_list):
# Path
load_model_path = par.load_model_path # choose the model you want to load
save_dir = '{}/results/'.format(par.save_path) # directory to save prediction answer
if not os.path.exists(save_dir):
os.makedirs(save_dir)
device = torch.device('cuda:1')
# Load model
M_deepvo = DeepVO(par.img_h, par.img_w, par.batch_norm)
use_cuda = torch.cuda.is_available()
if use_cuda:
M_deepvo = M_deepvo.to(par.device)
M_deepvo.load_state_dict(torch.load(load_model_path))
else:
M_deepvo.load_state_dict(torch.load(load_model_path, map_location={'cpu'}))
print('Load model from: ', load_model_path)
M_deepvo.eval()
for test_video in path_list:
data = get_data_info(folder_list=[test_video], seq_len=10, drop_last=False, overlap=1)
image_arr = data['image_path'] # image paths
groundtruth_arr = data['pose']
prev = None
answer = [[1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0], ]
ang_err_list = []
trans_err_list = []
for i in range(len(data['image_path'])):
print('{} / {}'.format(i, len(data['image_path'])), end='\r', flush=True)
# Load the test images
image_path_sequence = image_arr[i]
image_sequence = []
for img_path in image_path_sequence:
img_as_img = Image.open(img_path)
img_as_img = TF.resize(img_as_img, size=(par.img_h, par.img_w))
img_as_tensor = TF.to_tensor(img_as_img)-0.5
img_as_tensor = img_as_tensor.unsqueeze(0)
image_sequence.append(img_as_tensor)
image_sequence = torch.cat(image_sequence, 0).to(par.device)
gt_sequence = groundtruth_arr[i][:, :6]
with torch.no_grad():
x_in = image_sequence.unsqueeze(0)
angle, trans, hc = M_deepvo.forward(x_in, prev=prev)
prev = hc
angle = angle.squeeze().detach().cpu().numpy()
trans = trans.squeeze().detach().cpu().numpy()
pose_pred = np.hstack((angle, trans))
# Record the estimation error
ang_err_list.append(np.mean((gt_sequence[:, :3] - angle)**2))
trans_err_list.append(np.mean((gt_sequence[:, 3:] - trans)**2))
# Accumulate the relative poses
for index in range(angle.shape[0]):
poses_pre = answer[-1]
poses_pre = np.array(poses_pre).reshape(3, 4)
R_pre = poses_pre[:, :3]
t_pre = poses_pre[:, 3]
pose_rel = pose_pred[index, :]
Rt_rel = pose_to_SE3(pose_rel)
R_rel = Rt_rel[:, :3]
t_rel = Rt_rel[:, 3]
R = R_pre @ R_rel
t = R_pre.dot(t_rel) + t_pre
pose = np.concatenate((R, t.reshape(3, 1)), 1).flatten().tolist()
answer.append(pose)
print('len(answer): ', len(answer))
print('expect len: ', len(glob.glob('{}{}/*.png'.format(par.image_dir, test_video))))
ang_err_m = np.mean(ang_err_list)
trans_err_m = np.mean(trans_err_list)
print(f'Average angle loss: {ang_err_m}')
print(f'Average translation loss: {trans_err_m}')
# Save answer
with open('{}/{}_pred.txt'.format(save_dir, test_video), 'w') as f:
for pose in answer:
if type(pose) == list:
f.write(' '.join([str(r) for r in pose]))
else:
f.write(str(pose))
f.write('\n')
if __name__ == '__main__':
test(par.test_video)
#test(par.train_video)