-
Notifications
You must be signed in to change notification settings - Fork 17
/
test.py
137 lines (112 loc) · 4.9 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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
#%%
import tqdm
from dataset import PnP_Data_Simulator
import torch
import cv2
import os
import numpy as np
from dataset import *
from model import *
from train import quaternion2rotation
def pose_err_3d(kpts3d, r, t, gr, gt):
t = t.reshape(-1, 1)
gt = gt.reshape(-1, 1)
# error metric in 3D space
res1 = np.matmul(r, kpts3d.T) + t
res2 = np.matmul(gr, kpts3d.T)+ gt
diff = np.linalg.norm(res1-res2, axis=0)
return diff.mean()
def pose_err_2d(kpts3d, r, t, gr, gt, K):
t = t.reshape(-1, 1)
gt = gt.reshape(-1, 1)
# error metric in 3D space
res1 = np.matmul(K, np.matmul(r, kpts3d.T) + t)
x1 = res1[0]/res1[2]
y1 = res1[1]/res1[2]
res2 = np.matmul(K, np.matmul(gr, kpts3d.T)+ gt)
x2 = res2[0]/res2[2]
y2 = res2[1]/res2[2]
diff = np.sqrt((x1-x2)*(x1-x2)+(y1-y2)*(y1-y2))
return diff.mean()
def PnP_OpenCV_EPNP(p3d, p2d, intrinsics):
ptCnt = len(p3d)
retval, rot, trans, inliers = cv2.solvePnPRansac(p3d, p2d, intrinsics, None, flags=cv2.SOLVEPNP_EPNP)
# retval, rot, trans, inliers = cv2.solvePnPRansac(p3d, p2d, intrinsics, None, flags=cv2.SOLVEPNP_P3P)
if not retval:
print('PnP Failed')
R = np.eye(3)
T = np.array([0,0,1]).reshape(-1, 1)
return R, T
R = cv2.Rodrigues(rot)[0] # convert to rotation matrix
T = trans.reshape(-1, 1)
return R, T
def PnP_Learning(p3d, sxy, dxy, intrinsics, width, height, translation_min, translation_max, model=None):
p3d = torch.from_numpy(p3d).float().unsqueeze(dim=0)
sxy = torch.from_numpy(sxy).float().unsqueeze(dim=0)
dxy = torch.from_numpy(dxy).float().unsqueeze(dim=0)
sxy[..., 0] = sxy[..., 0] - (width/2)
sxy[..., 1] = sxy[..., 1] - (height/2)
sxy[..., 0] = sxy[..., 0] / width
sxy[..., 1] = sxy[..., 1] / height
# dxy = dxy / ( 2* dxy.norm(dim=-1).unsqueeze(-1) )
dxy[..., 0] = dxy[..., 0] / width
dxy[..., 1] = dxy[..., 1] / height
inData = torch.cat((sxy, dxy), dim=-1)
# inData = torch.cat((p3d, inData), dim=-1)
inData = inData.transpose(1,2)
out = model(inData)
predQ = out[:, :4]
predT = out[:, 4:]
predR = quaternion2rotation(predQ)
# recover predicted translation
minT = translation_min.type_as(out).view(-1,3)
maxT = translation_max.type_as(out).view(-1,3)
predT = (predT + 0.5) * (maxT - minT) + minT
return predR[0].detach().numpy(), predT[0].detach().numpy().reshape(-1,1)
if __name__ == "__main__":
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
batch_size = 32
model = SimplePnPNet(nIn=4)
model.eval()
model.load_state_dict(torch.load('./xydxdy.pth'))
min_noise = 0
max_noise = 31
print("Noise\tEPnP(2d)\tOurs(2d)\tEPnP(3d)\tOurs(3d)")
#
for noise in range(min_noise, max_noise):
dataset = PnP_Data_Simulator(sampleCnt=2000, minNoiseSigma=noise, maxNoiseSigma=noise, minOutlier=0.1, maxOutlier=0.1)
data_loader = torch.utils.data.DataLoader(dataset, shuffle=True, batch_size=batch_size, num_workers=4)
translation_min = torch.FloatTensor(dataset.translation_min)
translation_max = torch.FloatTensor(dataset.translation_max)
# pbar = tqdm.tqdm(enumerate(data_loader), total=len(data_loader)) # progress bar
err_3d_stat_opencv_epnp = []
err_3d_stat_opencv_ours = []
#
err_2d_stat_opencv_epnp = []
err_2d_stat_opencv_ours = []
# for batch_idx, batch_data in pbar:
for batch_idx, batch_data in enumerate(data_loader):
intrinsic, quat, trans, sxy, dxy, p3ds = batch_data
p3ds = p3ds.repeat(1,dataset.gridCnt,1)
Rs = quaternion2rotation(quat)
Ts = trans
count = len(Rs)
for i in range(count):
k = intrinsic[i].numpy()
gtR = Rs[i].numpy()
gtT = Ts[i].numpy().reshape(-1, 1)
p3d = p3ds[i].numpy()
p2d = (sxy + dxy)[i].numpy()
ptCnt = len(p3d)
R1, T1 = PnP_OpenCV_EPNP(p3d, p2d, k)
err = pose_err_3d(p3ds[i].numpy(), R1, T1, gtR, gtT)
err_3d_stat_opencv_epnp.append(err)
err = pose_err_2d(p3ds[i].numpy(), R1, T1, gtR, gtT, k)
err_2d_stat_opencv_epnp.append(err)
R2, T2 = PnP_Learning(p3ds[i].numpy(), sxy[i].numpy(), dxy[i].numpy(), k, dataset.width, dataset.height, \
translation_min, translation_max, model=model)
err = pose_err_3d(p3ds[i].numpy(), R2, T2, gtR, gtT)
err_3d_stat_opencv_ours.append(err)
err = pose_err_2d(p3ds[i].numpy(), R2, T2, gtR, gtT, k)
err_2d_stat_opencv_ours.append(err)
print("%d\t%.3f\t\t%.3f\t\t%.3f\t\t%.3f" % (noise, np.array(err_2d_stat_opencv_epnp).mean(), np.array(err_2d_stat_opencv_ours).mean(), np.array(err_3d_stat_opencv_epnp).mean(), np.array(err_3d_stat_opencv_ours).mean()))