Python quaternion.as_rotation_matrix() Examples

def test_as_rotation_matrix(Rs):
    def quat_mat(quat):
        return np.array([(quat * v * quat.inverse()).vec for v in [quaternion.x, quaternion.y, quaternion.z]]).T

    def quat_mat_vec(quats):
        mat_vec = np.array([quaternion.as_float_array(quats * v * np.reciprocal(quats))[..., 1:]
                            for v in [quaternion.x, quaternion.y, quaternion.z]])
        return np.transpose(mat_vec, tuple(range(mat_vec.ndim))[1:-1]+(-1, 0))

    with pytest.raises(ZeroDivisionError):
        quaternion.as_rotation_matrix(quaternion.zero)

    for R in Rs:
        # Test correctly normalized rotors:
        assert allclose(quat_mat(R), quaternion.as_rotation_matrix(R), atol=2*eps)
        # Test incorrectly normalized rotors:
        assert allclose(quat_mat(R), quaternion.as_rotation_matrix(1.1*R), atol=2*eps)

    Rs0 = Rs.copy()
    Rs0[Rs.shape[0]//2] = quaternion.zero
    with pytest.raises(ZeroDivisionError):
        quaternion.as_rotation_matrix(Rs0)

    # Test correctly normalized rotors:
    assert allclose(quat_mat_vec(Rs), quaternion.as_rotation_matrix(Rs), atol=2*eps)
    # Test incorrectly normalized rotors:
    assert allclose(quat_mat_vec(Rs), quaternion.as_rotation_matrix(1.1*Rs), atol=2*eps)

    # Simply test that this function succeeds and returns the right shape
    assert quaternion.as_rotation_matrix(Rs.reshape((2, 5, 10))).shape == (2, 5, 10, 3, 3) 

def test_from_rotation_matrix(Rs):
    try:
        from scipy import linalg
        have_linalg = True
    except ImportError:
        have_linalg = False

    for nonorthogonal in [True, False]:
        if nonorthogonal and have_linalg:
            rot_mat_eps = 10*eps
        else:
            rot_mat_eps = 5*eps
        for i, R1 in enumerate(Rs):
            R2 = quaternion.from_rotation_matrix(quaternion.as_rotation_matrix(R1), nonorthogonal=nonorthogonal)
            d = quaternion.rotation_intrinsic_distance(R1, R2)
            assert d < rot_mat_eps, (i, R1, R2, d)  # Can't use allclose here; we don't care about rotor sign

        Rs2 = quaternion.from_rotation_matrix(quaternion.as_rotation_matrix(Rs), nonorthogonal=nonorthogonal)
        for R1, R2 in zip(Rs, Rs2):
            d = quaternion.rotation_intrinsic_distance(R1, R2)
            assert d < rot_mat_eps, (R1, R2, d)  # Can't use allclose here; we don't care about rotor sign

        Rs3 = Rs.reshape((2, 5, 10))
        Rs4 = quaternion.from_rotation_matrix(quaternion.as_rotation_matrix(Rs3))
        for R3, R4 in zip(Rs3.flatten(), Rs4.flatten()):
            d = quaternion.rotation_intrinsic_distance(R3, R4)
            assert d < rot_mat_eps, (R3, R4, d)  # Can't use allclose here; we don't care about rotor sign 

def get_mat(translation=None, rotation=None):
    assert False
    mat = np.eye(4)
    if translation is not None:
        mat[:3, 3] = translation
    if rotation is not None:
        assert False  # Does not work?
        mat[:3, :3] = quaternion.as_rotation_matrix(np.quaternion(*rotation))
    return mat 

def __init__(self, polar_dist_range):
        #  self.angle = np.random.uniform(-np.pi, np.pi)
        self.angle = np.random.uniform(-np.pi, np.pi)
        self.velocity = np.random.uniform(0, 1)
        self.translation = np.array([np.sin(self.angle), np.cos(self.angle), 0]) * self.velocity
        #  # self.q = rand_quat()
        #  self.q = rand_quat_planar()
        self.rel_angle = rand_angle() / 2.0
        self.q = quaternion.from_rotation_vector(np.array([0, 0, 1.]) * self.rel_angle)

        polar_angle = np.random.uniform(-np.pi, np.pi)
        polar_distance = np.random.uniform(*polar_dist_range)
        self.start_position = np.array([np.sin(polar_angle), np.cos(polar_angle), 0]) * polar_distance
        #  self.q_start = rand_quat_planar()
        self.start_angle = rand_angle()
        self.q_start = quaternion.from_rotation_vector(np.array([0, 0, 1.]) * self.rel_angle)

        self.end_position = self.start_position + self.translation
        self.end_angle = self.start_angle + self.rel_angle
        #  self.q_end = self.q_start * self.q
        self.q_end = quaternion.from_rotation_vector(np.array([0, 0, 1.]) * self.end_angle)

        #  self.transform_start = np.eye(4)
        #  self.transform_start[:3,:3] = quaternion.as_rotation_matrix(self.q_start)
        #  self.transform_start[:3,3] = self.start_position
        self.transform_start = get_mat_angle(self.start_position, self.start_angle)

        #  self.rel_transform = np.eye(4)
        #  self.rel_transform[:3,:3] = quaternion.as_rotation_matrix(self.q)
        #  self.rel_transform[:3,3] = self.translation
        self.rel_transform = get_mat_angle(self.translation, self.rel_angle)

        #  self.transform_end = np.eye(4)
        #  self.transform_end[:3,3] = self.end_position
        #  self.transform_end[:3,:3] = quaternion.as_rotation_matrix(self.q_end)
        #  self.transform_end = np.matmul(self.rel_transform, self.transform_start)
        self.transform_end = get_mat_angle(self.end_position, self.end_angle) 

def compute_quat(self, predictions, targets, reduce_fn="mean"):
        """
        Compute the chosen metrics. Predictions and targets are assumed to be quaternions.
        Args:
            predictions: An np array of shape (n, seq_length, n_joints*4)
            targets: An np array of the same shape as `predictions`
            reduce_fn: Which reduce function to apply to the joint dimension, if applicable. Choices are [mean, sum].

        Returns:
            A dictionary {metric_name -> values} where the values are given per batch entry and frame as an np array
            of shape (n, seq_length). `reduce_fn` is only applied to metrics where it makes sense, i.e. not to PCK
            and euler angle differences.
        """
        assert predictions.shape[-1] % 4 == 0, "predictions are not quaternions"
        assert targets.shape[-1] % 4 == 0, "targets are not quaternions"
        assert reduce_fn in ["mean", "sum"]
        assert not self._should_call_reset, "you should reset the state of this class after calling `finalize`"
        dof = 4
        batch_size = predictions.shape[0]
        seq_length = predictions.shape[1]

        # for simplicity we just convert quaternions to rotation matrices
        pred_q = quaternion.from_float_array(np.reshape(predictions, [batch_size, seq_length, -1, dof]))
        targ_q = quaternion.from_float_array(np.reshape(targets, [batch_size, seq_length, -1, dof]))
        pred_rots = quaternion.as_rotation_matrix(pred_q)
        targ_rots = quaternion.as_rotation_matrix(targ_q)

        preds = np.reshape(pred_rots, [batch_size, seq_length, -1])
        targs = np.reshape(targ_rots, [batch_size, seq_length, -1])
        return self.compute_rotmat(preds, targs, reduce_fn) 

def from_quat(self, joint_angles):
        """
        Get joint positions from quaternion representations in shape (N, H36M_NR_JOINTS*4)
        """
        qs = quaternion.from_float_array(np.reshape(joint_angles, [-1, H36M_NR_JOINTS, 4]))
        aa = quaternion.as_rotation_matrix(qs)
        return self.fk(np.reshape(aa, [-1, H36M_NR_JOINTS * 3])) 

def visualize_quat(self, seed, prediction, target, title):
        assert seed.shape[-1] == prediction.shape[-1] == target.shape[-1] == self.expected_n_input_joints * 4
        assert prediction.shape[0] == target.shape[0]
        dof = 4

        def _to_rotmat(x):
            b = x.shape[0]
            xq = quaternion.from_float_array(np.reshape(x, [b, -1, dof]))
            xr = quaternion.as_rotation_matrix(xq)
            return np.reshape(xr, [b, -1])

        self.visualize_rotmat(_to_rotmat(seed), _to_rotmat(prediction), _to_rotmat(target), title) 

def local_rot_to_global(joint_angles, parents, rep="rotmat", left_mult=False):
    """
    Converts local rotations into global rotations by "unrolling" the kinematic chain.
    Args:
        joint_angles: An np array of rotation matrices of shape (N, nr_joints*dof)
        parents: A np array specifying the parent for each joint
        rep: Which representation is used for `joint_angles`
        left_mult: If True the local matrix is multiplied from the left, rather than the right

    Returns:
        The global rotations as an np array of rotation matrices in format (N, nr_joints, 3, 3)
    """
    assert rep in ["rotmat", "quat", "aa"]
    n_joints = len(parents)
    if rep == "rotmat":
        rots = np.reshape(joint_angles, [-1, n_joints, 3, 3])
    elif rep == "quat":
        rots = quaternion.as_rotation_matrix(quaternion.from_float_array(
            np.reshape(joint_angles, [-1, n_joints, 4])))
    else:
        rots = quaternion.as_rotation_matrix(quaternion.from_rotation_vector(
            np.reshape(joint_angles, [-1, n_joints, 3])))

    out = np.zeros_like(rots)
    dof = rots.shape[-3]
    for j in range(dof):
        if parents[j] < 0:
            # root rotation
            out[..., j, :, :] = rots[..., j, :, :]
        else:
            parent_rot = out[..., parents[j], :, :]
            local_rot = rots[..., j, :, :]
            lm = local_rot if left_mult else parent_rot
            rm = parent_rot if left_mult else local_rot
            out[..., j, :, :] = np.matmul(lm, rm)
    return out 

def orientation_to_angles(ori):
    """
    Covert an array of quaternions to an array of Euler angles. Calculations are from Android source code:

    https://android.googlesource.com/platform/frameworks/base/+/master/core/java/android/hardware/SensorManager.java
    Function "getOrientation(float[] R, float[] values)"

    Note that this function DOES NOT consider singular configurations, such as Gimbal Lock.

    Args:
        ori: an array of N quaternions.

    Returns:
        A Nx3 array. With Android's game rotation vector or rotation vector, each group of three values
        correspond to: azimuth(yaw), pitch and roll, respectively.
    """
    if ori.dtype != quaternion.quaternion:
        ori = quaternion.from_float_array(ori)

    rm = quaternion.as_rotation_matrix(ori)
    angles = np.zeros([ori.shape[0], 3])
    angles[:, 0] = adjust_angle_array(np.arctan2(rm[:, 0, 1], rm[:, 1, 1]))
    angles[:, 1] = adjust_angle_array(np.arcsin(-rm[:, 2, 1]))
    angles[:, 2] = adjust_angle_array(np.arctan2(-rm[:, 2, 0], rm[:, 2, 2]))

    return angles 

def __call__(self, feat, targ, **kwargs):
        rv = np.random.random(3)
        na = np.linalg.norm(rv)
        if na < 1e-06:
            return feat

        angle = np.random.random() * self.max_angle * math.pi / 180
        rv = rv / na * math.sin(angle / 2.0)
        rot = quaternion.as_rotation_matrix(quaternion.from_rotation_vector(rv))
        rows = feat.shape[0]
        feat = np.matmul(rot, feat.reshape([-1, 3]).T).T
        return feat.reshape([rows, -1]), targ 

def aa_to_rot_matrix(data):
    """
    Converts the orientation data to represent angle axis as rotation matrices. `data` is expected in format
    (seq_length, n*3). Returns an array of shape (seq_length, n*9).
    """
    # reshape to have sensor values explicit
    data_c = np.array(data, copy=True)
    seq_length, n = data_c.shape[0], data_c.shape[1] // 3
    data_r = np.reshape(data_c, [seq_length, n, 3])

    qs = quaternion.from_rotation_vector(data_r)
    rot = np.reshape(quaternion.as_rotation_matrix(qs), [seq_length, n, 9])

    return np.reshape(rot, [seq_length, 9*n]) 

def compute_metrics(prediction, target, compute_positional_error=False):
    """
    Compute the metrics on the predictions. The function can handle variable sequence lengths for each pair of
    prediction-target array. The pose parameters can either be represented as rotation matrices (dof = 9) or
    quaternions (dof = 4)
    :param prediction: a list of np arrays of size (seq_length, 24*dof)
    :param target: a list of np arrays of size (seq_length, 24*dof)
    :param compute_positional_error: if set, the euclidean pose error is calculated which can take some time.
    """
    assert len(prediction) == len(target)
    dof = prediction[0].shape[1] // SMPL_NR_JOINTS
    assert dof == 9 or dof == 4

    # because we are interested in difference per frame, flatten inputs
    pred = np.concatenate(prediction, axis=0)
    targ = np.concatenate(target, axis=0)

    if dof == 4:
        def to_rot(x):
            seq_length = x.shape[0]
            x_ = np.reshape(x, [seq_length, -1, dof])
            x_ = quaternion.as_rotation_matrix(quaternion.as_quat_array(x_))
            return np.reshape(x_, [seq_length, -1])

        # convert quaternions to rotation matrices
        pred = to_rot(pred)
        targ = to_rot(targ)

    pred_g = smpl_rot_to_global(pred)
    targ_g = smpl_rot_to_global(targ)

    angles = joint_angle_error(pred_g, targ_g)
    mm = np.zeros([1, SMPL_NR_JOINTS])  # Ignore positional loss in Python 3.

    return angles, mm 

def write_ply_to_file(path, position, orientation, acceleration=None,
                      global_rotation=np.identity(3, float), local_axis=None,
                      trajectory_color=None, num_axis=3,
                      length=1.0, kpoints=100, interval=100):
    """
    Visualize camera trajectory as ply file
    :param path: path to save
    :param position: Nx3 array of positions
    :param orientation: Nx4 array or orientation as quaternion
    :param acceleration: (optional) Nx3 array of acceleration
    :param global_rotation: (optional) global rotation
    :param local_axis: (optional) local axis vector
    :param trajectory_color: (optional) the color of the trajectory. The default is [255, 0, 0] (red)
    :return: None
    """
    num_cams = position.shape[0]
    assert orientation.shape[0] == num_cams

    max_acceleration = 1.0
    if acceleration is not None:
        assert acceleration.shape[0] == num_cams
        max_acceleration = max(np.linalg.norm(acceleration, axis=1))
        print('max_acceleration: ', max_acceleration)
        num_axis = 4

    sample_pt = np.arange(0, num_cams, interval, dtype=int)
    num_sample = sample_pt.shape[0]

    # Define the optional transformation. Default is set w.r.t tango coordinate system
    position_transformed = np.matmul(global_rotation, np.array(position).transpose()).transpose()
    if local_axis is None:
        local_axis = np.array([[1.0, 0.0, 0.0, 0.0],
                               [0.0, 1.0, 0.0, 0.0],
                               [0.0, 0.0, 1.0, 0.0]])
    if trajectory_color is None:
        trajectory_color = [0, 255, 255]
    axis_color = [[255, 0, 0], [0, 255, 0], [0, 0, 255], [255, 0, 255]]
    vertex_type = [('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('red', 'u1'), ('green', 'u1'), ('blue', 'u1')]
    positions_data = np.empty((position_transformed.shape[0],), dtype=vertex_type)
    positions_data[:] = [tuple([*i, *trajectory_color]) for i in position_transformed]

    app_vertex = np.empty([num_axis * kpoints], dtype=vertex_type)
    for i in range(num_sample):
        q = quaternion.quaternion(*orientation[sample_pt[i]])
        if acceleration is not None:
            local_axis[:, -1] = acceleration[sample_pt[i]].flatten() / max_acceleration
        global_axes = np.matmul(global_rotation, np.matmul(quaternion.as_rotation_matrix(q), local_axis))
        for k in range(num_axis):
            for j in range(kpoints):
                axes_pts = position_transformed[sample_pt[i]].flatten() +\
                           global_axes[:, k].flatten() * j * length / kpoints
                app_vertex[k*kpoints + j] = tuple([*axes_pts, *axis_color[k]])

        positions_data = np.concatenate([positions_data, app_vertex], axis=0)
    vertex_element = plyfile.PlyElement.describe(positions_data, 'vertex')
    plyfile.PlyData([vertex_element], text=True).write(path)