26 lines
841 B
Python
26 lines
841 B
Python
|
import numpy as np
|
||
|
|
from scipy.spatial.transform import Rotation as R
|
||
|
|
|
||
|
|
|
||
|
|
def get_gravity_orientation(quaternion):
|
||
|
|
qw = quaternion[0]
|
||
|
|
qx = quaternion[1]
|
||
|
|
qy = quaternion[2]
|
||
|
|
qz = quaternion[3]
|
||
|
|
|
||
|
|
gravity_orientation = np.zeros(3)
|
||
|
|
|
||
|
|
gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
|
||
|
|
gravity_orientation[1] = -2 * (qz * qy + qw * qx)
|
||
|
|
gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)
|
||
|
|
|
||
|
|
return gravity_orientation
|
||
|
|
|
||
|
|
|
||
|
|
def transform_imu_data(waist_yaw, waist_yaw_omega, imu_quat, imu_omega):
|
||
|
|
RzWaist = R.from_euler("z", waist_yaw).as_matrix()
|
||
|
|
R_torso = R.from_quat([imu_quat[1], imu_quat[2], imu_quat[3], imu_quat[0]]).as_matrix()
|
||
|
|
R_pelvis = np.dot(R_torso, RzWaist.T)
|
||
|
|
w = np.dot(RzWaist, imu_omega[0]) - np.array([0, 0, waist_yaw_omega])
|
||
|
|
return R.from_matrix(R_pelvis).as_quat()[[3, 0, 1, 2]], w
|