31 lines
1.0 KiB
Python
31 lines
1.0 KiB
Python
|
import os
|
||
|
|
import numpy as np
|
||
|
|
import random
|
||
|
|
import torch
|
||
|
|
|
||
|
|
@torch.jit.script
|
||
|
|
def copysign(a, b):
|
||
|
|
# type: (float, Tensor) -> Tensor
|
||
|
|
a = torch.tensor(a, device=b.device, dtype=torch.float).repeat(b.shape[0])
|
||
|
|
return torch.abs(a) * torch.sign(b)
|
||
|
|
def get_euler_xyz(q):
|
||
|
|
qx, qy, qz, qw = 0, 1, 2, 3
|
||
|
|
# roll (x-axis rotation)
|
||
|
|
sinr_cosp = 2.0 * (q[:, qw] * q[:, qx] + q[:, qy] * q[:, qz])
|
||
|
|
cosr_cosp = q[:, qw] * q[:, qw] - q[:, qx] * \
|
||
|
|
q[:, qx] - q[:, qy] * q[:, qy] + q[:, qz] * q[:, qz]
|
||
|
|
roll = torch.atan2(sinr_cosp, cosr_cosp)
|
||
|
|
|
||
|
|
# pitch (y-axis rotation)
|
||
|
|
sinp = 2.0 * (q[:, qw] * q[:, qy] - q[:, qz] * q[:, qx])
|
||
|
|
pitch = torch.where(
|
||
|
|
torch.abs(sinp) >= 1, copysign(np.pi / 2.0, sinp), torch.asin(sinp))
|
||
|
|
|
||
|
|
# yaw (z-axis rotation)
|
||
|
|
siny_cosp = 2.0 * (q[:, qw] * q[:, qz] + q[:, qx] * q[:, qy])
|
||
|
|
cosy_cosp = q[:, qw] * q[:, qw] + q[:, qx] * \
|
||
|
|
q[:, qx] - q[:, qy] * q[:, qy] - q[:, qz] * q[:, qz]
|
||
|
|
yaw = torch.atan2(siny_cosp, cosy_cosp)
|
||
|
|
|
||
|
|
return torch.stack((roll, pitch, yaw), dim=-1)
|