From e9638630ac4a04e7f027eb7151b764f14a825a4e Mon Sep 17 00:00:00 2001 From: Cathy Date: Wed, 27 Dec 2023 10:56:04 -0500 Subject: [PATCH] add pose util --- .../pyrcareworld/utils/pose_converter.py | 352 ++++++++++++++++++ 1 file changed, 352 insertions(+) create mode 100644 pyrcareworld/pyrcareworld/utils/pose_converter.py diff --git a/pyrcareworld/pyrcareworld/utils/pose_converter.py b/pyrcareworld/pyrcareworld/utils/pose_converter.py new file mode 100644 index 00000000..c8208314 --- /dev/null +++ b/pyrcareworld/pyrcareworld/utils/pose_converter.py @@ -0,0 +1,352 @@ +import numpy as np + + +class CoordinateSystemConverter(): + """ + + Coordinate System Converter class. + + Args: + cs1_direction: list, The visual direction corresponding to the xyz axis, + can be: + ["left"/"right"/"up"/"down"/"forward"/"back"/ + "l"/"r"/"u"/"d"/"f"/"b"/ + "-left"/"-right"/"-up"/"-down"/"-forward"/"-back"/ + "-l"/"-r"/"-u"/"-d"/"-f"/"-b"] + + cs2_direction: list, The visual direction corresponding to the xyz axis, + can be: + ["left"/"right"/"up"/"down"/"forward"/"back"/ + "l"/"r"/"u"/"d"/"f"/"b"/ + "-left"/"-right"/"-up"/"-down"/"-forward"/"-back"/ + "-l"/"-r"/"-u"/"-d"/"-f"/"-b"] + """ + def __init__(self, cs1_direction=["right", "up", "forward"], cs2_direction=["right", "up", "forward"]): + self.cs1_dir, self.cs1_axis, self.cs1_sign, self.cs1_system = self._preprocessing(cs1_direction) + self.cs2_dir, self.cs2_axis, self.cs2_sign, self.cs2_system = self._preprocessing(cs2_direction) + self.cs1_to_cs2_map = [ + self.cs1_axis[self.cs2_dir[0]], + self.cs1_axis[self.cs2_dir[1]], + self.cs1_axis[self.cs2_dir[2]], + ] + + self.cs1_to_cs2_sign = [ + self.cs1_sign[self.cs2_dir[0]] * self.cs2_sign[self.cs2_dir[0]], + self.cs1_sign[self.cs2_dir[1]] * self.cs2_sign[self.cs2_dir[1]], + self.cs1_sign[self.cs2_dir[2]] * self.cs2_sign[self.cs2_dir[2]], + ] + self.cs2_to_cs1_map = [ + self.cs2_axis[self.cs1_dir[0]], + self.cs2_axis[self.cs1_dir[1]], + self.cs2_axis[self.cs1_dir[2]], + ] + + self.cs2_to_cs1_sign = [ + self.cs1_sign[self.cs1_dir[0]] * self.cs2_sign[self.cs1_dir[0]], + self.cs1_sign[self.cs1_dir[1]] * self.cs2_sign[self.cs1_dir[1]], + self.cs1_sign[self.cs1_dir[2]] * self.cs2_sign[self.cs1_dir[2]], + ] + + self.cs1_left_or_right = self.cs1_sign['right'] * self.cs1_sign['up'] * self.cs1_sign['forward'] + self.cs2_left_or_right = self.cs2_sign['right'] * self.cs2_sign['up'] * self.cs2_sign['forward'] + + + + def _preprocessing(self, direction): + dir = ["", "", ""] + axis = {} + sign = {} + + direction_in_left = [[], [], []] + for i in range(3): + if direction[i] in ["right", "r", "-l", "-left"]: + dir[i] = "right" + axis["right"] = i + sign["right"] = 1 + direction_in_left[i] = [1, 0, 0] + + elif direction[i] in ["up", "u", "-d", "-down"]: + dir[i] = "up" + axis["up"] = i + sign["up"] = 1 + direction_in_left[i] = [0, 1, 0] + + elif direction[i] in ["forward", "f", "-b", "-back"]: + dir[i] = "forward" + axis["forward"] = i + sign["forward"] = 1 + direction_in_left[i] = [0, 0, 1] + + elif direction[i] in ["-right", "-r", "l", "left"]: + dir[i] = "right" + axis["right"] = i + sign["right"] = -1 + direction_in_left[i] = [-1, 0, 0] + + elif direction[i] in ["-up", "-u", "d", "down"]: + dir[i] = "up" + axis["up"] = i + sign["up"] = -1 + direction_in_left[i] = [0, -1, 0] + + elif direction[i] in ["-forward", "-f", "b", "back"]: + dir[i] = "forward" + axis["forward"] = i + sign["forward"] = -1 + direction_in_left[i] = [0, 0, -1] + + else: + raise f"{direction[i]} is Unrecognized axis" + + + + z = self._cross(direction_in_left[0], direction_in_left[1]) + if z == direction_in_left[2]: + system = 1 + + else: + system = -1 + + return dir, axis, sign, system + + + + def _cross(self, lhs: list, rhs: list) -> list: + return [lhs[1] * rhs[2] - lhs[2] * rhs[1], lhs[2] * rhs[0] - lhs[0] * rhs[2], lhs[0] * rhs[1] - lhs[1] * rhs[0]] + + + + def cs1_pos_to_cs2_pos(self, pos: list) -> list: + + """ + Convert position form Coordinate System 1 to Coordinate System 2. + Args: + pos: List of length 3, position of Coordinate System 1. + + + + Return: + list: List of length 3, position of Coordinate System 2. + + """ + + x = pos[self.cs1_to_cs2_map[0]] * self.cs1_to_cs2_sign[0] + y = pos[self.cs1_to_cs2_map[1]] * self.cs1_to_cs2_sign[1] + z = pos[self.cs1_to_cs2_map[2]] * self.cs1_to_cs2_sign[2] + + return [x, y, z] + + + + def cs2_pos_to_cs1_pos(self, pos: list) -> list: + + """ + Convert position form Coordinate System 2 to Coordinate System 1. + + Args: + pos: List of length 3, position of Coordinate System 2. + + + + Return: + list: List of length 3, position of Coordinate System 1. + + """ + + x = pos[self.cs2_to_cs1_map[0]] * self.cs2_to_cs1_sign[0] + + y = pos[self.cs2_to_cs1_map[1]] * self.cs2_to_cs1_sign[1] + + z = pos[self.cs2_to_cs1_map[2]] * self.cs2_to_cs1_sign[2] + + return [x, y, z] + + + + def cs1_quat_to_cs2_quat(self, quat: list) -> list: + + """ + Convert quaternion form Coordinate System 1 to Coordinate System 2. + + Args: + quat: List of length 3, quaternion of Coordinate System 1. + + + + Return: + list: List of length 3, quaternion of Coordinate System 2. + + """ + + x = quat[self.cs1_to_cs2_map[0]] * self.cs1_to_cs2_sign[0] + y = quat[self.cs1_to_cs2_map[1]] * self.cs1_to_cs2_sign[1] + z = quat[self.cs1_to_cs2_map[2]] * self.cs1_to_cs2_sign[2] + w = quat[3] * self.cs1_system * self.cs2_system + + return [x, y, z, w] + + + + def cs2_quat_to_cs1_quat(self, quat: list) -> list: + + """ + Convert quaternion form Coordinate System 2 to Coordinate System 1. + Args: + quat: List of length 4, quaternion[x,y,z,w] of Coordinate System 2. + + + + Return: + list: List of length 4, quaternion[x,y,z,w] of Coordinate System 1. + + """ + + x = quat[self.cs2_to_cs1_map[0]] * self.cs2_to_cs1_sign[0] + y = quat[self.cs2_to_cs1_map[1]] * self.cs2_to_cs1_sign[1] + z = quat[self.cs2_to_cs1_map[2]] * self.cs2_to_cs1_sign[2] + w = quat[3] * self.cs1_system * self.cs2_system + + return [x, y, z, w] + + + + def cs1_scale_to_cs2_scale(self, scale: list) -> list: + + """ + Convert scale form Coordinate System 1 to Coordinate System 2. + + Args: + scale: List of length 3, scale of Coordinate System 1. + + Return: + list: List of length 3, scale of Coordinate System 2. + + """ + + x = scale[self.cs1_to_cs2_map[0]] + y = scale[self.cs1_to_cs2_map[1]] + z = scale[self.cs1_to_cs2_map[2]] + + return [x, y, z] + + + + def cs2_scale_to_cs1_scale(self, scale: list) -> list: + + """ + Convert scale form Coordinate System 2 to Coordinate System 1. + + Args: + scale: List of length 3, scale of Coordinate System 2. + + Return: + list: List of length 3, scale of Coordinate System 1. + + """ + + x = scale[self.cs2_to_cs1_map[0]] + y = scale[self.cs2_to_cs1_map[1]] + z = scale[self.cs2_to_cs1_map[2]] + + return [x, y, z] + + + + def cs1_matrix_to_cs2_matrix(self, matrix) -> np.ndarray: + + """ + Convert rotation matrix form Coordinate System 1 to Coordinate System 2. + Args: + + matrix: list or np.ndarray shape[3,3], rotation matrix of Coordinate System 1. + + Return: + np.ndarray: shape[3,3], rotation matrix of Coordinate System 2. + + """ + + quaternion = self.matrix_to_quat(matrix) + quaternion = self.cs1_quat_to_cs2_quat(quaternion) + return self.quat_to_matrix(quaternion) + + + + def cs2_matrix_to_cs1_matrix(self, matrix) -> np.ndarray: + + """ + Convert rotation matrix form Coordinate System 2 to Coordinate System 1. + + Args: + matrix: list or np.ndarray shape[3,3], rotation matrix of Coordinate System 2. + + Return: + np.ndarray: shape[3,3], rotation matrix of Coordinate System 1. + + """ + + quaternion = self.matrix_to_quat(matrix) + quaternion = self.cs2_quat_to_cs1_quat(quaternion) + + return self.quat_to_matrix(quaternion) + + + + def quat_to_matrix(self, quat=[0,0,0,1]) -> np.ndarray: + """ + Convert quaternion to rotation matrix. + Args: + quat: List of length 4, representing quaternion in [x, y, z, w] order. + + Return: + np.ndarray: shape[3,3], rotation matrix. + """ + return np.array([ + [1 - 2 * (quat[1] ** 2 + quat[2] ** 2), 2 * (quat[0] * quat[1] - quat[2] * quat[3]), 2 * (quat[0] * quat[2] + quat[1] * quat[3])], + [2 * (quat[0] * quat[1] + quat[2] * quat[3]), 1 - 2 * (quat[0] ** 2 + quat[2] ** 2), 2 * (quat[1] * quat[2] - quat[0] * quat[3])], + [2 * (quat[0] * quat[2] - quat[1] * quat[3]), 2 * (quat[1] * quat[2] + quat[0] * quat[3]), 1 - 2 * (quat[0] ** 2 + quat[1] ** 2)] + + ]) + + + + def matrix_to_quat(self, matrix) -> list: + """ + Convert rotation matrix to quaternion. + Args: + matrix: list or np.ndarray shape[3,3], rotation matrix. + + Return: + list: List of length 4, representing quaternion in [x, y, z, w] order. + """ + m = np.array(matrix, dtype=float) + trace = np.trace(m) + + if trace > 0: + s = 0.5 / np.sqrt(trace + 1.0) + w = 0.25 / s + x = (m[2, 1] - m[1, 2]) * s + y = (m[0, 2] - m[2, 0]) * s + z = (m[1, 0] - m[0, 1]) * s + + elif m[0, 0] > m[1, 1] and m[0, 0] > m[2, 2]: + s = 2.0 * np.sqrt(1.0 + m[0, 0] - m[1, 1] - m[2, 2]) + w = (m[2, 1] - m[1, 2]) / s + x = 0.25 * s + y = (m[0, 1] + m[1, 0]) / s + z = (m[0, 2] + m[2, 0]) / s + + elif m[1, 1] > m[2, 2]: + s = 2.0 * np.sqrt(1.0 + m[1, 1] - m[0, 0] - m[2, 2]) + w = (m[0, 2] - m[2, 0]) / s + x = (m[0, 1] + m[1, 0]) / s + y = 0.25 * s + z = (m[1, 2] + m[2, 1]) / s + + else: + s = 2.0 * np.sqrt(1.0 + m[2, 2] - m[0, 0] - m[1, 1]) + w = (m[1, 0] - m[0, 1]) / s + x = (m[0, 2] + m[2, 0]) / s + y = (m[1, 2] + m[2, 1]) / s + z = 0.25 * s + + return [x, y, z, w] \ No newline at end of file