-
Notifications
You must be signed in to change notification settings - Fork 1
/
geometric_utils.py
112 lines (100 loc) · 3.45 KB
/
geometric_utils.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
import numpy as np
import torch
def numpy_quaternion_to_rot_mat(q):
q /= np.linalg.norm(q)
qq = 2 * q[:3].reshape(1, -1) * q[:3].reshape(-1, 1)
qq[0, 1] -= 2 * q[3] * q[2]
qq[0, 2] += 2 * q[3] * q[1]
qq[1, 2] -= 2 * q[3] * q[0]
qq[2, 1] += 2 * q[3] * q[0]
qq[2, 0] -= 2 * q[3] * q[1]
qq[1, 0] += 2 * q[3] * q[2]
q2 = q ** 2
qq[0, 0] = 1 - 2 * (q2[1] + q2[2])
qq[1, 1] = 1 - 2 * (q2[0] + q2[2])
qq[2, 2] = 1 - 2 * (q2[0] + q2[1])
return qq
def torch_batch_quaternion_to_rot_mat(q_unnorm):
# q is (-1, 4)
q = q_unnorm / torch.norm(q_unnorm, dim=1).view(-1, 1)
qq = torch.empty((q.shape[0], 3, 3), dtype=q.dtype, device=q.device)
qq[:, 0, 1] = 2 * (q[:, 0] * q[:, 1] - q[:, 3] * q[:, 2])
qq[:, 0, 2] = 2 * (q[:, 0] * q[:, 2] + q[:, 3] * q[:, 1])
qq[:, 1, 2] = 2 * (q[:, 1] * q[:, 2] - q[:, 3] * q[:, 0])
qq[:, 2, 1] = 2 * (q[:, 1] * q[:, 2] + q[:, 3] * q[:, 0])
qq[:, 2, 0] = 2 * (q[:, 0] * q[:, 2] - q[:, 3] * q[:, 1])
qq[:, 1, 0] = 2 * (q[:, 0] * q[:, 1] + q[:, 3] * q[:, 2])
q2 = q[:, :3] ** 2
qq[:, 0, 0] = 1 - 2 * (q2[:, 1] + q2[:, 2])
qq[:, 1, 1] = 1 - 2 * (q2[:, 0] + q2[:, 2])
qq[:, 2, 2] = 1 - 2 * (q2[:, 0] + q2[:, 1])
return qq
def numpy_euler_to_rotation(angles):
# note angles are z,y,x intrinsic
ax = angles[0]
cx = np.cos(ax)
sx = np.sin(ax)
ay = angles[1]
cy = np.cos(ay)
sy = np.sin(ay)
az = angles[2]
cz = np.cos(az)
sz = np.sin(az)
Ax = np.array([[1, 0, 0], [0, cx, -sx], [0, sx, cx]])
Ay = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]])
Az = np.array([[cz, -sz, 0], [sz, cz, 0], [0, 0, 1]])
R = np.matmul(Az, np.matmul(Ay, Ax))
return R
def numpy_quaternion_to_angles(q):
# note angles are z,y,x intrinsic
q /= np.linalg.norm(q)
roll = np.arctan2(2 * (q[3] * q[0] + q[1] * q[2]), 1 - 2 * (q[0] ** 2 + q[1] ** 2))
pitch = np.arcsin(2 * (q[3] * q[1] - q[0] * q[2]))
yaw = np.arctan2(2 * (q[3] * q[2] + q[0] * q[1]), 1 - 2 * (q[1] ** 2 + q[2] ** 2))
return np.array([roll, pitch, yaw])
def torch_batch_euler_to_rotation(angles):
# note angles are z,y,x intrinsic
ax = angles[:, 0]
cx = torch.cos(ax)
sx = torch.sin(ax)
ay = angles[:, 1]
cy = torch.cos(ay)
sy = torch.sin(ay)
az = angles[:, 2]
cz = torch.cos(az)
sz = torch.sin(az)
Ax = torch.zeros((angles.shape[0], 3, 3), dtype=angles.dtype, device=angles.device)
Ax[:, 0, 0] = 1
Ax[:, 1, 1] = cx
Ax[:, 1, 2] = -sx
Ax[:, 2, 1] = sx
Ax[:, 2, 2] = cx
Ay = torch.zeros((angles.shape[0], 3, 3), dtype=angles.dtype, device=angles.device)
Ay[:, 1, 1] = 1
Ay[:, 0, 0] = cy
Ay[:, 0, 2] = sy
Ay[:, 2, 0] = -sy
Ay[:, 2, 2] = cy
Az = torch.zeros((angles.shape[0], 3, 3), dtype=angles.dtype, device=angles.device)
Az[:, 2, 2] = 1
Az[:, 0, 0] = cz
Az[:, 0, 1] = -sz
Az[:, 1, 0] = sz
Az[:, 1, 1] = cz
ret = torch.matmul(Az, torch.matmul(Ay, Ax))
return ret
def numpy_aet_to_rot(aet):
a = aet[0]
e = aet[1]
t = aet[2]
Ra = np.array([[np.cos(a), -np.sin(a), 0],
[np.sin(a), np.cos(a), 0],
[0, 0, 1]])
Re = np.array([[np.cos(e), 0, -np.sin(e)],
[0, 1, 0],
[np.sin(a), 0, np.cos(e)]])
Rt = np.array([[np.cos(t), -np.sin(t), 0],
[np.sin(t), np.cos(t), 0],
[0, 0, 1]])
R = np.matmul(Rt, np.matmul(Re, Ra))
return R