559 lines
17 KiB
Python
559 lines
17 KiB
Python
def quaternion_rotate(quaternion, axis, angle):
|
||
"""
|
||
Rotate a quaternion around a specified axis by a given angle.
|
||
|
||
Parameters:
|
||
quaternion (numpy array): The input quaternion [w, x, y, z].
|
||
axis (str): The axis to rotate around ('x', 'y', or 'z').
|
||
angle (float): The rotation angle in degrees.
|
||
|
||
Returns:
|
||
numpy array: The rotated quaternion.
|
||
"""
|
||
# Convert angle from degrees to radians
|
||
angle_rad = np.radians(angle)
|
||
|
||
# Calculate the rotation quaternion based on the specified axis
|
||
cos_half_angle = np.cos(angle_rad / 2)
|
||
sin_half_angle = np.sin(angle_rad / 2)
|
||
|
||
if axis == 'x':
|
||
q_axis = np.array([cos_half_angle, sin_half_angle, 0, 0])
|
||
elif axis == 'y':
|
||
q_axis = np.array([cos_half_angle, 0, sin_half_angle, 0])
|
||
elif axis == 'z':
|
||
q_axis = np.array([cos_half_angle, 0, 0, sin_half_angle])
|
||
else:
|
||
raise ValueError("Unsupported axis. Use 'x', 'y', or 'z'.")
|
||
|
||
# Extract components of the input quaternion
|
||
w1, x1, y1, z1 = quaternion
|
||
w2, x2, y2, z2 = q_axis
|
||
|
||
# Quaternion multiplication
|
||
w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
|
||
x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
|
||
y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
|
||
z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
|
||
|
||
return np.array([w, x, y, z])
|
||
|
||
|
||
def axis_to_quaternion(axis, target_axis='y'):
|
||
"""
|
||
Calculate the quaternion that rotates a given axis to the target axis.
|
||
|
||
Parameters:
|
||
axis (str): The axis in the object's local coordinate system ('x', 'y', or 'z').
|
||
target_axis (str): The target axis in the world coordinate system ('x', 'y', or 'z').
|
||
|
||
Returns:
|
||
numpy array: The quaternion representing the rotation.
|
||
"""
|
||
# Define unit vectors for each axis
|
||
unit_vectors = {
|
||
'x': np.array([1, 0, 0]),
|
||
'y': np.array([0, 1, 0]),
|
||
'z': np.array([0, 0, 1])
|
||
}
|
||
|
||
if axis not in unit_vectors or target_axis not in unit_vectors:
|
||
raise ValueError("Unsupported axis. Use 'x', 'y', or 'z'.")
|
||
|
||
v1 = unit_vectors[axis]
|
||
v2 = unit_vectors[target_axis]
|
||
|
||
# Calculate the cross product and dot product
|
||
cross_prod = np.cross(v1, v2)
|
||
dot_prod = np.dot(v1, v2)
|
||
|
||
# Calculate the quaternion
|
||
w = np.sqrt((np.linalg.norm(v1) ** 2) * (np.linalg.norm(v2) ** 2)) + dot_prod
|
||
x, y, z = cross_prod
|
||
|
||
# Normalize the quaternion
|
||
q = np.array([w, x, y, z])
|
||
q = q / np.linalg.norm(q)
|
||
|
||
return q
|
||
|
||
|
||
def is_y_axis_up(pose_matrix):
|
||
"""
|
||
判断物体在世界坐标系中的 y 轴是否朝上。
|
||
|
||
参数:
|
||
pose_matrix (numpy.ndarray): 4x4 齐次变换矩阵。
|
||
|
||
返回:
|
||
bool: True 如果 y 轴朝上，False 如果 y 轴朝下。
|
||
"""
|
||
# 提取旋转矩阵的第二列
|
||
y_axis_vector = pose_matrix[:3, 1]
|
||
|
||
# 世界坐标系的 y 轴
|
||
world_y_axis = np.array([0, 1, 0])
|
||
|
||
# 计算点积
|
||
dot_product = np.dot(y_axis_vector, world_y_axis)
|
||
|
||
# 返回 True 如果朝上，否则返回 False
|
||
return dot_product > 0
|
||
|
||
|
||
def is_local_axis_facing_world_axis(pose_matrix, local_axis='y', world_axis='z'):
|
||
# 定义局部坐标系的轴索引
|
||
local_axis_index = {'x': 0, 'y': 1, 'z': 2}
|
||
|
||
# 定义世界坐标系的轴向量
|
||
world_axes = {
|
||
'x': np.array([1, 0, 0]),
|
||
'y': np.array([0, 1, 0]),
|
||
'z': np.array([0, 0, 1])
|
||
}
|
||
|
||
# 提取局部坐标系的指定轴
|
||
local_axis_vector = pose_matrix[:3, local_axis_index[local_axis]]
|
||
|
||
# 获取世界坐标系的指定轴向量
|
||
world_axis_vector = world_axes[world_axis]
|
||
|
||
# 计算点积
|
||
dot_product = np.dot(local_axis_vector, world_axis_vector)
|
||
|
||
# 返回 True 如果朝向指定世界轴，否则返回 False
|
||
return dot_product > 0
|
||
|
||
|
||
def rotate_180_along_axis(target_affine, rot_axis='z'):
|
||
'''
|
||
gripper是对称结构，绕Z轴旋转180度前后是等效的。选择离当前pose更近的target pose以避免不必要的旋转
|
||
'''
|
||
if rot_axis == 'z':
|
||
R_180 = np.array([[-1, 0, 0],
|
||
[0, -1, 0],
|
||
[0, 0, 1]])
|
||
elif rot_axis == 'y':
|
||
R_180 = np.array([[-1, 0, 0],
|
||
[0, 1, 0],
|
||
[0, 0, -1]])
|
||
elif rot_axis == 'x':
|
||
R_180 = np.array([[1, 0, 0],
|
||
[0, -1, 0],
|
||
[0, 0, -1]])
|
||
else:
|
||
assert False, "Invalid rotation axis. Please choose from 'x', 'y', 'z'."
|
||
|
||
# 提取旋转部分（3x3矩阵）
|
||
target_rotation = target_affine[:3, :3]
|
||
# 将target_rotation绕其自身的Z轴转180度，得到target_rotation_2
|
||
target_rotation_2 = np.dot(target_rotation, R_180)
|
||
|
||
# 重新组合旋转矩阵target_rotation_2和原始的平移部分
|
||
target_affine_2 = np.eye(4)
|
||
target_affine_2[:3, :3] = target_rotation_2
|
||
target_affine_2[:3, 3] = target_affine[:3, 3] # 保留原始的平移部分
|
||
return target_affine_2
|
||
|
||
|
||
def rotate_along_axis(target_affine, angle_degrees, rot_axis='z', use_local=True):
|
||
'''
|
||
根据指定的角度和旋转轴来旋转target_affine。
|
||
参数:
|
||
- target_affine: 4x4 仿射变换矩阵
|
||
- angle_degrees: 旋转角度（以度为单位）
|
||
- rot_axis: 旋转轴，'x'、'y' 或 'z'
|
||
'''
|
||
# 将角度转换为弧度
|
||
angle_radians = np.deg2rad(angle_degrees)
|
||
|
||
# 创建旋转对象
|
||
if rot_axis == 'z':
|
||
rotation_vector = np.array([0, 0, angle_radians])
|
||
elif rot_axis == 'y':
|
||
rotation_vector = np.array([0, angle_radians, 0])
|
||
elif rot_axis == 'x':
|
||
rotation_vector = np.array([angle_radians, 0, 0])
|
||
else:
|
||
raise ValueError("Invalid rotation axis. Please choose from 'x', 'y', 'z'.")
|
||
|
||
# 生成旋转矩阵
|
||
R_angle = R.from_rotvec(rotation_vector).as_matrix()
|
||
|
||
# 提取旋转部分（3x3矩阵）
|
||
target_rotation = target_affine[:3, :3]
|
||
|
||
# 将 target_rotation 绕指定轴旋转指定角度，得到 target_rotation_2
|
||
if use_local:
|
||
target_rotation_2 = np.dot(target_rotation, R_angle)
|
||
else:
|
||
target_rotation_2 = np.dot(R_angle, target_rotation)
|
||
|
||
# 重新组合旋转矩阵 target_rotation_2 和原始的平移部分
|
||
target_affine_2 = np.eye(4)
|
||
target_affine_2[:3, :3] = target_rotation_2
|
||
target_affine_2[:3, 3] = target_affine[:3, 3] # 保留原始的平移部分
|
||
|
||
return target_affine_2
|
||
|
||
|
||
def rotate_along_axis(target_affine, angle_degrees, rot_axis='z', use_local=True):
|
||
'''
|
||
根据指定的角度和旋转轴来旋转target_affine。
|
||
参数:
|
||
- target_affine: 4x4 仿射变换矩阵
|
||
- angle_degrees: 旋转角度（以度为单位）
|
||
- rot_axis: 旋转轴，'x'、'y' 或 'z'
|
||
'''
|
||
# 将角度转换为弧度
|
||
angle_radians = np.deg2rad(angle_degrees)
|
||
|
||
# 创建旋转对象
|
||
if rot_axis == 'z':
|
||
rotation_vector = np.array([0, 0, angle_radians])
|
||
elif rot_axis == 'y':
|
||
rotation_vector = np.array([0, angle_radians, 0])
|
||
elif rot_axis == 'x':
|
||
rotation_vector = np.array([angle_radians, 0, 0])
|
||
else:
|
||
raise ValueError("Invalid rotation axis. Please choose from 'x', 'y', 'z'.")
|
||
|
||
# 生成旋转矩阵
|
||
R_angle = R.from_rotvec(rotation_vector).as_matrix()
|
||
|
||
# 提取旋转部分（3x3矩阵）
|
||
target_rotation = target_affine[:3, :3]
|
||
|
||
# 将 target_rotation 绕指定轴旋转指定角度，得到 target_rotation_2
|
||
if use_local:
|
||
target_rotation_2 = np.dot(target_rotation, R_angle)
|
||
else:
|
||
target_rotation_2 = np.dot(R_angle, target_rotation)
|
||
|
||
# 重新组合旋转矩阵 target_rotation_2 和原始的平移部分
|
||
target_affine_2 = np.eye(4)
|
||
target_affine_2[:3, :3] = target_rotation_2
|
||
target_affine_2[:3, 3] = target_affine[:3, 3] # 保留原始的平移部分
|
||
|
||
return target_affine_2
|
||
|
||
|
||
import numpy as np
|
||
from scipy.spatial.transform import Rotation as R
|
||
|
||
|
||
def get_quaternion_wxyz_from_rotation_matrix(rotation_matrix):
|
||
"""
|
||
Convert a 3x3 rotation matrix to a quaternion in the wxyz format.
|
||
|
||
Parameters:
|
||
R (numpy array): A 3x3 rotation matrix.
|
||
|
||
Returns:
|
||
numpy array: A 4x1 quaternion in the wxyz format.
|
||
"""
|
||
# Convert the rotation matrix to a quaternion
|
||
rot = R.from_matrix(rotation_matrix)
|
||
quat = rot.as_quat()
|
||
|
||
# Reorder the quaternion to the wxyz format
|
||
if quat.shape[0] == 4:
|
||
quaternions_wxyz = quat[[3, 0, 1, 2]]
|
||
else:
|
||
quaternions_wxyz = quat[:, [3, 0, 1, 2]]
|
||
return quaternions_wxyz
|
||
|
||
|
||
def get_quaternion_from_rotation_matrix(R):
|
||
assert R.shape == (3, 3)
|
||
|
||
# 计算四元数分量
|
||
trace = np.trace(R)
|
||
if trace > 0:
|
||
S = np.sqrt(trace + 1.0) * 2 # S=4*qw
|
||
qw = 0.25 * S
|
||
qx = (R[2, 1] - R[1, 2]) / S
|
||
qy = (R[0, 2] - R[2, 0]) / S
|
||
qz = (R[1, 0] - R[0, 1]) / S
|
||
elif (R[0, 0] > R[1, 1]) and (R[0, 0] > R[2, 2]):
|
||
S = np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2]) * 2 # S=4*qx
|
||
qw = (R[2, 1] - R[1, 2]) / S
|
||
qx = 0.25 * S
|
||
qy = (R[0, 1] + R[1, 0]) / S
|
||
qz = (R[0, 2] + R[2, 0]) / S
|
||
elif R[1, 1] > R[2, 2]:
|
||
S = np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2]) * 2 # S=4*qy
|
||
qw = (R[0, 2] - R[2, 0]) / S
|
||
qx = (R[0, 1] + R[1, 0]) / S
|
||
qy = 0.25 * S
|
||
qz = (R[1, 2] + R[2, 1]) / S
|
||
else:
|
||
S = np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1]) * 2 # S=4*qz
|
||
qw = (R[1, 0] - R[0, 1]) / S
|
||
qx = (R[0, 2] + R[2, 0]) / S
|
||
qy = (R[1, 2] + R[2, 1]) / S
|
||
qz = 0.25 * S
|
||
|
||
return np.array([qw, qx, qy, qz])
|
||
|
||
|
||
# @nb.jit(nopython=True)
|
||
def get_rotation_matrix_from_quaternion(quat: np.ndarray) -> np.ndarray:
|
||
"""Convert a quaternion to a rotation matrix.
|
||
|
||
Args:
|
||
quat (np.ndarray): A 4x1 vector in order (w, x, y, z)
|
||
|
||
Returns:
|
||
np.ndarray: The resulting 3x3 rotation matrix.
|
||
"""
|
||
w, x, y, z = quat
|
||
rot = np.array(
|
||
[
|
||
[2 * (w ** 2 + x ** 2) - 1, 2 * (x * y - w * z), 2 * (x * z + w * y)],
|
||
[2 * (x * y + w * z), 2 * (w ** 2 + y ** 2) - 1, 2 * (y * z - w * x)],
|
||
[2 * (x * z - w * y), 2 * (y * z + w * x), 2 * (w ** 2 + z ** 2) - 1],
|
||
]
|
||
)
|
||
return rot
|
||
|
||
|
||
# @nb.jit(nopython=True)
|
||
def get_xyz_euler_from_quaternion(quat: np.ndarray) -> np.ndarray:
|
||
"""Convert a quaternion to XYZ euler angles.
|
||
|
||
Args:
|
||
quat (np.ndarray): A 4x1 vector in order (w, x, y, z).
|
||
|
||
Returns:
|
||
np.ndarray: A 3x1 vector containing (roll, pitch, yaw).
|
||
"""
|
||
w, x, y, z = quat
|
||
y_sqr = y * y
|
||
|
||
t0 = +2.0 * (w * x + y * z)
|
||
t1 = +1.0 - 2.0 * (x * x + y_sqr)
|
||
eulerx = np.arctan2(t0, t1)
|
||
|
||
t2 = +2.0 * (w * y - z * x)
|
||
t2 = +1.0 if t2 > +1.0 else t2
|
||
t2 = -1.0 if t2 < -1.0 else t2
|
||
eulery = np.arcsin(t2)
|
||
|
||
t3 = +2.0 * (w * z + x * y)
|
||
t4 = +1.0 - 2.0 * (y_sqr + z * z)
|
||
eulerz = np.arctan2(t3, t4)
|
||
|
||
result = np.zeros(3)
|
||
result[0] = eulerx
|
||
result[1] = eulery
|
||
result[2] = eulerz
|
||
|
||
return result
|
||
|
||
|
||
# @nb.jit(nopython=True)
|
||
def get_quaternion_from_euler(euler: np.ndarray, order: str = "XYZ") -> np.ndarray:
|
||
"""Convert an euler angle to a quaternion based on specified euler angle order.
|
||
|
||
Supported Euler angle orders: {'XYZ', 'YXZ', 'ZXY', 'ZYX', 'YZX', 'XZY'}.
|
||
|
||
Args:
|
||
euler (np.ndarray): A 3x1 vector with angles in radians.
|
||
order (str, optional): The specified order of input euler angles. Defaults to "XYZ".
|
||
|
||
Raises:
|
||
ValueError: If input order is not valid.
|
||
|
||
Reference:
|
||
[1] https://github.com/mrdoob/three.js/blob/master/src/math/Quaternion.js
|
||
"""
|
||
# extract input angles
|
||
r, p, y = euler
|
||
# compute constants
|
||
y = y / 2.0
|
||
p = p / 2.0
|
||
r = r / 2.0
|
||
c3 = np.cos(y)
|
||
s3 = np.sin(y)
|
||
c2 = np.cos(p)
|
||
s2 = np.sin(p)
|
||
c1 = np.cos(r)
|
||
s1 = np.sin(r)
|
||
# convert to quaternion based on order
|
||
if order == "XYZ":
|
||
result = np.array(
|
||
[
|
||
c1 * c2 * c3 - s1 * s2 * s3,
|
||
c1 * s2 * s3 + c2 * c3 * s1,
|
||
c1 * c3 * s2 - s1 * c2 * s3,
|
||
c1 * c2 * s3 + s1 * c3 * s2,
|
||
]
|
||
)
|
||
if result[0] < 0:
|
||
result = -result
|
||
return result
|
||
elif order == "YXZ":
|
||
result = np.array(
|
||
[
|
||
c1 * c2 * c3 + s1 * s2 * s3,
|
||
s1 * c2 * c3 + c1 * s2 * s3,
|
||
c1 * s2 * c3 - s1 * c2 * s3,
|
||
c1 * c2 * s3 - s1 * s2 * c3,
|
||
]
|
||
)
|
||
return result
|
||
elif order == "ZXY":
|
||
result = np.array(
|
||
[
|
||
c1 * c2 * c3 - s1 * s2 * s3,
|
||
s1 * c2 * c3 - c1 * s2 * s3,
|
||
c1 * s2 * c3 + s1 * c2 * s3,
|
||
c1 * c2 * s3 + s1 * s2 * c3,
|
||
]
|
||
)
|
||
return result
|
||
elif order == "ZYX":
|
||
result = np.array(
|
||
[
|
||
c1 * c2 * c3 + s1 * s2 * s3,
|
||
s1 * c2 * c3 - c1 * s2 * s3,
|
||
c1 * s2 * c3 + s1 * c2 * s3,
|
||
c1 * c2 * s3 - s1 * s2 * c3,
|
||
]
|
||
)
|
||
return result
|
||
elif order == "YZX":
|
||
result = np.array(
|
||
[
|
||
c1 * c2 * c3 - s1 * s2 * s3,
|
||
s1 * c2 * c3 + c1 * s2 * s3,
|
||
c1 * s2 * c3 + s1 * c2 * s3,
|
||
c1 * c2 * s3 - s1 * s2 * c3,
|
||
]
|
||
)
|
||
return result
|
||
elif order == "XZY":
|
||
result = np.array(
|
||
[
|
||
c1 * c2 * c3 + s1 * s2 * s3,
|
||
s1 * c2 * c3 - c1 * s2 * s3,
|
||
c1 * s2 * c3 - s1 * c2 * s3,
|
||
c1 * c2 * s3 + s1 * s2 * c3,
|
||
]
|
||
)
|
||
return result
|
||
else:
|
||
raise ValueError("Input euler angle order is meaningless.")
|
||
|
||
|
||
# @nb.jit(nopython=True)
|
||
def get_rotation_matrix_from_euler(euler: np.ndarray, order: str = "XYZ") -> np.ndarray:
|
||
quat = get_quaternion_from_euler(euler, order)
|
||
return get_rotation_matrix_from_quaternion(quat)
|
||
|
||
|
||
# @nb.jit(nopython=True)
|
||
def quat_multiplication(q: np.ndarray, p: np.ndarray) -> np.ndarray:
|
||
"""Compute the product of two quaternions.
|
||
|
||
Args:
|
||
q (np.ndarray): First quaternion in order (w, x, y, z).
|
||
p (np.ndarray): Second quaternion in order (w, x, y, z).
|
||
|
||
Returns:
|
||
np.ndarray: A 4x1 vector representing a quaternion in order (w, x, y, z).
|
||
"""
|
||
quat = np.array(
|
||
[
|
||
p[0] * q[0] - p[1] * q[1] - p[2] * q[2] - p[3] * q[3],
|
||
p[0] * q[1] + p[1] * q[0] - p[2] * q[3] + p[3] * q[2],
|
||
p[0] * q[2] + p[1] * q[3] + p[2] * q[0] - p[3] * q[1],
|
||
p[0] * q[3] - p[1] * q[2] + p[2] * q[1] + p[3] * q[0],
|
||
]
|
||
)
|
||
return quat
|
||
|
||
|
||
# @nb.jit(nopython=True)
|
||
def skew(vector: np.ndarray) -> np.ndarray:
|
||
"""Convert vector to skew symmetric matrix.
|
||
|
||
This function returns a skew-symmetric matrix to perform cross-product
|
||
as a matrix multiplication operation, i.e.:
|
||
|
||
np.cross(a, b) = np.dot(skew(a), b)
|
||
|
||
|
||
Args:
|
||
vector (np.ndarray): A 3x1 vector.
|
||
|
||
Returns:
|
||
np.ndarray: The resluting skew-symmetric matrix.
|
||
"""
|
||
mat = np.array([[0, -vector[2], vector[1]], [vector[2], 0, -vector[0]], [-vector[1], vector[0], 0]])
|
||
return mat
|
||
|
||
|
||
import math
|
||
|
||
_POLE_LIMIT = 1.0 - 1e-6
|
||
|
||
|
||
def matrix_to_euler_angles(mat: np.ndarray, degrees: bool = False, extrinsic: bool = True) -> np.ndarray:
|
||
"""Convert rotation matrix to Euler XYZ extrinsic or intrinsic angles.
|
||
|
||
Args:
|
||
mat (np.ndarray): A 3x3 rotation matrix.
|
||
degrees (bool, optional): Whether returned angles should be in degrees.
|
||
extrinsic (bool, optional): True if the rotation matrix follows the extrinsic matrix
|
||
convention (equivalent to ZYX ordering but returned in the reverse) and False if it follows
|
||
the intrinsic matrix conventions (equivalent to XYZ ordering).
|
||
Defaults to True.
|
||
|
||
Returns:
|
||
np.ndarray: Euler XYZ angles (intrinsic form) if extrinsic is False and Euler XYZ angles (extrinsic form) if extrinsic is True.
|
||
"""
|
||
if extrinsic:
|
||
if mat[2, 0] > _POLE_LIMIT:
|
||
roll = np.arctan2(mat[0, 1], mat[0, 2])
|
||
pitch = -np.pi / 2
|
||
yaw = 0.0
|
||
return np.array([roll, pitch, yaw])
|
||
|
||
if mat[2, 0] < -_POLE_LIMIT:
|
||
roll = np.arctan2(mat[0, 1], mat[0, 2])
|
||
pitch = np.pi / 2
|
||
yaw = 0.0
|
||
return np.array([roll, pitch, yaw])
|
||
|
||
roll = np.arctan2(mat[2, 1], mat[2, 2])
|
||
pitch = -np.arcsin(mat[2, 0])
|
||
yaw = np.arctan2(mat[1, 0], mat[0, 0])
|
||
if degrees:
|
||
roll = math.degrees(roll)
|
||
pitch = math.degrees(pitch)
|
||
yaw = math.degrees(yaw)
|
||
return np.array([roll, pitch, yaw])
|
||
else:
|
||
if mat[0, 2] > _POLE_LIMIT:
|
||
roll = np.arctan2(mat[1, 0], mat[1, 1])
|
||
pitch = np.pi / 2
|
||
yaw = 0.0
|
||
return np.array([roll, pitch, yaw])
|
||
|
||
if mat[0, 2] < -_POLE_LIMIT:
|
||
roll = np.arctan2(mat[1, 0], mat[1, 1])
|
||
pitch = -np.pi / 2
|
||
yaw = 0.0
|
||
return np.array([roll, pitch, yaw])
|
||
roll = -math.atan2(mat[1, 2], mat[2, 2])
|
||
pitch = math.asin(mat[0, 2])
|
||
yaw = -math.atan2(mat[0, 1], mat[0, 0])
|
||
|
||
if degrees:
|
||
roll = math.degrees(roll)
|
||
pitch = math.degrees(pitch)
|
||
yaw = math.degrees(yaw)
|
||
return np.array([roll, pitch, yaw])
|