Init

deploy/deploy_real/common/rotation_helper.py

@@ -0,0 +1,25 @@
+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