-
Notifications
You must be signed in to change notification settings - Fork 25
/
Copy pathpose.py
97 lines (89 loc) · 3.48 KB
/
pose.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
import numpy as np
import torch
import easydict
def params_to_mtrx(opt,sim3):
R = get_lie_rotation_matrix(opt,sim3.rot)
mtrx = torch.cat([sim3.scale.exp()*R,sim3.trans[:,None]],dim=1)
return mtrx
def apply_3Dsim(opt,points,sim3_mtrx,inv=False):
sR,t = sim3_mtrx[:,:3],sim3_mtrx[:,3]
points = (points-t)@sR.inverse().t() if inv else \
points@sR.t()+t
return points
def get_lie_rotation_matrix(opt,r):
O = torch.tensor(0.0,device=opt.device,dtype=torch.float32)
rx = torch.stack([torch.stack([O,-r[2],r[1]]),
torch.stack([r[2],O,-r[0]]),
torch.stack([-r[1],r[0],O])],dim=0)
# matrix exponential
R = torch.eye(3,device=opt.device,dtype=torch.float32)
numer = torch.eye(3,device=opt.device,dtype=torch.float32)
denom = 1.0
for i in range(1,20):
numer = numer.matmul(rx)
denom *= i
R += numer/denom
return R
def add_noise(opt,cs_map_mtrx):
noise = easydict.EasyDict({
"scale": opt.noise*torch.randn(1,device=opt.device)[0],
"rot": opt.noise*torch.randn(3,device=opt.device),
"trans": opt.noise*torch.randn(3,device=opt.device),
})
noise_mtrx = params_to_mtrx(opt,noise)
hom = torch.cat([torch.zeros(1,3,device=opt.device),
torch.ones(1,1,device=opt.device)],dim=1)
cs_map_mtrx_hom = torch.cat([cs_map_mtrx,hom],dim=0)
noisy_cs_map_mtrx = noise_mtrx@cs_map_mtrx_hom
print("noise -- scale: {0:.4f}, rot: [{1:.4f},{2:.4f},{3:.4f}], trans: [{4:.4f},{5:.4f},{6:.4f}]"
.format(noise.scale,noise.rot[0],noise.rot[1],noise.rot[2],noise.trans[0],noise.trans[1],noise.trans[2]))
return noisy_cs_map_mtrx
def interpolate_camera(extr1,extr2,alpha=0.5): # [B,3,4]
R1,t1 = extr1[...,:3],extr1[...,3:]
R2,t2 = extr2[...,:3],extr2[...,3:]
q1 = rotation_matrix_to_quaternion(R1)
q2 = rotation_matrix_to_quaternion(R2)
q_intp = interpolate_quaternion(q1,q2,alpha)
R_intp = quaternion_to_rotation_matrix(q_intp)
t_intp = (1-alpha)*t1+alpha*t2
extr_intp = torch.cat([R_intp,t_intp],dim=-1)
return extr_intp
def interpolate_quaternion(q1,q2,alpha=0.5):
cos_angle = (q1*q2).sum(dim=-1)
flip = cos_angle<0
q1[flip] *= -1
cos_angle[flip] *= -1
theta = cos_angle.acos()[...,None]
slerp = (((1-alpha)*theta).sin()*q1+(alpha*theta).sin()*q2)/theta.sin()
return slerp
def rotation_matrix_to_quaternion(R): # [B,3,3]
row0,row1,row2 = torch.unbind(R,dim=-2)
R00,R01,R02 = torch.unbind(row0,dim=-1)
R10,R11,R12 = torch.unbind(row1,dim=-1)
R20,R21,R22 = torch.unbind(row2,dim=-1)
t = R[...,0,0]+R[...,1,1]+R[...,2,2]
r = (1+t).sqrt()
qa = 0.5*r
qb = (R21-R12).sign()*0.5*(1+R00-R11-R22).sqrt()
qc = (R02-R20).sign()*0.5*(1-R00+R11-R22).sqrt()
qd = (R10-R01).sign()*0.5*(1-R00-R11+R22).sqrt()
q = torch.stack([qa,qb,qc,qd],dim=-1)
for i,qi in enumerate(q):
if torch.isnan(qi).any():
print(i)
K = torch.stack([torch.stack([R00-R11-R22,R10+R01,R20+R02,R12-R21],dim=-1),
torch.stack([R10+R01,R11-R00-R22,R21+R12,R20-R20],dim=-1),
torch.stack([R20+R02,R21+R12,R22-R00-R11,R01-R10],dim=-1),
torch.stack([R12-R21,R20-R02,R01-R10,R00+R11+R22],dim=-1)],dim=-2)/3.0
K = K[i]
eigval,eigvec = K.eig(eigenvectors=True)
idx = eigval[:,0].argmax()
V = eigvec[:,idx]
q[i] = torch.stack([V[3],V[0],V[1],V[2]])
return q
def quaternion_to_rotation_matrix(q): # [B,4]
qa,qb,qc,qd = torch.unbind(q,dim=-1)
R = torch.stack([torch.stack([1-2*(qc**2+qd**2),2*(qb*qc-qa*qd),2*(qa*qc+qb*qd)],dim=-1),
torch.stack([2*(qb*qc+qa*qd),1-2*(qb**2+qd**2),2*(qc*qd-qa*qb)],dim=-1),
torch.stack([2*(qb*qd-qa*qc),2*(qa*qb+qc*qd),1-2*(qb**2+qc**2)],dim=-1)],dim=-2)
return R