OpenMotionLab/MotionGPT

1from .geometry_tools import *2 3 4def rep_to_rep(oldtype, newtype, rotations):5    if newtype in ["matrix"]:6        return to_matrix(oldtype, rotations)7    8    if oldtype in ["rotvec", "axisangle"]:9        return axis_angle_to(newtype, rotations)10    elif oldtype in ["matrix"]:11        return matrix_to(newtype, rotations)12    else:13        raise NotImplementedError("Only rotvec and matrix are supported.")14 15def nfeats_of(rottype):16    if rottype in ["rotvec", "axisangle"]:17        return 318    elif rottype in ["rotquat", "quaternion"]:19        return 420    elif rottype in ["rot6d", "6drot", "rotation6d"]:21        return 622    elif rottype in ["rotmat"]:23        return 924    else:25        return TypeError("This rotation type doesn't have features.")26 27 28def axis_angle_to(newtype, rotations):29    if newtype in ["matrix"]:30        rotations = axis_angle_to_matrix(rotations)31        return rotations32    elif newtype in ["rotmat"]:33        rotations = axis_angle_to_matrix(rotations)34        rotations = matrix_to("rotmat", rotations)35        return rotations36    elif newtype in ["rot6d", "6drot", "rotation6d"]:37        rotations = axis_angle_to_matrix(rotations)38        rotations = matrix_to("rot6d", rotations)39        return rotations40    elif newtype in ["rotquat", "quaternion"]:41        rotations = axis_angle_to_quaternion(rotations)42        return rotations43    elif newtype in ["rotvec", "axisangle"]:44        return rotations45    else:46        raise NotImplementedError47 48 49def matrix_to(newtype, rotations):50    if newtype in ["matrix"]:51        return rotations52    if newtype in ["rotmat"]:53        rotations = rotations.reshape((*rotations.shape[:-2], 9))54        return rotations55    elif newtype in ["rot6d", "6drot", "rotation6d"]:56        rotations = matrix_to_rotation_6d(rotations)57        return rotations58    elif newtype in ["rotquat", "quaternion"]:59        rotations = matrix_to_quaternion(rotations)60        return rotations61    elif newtype in ["rotvec", "axisangle"]:62        rotations = matrix_to_axis_angle(rotations)63        return rotations64    else:65        raise NotImplementedError66 67 68def to_matrix(oldtype, rotations):69    if oldtype in ["matrix"]:70        return rotations71    if oldtype in ["rotmat"]:72        rotations = rotations.reshape((*rotations.shape[:-2], 3, 3))73        return rotations74    elif oldtype in ["rot6d", "6drot", "rotation6d"]:75        rotations = rotation_6d_to_matrix(rotations)76        return rotations77    elif oldtype in ["rotquat", "quaternion"]:78        rotations = quaternion_to_matrix(rotations)79        return rotations80    elif oldtype in ["rotvec", "axisangle"]:81        rotations = axis_angle_to_matrix(rotations)82        return rotations83    else:84        raise NotImplementedError85