quaternion_to_aa function

[JianqiangRen]

I want to get SMPL theta from pred_theta_mats, I guess the pred_theta_mats is quaternion format, so could you please provide the quaternion_to_aa function code?

[Jeff-sjtu]

Hi @JianqiangRen, the following code might be helpful to you.

def quaternion_to_aa(quaternion):
    '''
    input quaternion: [B, 96]
    return axis-angle: [B, 72]
    '''
    batch_size = quaternion.shape[0]
    quaternion = quaternion.reshape(batch_size, 24, 4)
    angle = np.arccos(quaternion[:, :, 0:1]) * 2
    axis = quaternion[:, :, 1:] / \
        np.linalg.norm(quaternion[:, :, 1:], axis=2, keepdims=True)

    aa = angle * axis
    aa = aa.reshape(batch_size, -1)
    return aa