Spaces:
Runtime error
Runtime error
File size: 2,392 Bytes
35e2073 |
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 |
import numpy as np
import torch
from scipy.spatial.transform import Rotation as RotLib
def SO3_to_quat(R):
"""
:param R: (N, 3, 3) or (3, 3) np
:return: (N, 4, ) or (4, ) np
"""
x = RotLib.from_matrix(R)
quat = x.as_quat()
return quat
def quat_to_SO3(quat):
"""
:param quat: (N, 4, ) or (4, ) np
:return: (N, 3, 3) or (3, 3) np
"""
x = RotLib.from_quat(quat)
R = x.as_matrix()
return R
def convert3x4_4x4(input):
"""
:param input: (N, 3, 4) or (3, 4) torch or np
:return: (N, 4, 4) or (4, 4) torch or np
"""
if torch.is_tensor(input):
if len(input.shape) == 3:
output = torch.cat([input, torch.zeros_like(input[:, 0:1])], dim=1) # (N, 4, 4)
output[:, 3, 3] = 1.0
else:
output = torch.cat([input, torch.tensor([[0,0,0,1]], dtype=input.dtype, device=input.device)], dim=0) # (4, 4)
else:
if len(input.shape) == 3:
output = np.concatenate([input, np.zeros_like(input[:, 0:1])], axis=1) # (N, 4, 4)
output[:, 3, 3] = 1.0
else:
output = np.concatenate([input, np.array([[0,0,0,1]], dtype=input.dtype)], axis=0) # (4, 4)
output[3, 3] = 1.0
return output
def vec2skew(v):
"""
:param v: (3, ) torch tensor
:return: (3, 3)
"""
zero = torch.zeros(1, dtype=torch.float32, device=v.device)
skew_v0 = torch.cat([ zero, -v[2:3], v[1:2]]) # (3, 1)
skew_v1 = torch.cat([ v[2:3], zero, -v[0:1]])
skew_v2 = torch.cat([-v[1:2], v[0:1], zero])
skew_v = torch.stack([skew_v0, skew_v1, skew_v2], dim=0) # (3, 3)
return skew_v # (3, 3)
def Exp(r):
"""so(3) vector to SO(3) matrix
:param r: (3, ) axis-angle, torch tensor
:return: (3, 3)
"""
skew_r = vec2skew(r) # (3, 3)
norm_r = r.norm() + 1e-15
eye = torch.eye(3, dtype=torch.float32, device=r.device)
R = eye + (torch.sin(norm_r) / norm_r) * skew_r + ((1 - torch.cos(norm_r)) / norm_r**2) * (skew_r @ skew_r)
return R
def make_c2w(r, t):
"""
:param r: (3, ) axis-angle torch tensor
:param t: (3, ) translation vector torch tensor
:return: (4, 4)
"""
R = Exp(r) # (3, 3)
c2w = torch.cat([R, t.unsqueeze(1)], dim=1) # (3, 4)
c2w = convert3x4_4x4(c2w) # (4, 4)
return c2w
|