Spaces:

IsshikiHugh
/

HSMR

Sleeping

File size: 14,923 Bytes

5ac1897

from lib.kits.basic import *

import torch
import numpy as np
import torch.nn.functional as F

from .definition import JOINTS_DEF, N_JOINTS, JID2DOF, JID2QIDS, DoF1_JIDS, DoF2_JIDS, DoF3_JIDS, DoF1_QIDS, DoF2_QIDS, DoF3_QIDS

from lib.utils.data import to_tensor
from lib.utils.geometry.rotation import (
    matrix_to_euler_angles,
    matrix_to_rotation_6d,
    euler_angles_to_matrix,
    rotation_6d_to_matrix,
)


# ====== Internal Utils ======


def axis2convention(axis:List):
    ''' [1,0,0] -> 'X', [0,1,0] -> 'Y', [0,0,1] -> 'Z' '''
    if axis == [1, 0, 0]:
        return 'X'
    elif axis == [0, 1, 0]:
        return 'Y'
    elif axis == [0, 0, 1]:
        return 'Z'
    else:
        raise ValueError(f'Unsupported axis: {axis}.')


def rotation_2d_to_angle(r2d:torch.Tensor):
    '''
    Extract single angle from a 2D rotation vector, which is the first column of a 2x2 rotation matrix.

    ### Args
    - r2d: torch.Tensor
        - shape = (...B, 2)

    ### Returns
    - shape = (...B,)
    '''
    cos, sin = r2d[..., [0]], -r2d[..., [1]]
    return torch.atan2(sin, cos)

# ====== Tools ======


OS2S_FLIP = [-1, 1, 1]
OS2S_CONV = 'YZX'
def real_orient_mat2q(orient_mat:torch.Tensor) -> torch.Tensor:
    '''
    The rotation matrix that SKEL uses is different from the SMPL's orientation matrix.
    The rotation to representation functions below can not be used to transform the rotaiton matrix.
    This function is used to convert the SMPL's orientation matrix to the SKEL's orientation q.
    BUT, is that really important? Maybe we shouldn't align SMPL's orientation with SKEL's, can they be different? 

    ### Args
    - orient_mat: torch.Tensor, shape = (..., 3, 3)

    ### Returns
    - orient_q: torch.Tensor, shape = (..., 3)
    '''
    device = orient_mat.device
    flip = to_tensor(OS2S_FLIP, device=device)  # (3,)
    orient_ea = matrix_to_euler_angles(orient_mat.clone(), OS2S_CONV)  # (..., 3)
    orient_ea = orient_ea[..., [2, 1, 0]]  # Re-permuting the order.
    orient_q = orient_ea * flip[None]
    return orient_q


def real_orient_q2mat(orient_q:torch.Tensor) -> torch.Tensor:
    '''
    The rotation matrix that SKEL uses is different from the SMPL's orientation matrix.
    The rotation to representation functions below can not be used to transform the rotation matrix.
    This function is used to convert the SKEL's orientation q to the SMPL's orientation matrix.
    BUT, is that really important? Maybe we shouldn't align SMPL's orientation with SKEL's, can they be different? 

    ### Args
    - orient_q: torch.Tensor, shape = (..., 3)
    
    ### Returns
    - orient_mat: torch.Tensor, shape = (..., 3, 3)
    '''
    device = orient_q.device
    flip = to_tensor(OS2S_FLIP, device=device)  # (3,)
    orient_ea = orient_q * flip[None]
    orient_ea = orient_ea[..., [2, 1, 0]]  # Re-permuting the order.
    orient_mat = euler_angles_to_matrix(orient_ea, OS2S_CONV)
    return orient_mat


def flip_params_lr(poses:torch.Tensor) -> torch.Tensor:
    '''
    It flips the skel through exchanging the params of left part and right part of the body. It's useful for
    data augmentation. Note that the 'left & right' defined when the body is facing z+ direction, this is
    only important for the orientation.

    ### Args
    - poses: torch.Tensor, shape = (B, L, 46) or (L, 46)

    ### Returns
    - flipped_poses: torch.Tensor, shape = (B, L, 46) or (L, 46)
    '''
    assert len(poses.shape) in [2, 3] and poses.shape[-1] == 46, f'Shape of poses should be (B, L, 46) or (L, 46) but get {poses.shape}.'

    # 1. Switch the value of each pair through re-permuting.
    flipped_perm = [
             0,  1,  2,  # pelvis
            10, 11, 12,  # femur-r -> femur-l
            13,          # tibia-r -> tibia-l
            14,          # talus-r -> talus-l
            15,          # calcn-r -> calcn-l
            16,          # toes-r -> toes-l
             3,  4,  5,  # femur-l -> femur-r
             6,          # tibia-l -> tibia-r
             7,          # talus-l -> talus-r
             8,          # calcn-l -> calcn-r
             9,          # toes-l -> toes-r
            17, 18, 19,  # lumbar
            20, 21, 22,  # thorax
            23, 24, 25,  # head
            36, 37, 38,  # scapula-r -> scapula-l
            39, 40, 41,  # humerus-r -> humerus-l
            42,          # ulna-r -> ulna-l
            43,          # radius-r -> radius-l
            44, 45,      # hand-r -> hand-l
            26, 27, 28,  # scapula-l -> scapula-r
            29, 30, 31,  # humerus-l -> humerus-r
            32,          # ulna-l -> ulna-r
            33,          # radius-l -> radius-r
            34, 35       # hand-l -> hand-r
        ]

    flipped_poses = poses[..., flipped_perm]

    # 2. Mirror the rotation direction through apply -1.
    flipped_sign = [
             1, -1, -1,  # pelvis
             1,  1,  1,  # femur-r'
             1,          # tibia-r'
             1,          # talus-r'
             1,          # calcn-r'
             1,          # toes-r'
             1,  1,  1,  # femur-l'
             1,          # tibia-l'
             1,          # talus-l'
             1,          # calcn-l'
             1,          # toes-l'
            -1,  1, -1,  # lumbar
            -1,  1, -1,  # thorax
            -1,  1, -1,  # head
            -1, -1,  1,  # scapula-r'
            -1, -1,  1,  # humerus-r'
             1,          # ulna-r'
             1,          # radius-r'
             1,  1,      # hand-r'
            -1, -1,  1,  # scapula-l'
            -1, -1,  1,  # humerus-l'
             1,          # ulna-l'
             1,          # radius-l'
             1,  1       # hand-l'
        ]
    flipped_sign = torch.tensor(flipped_sign, dtype=poses.dtype, device=poses.device)  # (46,)
    flipped_poses = flipped_sign * flipped_poses

    return flipped_poses



# def rotate_orient_around_z(q, rot):
#     """
#     Rotate SKEL orientation.
#     Args:
#         q (np.ndarray): SKEL style rotation representation (3,).
#         rot (np.ndarray): Rotation angle in degrees.
#     Returns:
#         np.ndarray: Rotated axis-angle vector.
#     """
#     import torch
#     from lib.body_models.skel.osim_rot import CustomJoint
#     # q to mat
#     root = CustomJoint(axis=[[0,0,1], [1,0,0], [0,1,0]], axis_flip=[1, 1, 1])  # pelvis
#     q = torch.from_numpy(q).unsqueeze(0)
#     q = q[:, [2, 1, 0]]
#     Rp = euler_angles_to_matrix(q, convention="YXZ")
#     # rotate around z
#     R = torch.Tensor([[np.deg2rad(-rot), 0, 0]])
#     R = axis_angle_to_matrix(R)
#     R = torch.matmul(R, Rp)
#     # mat to q
#     q = matrix_to_euler_angles(R, convention="YXZ")
#     q = q[:, [2, 1, 0]]
#     q = q.numpy().squeeze()

#     return q.astype(np.float32)


def params_q2rot(params_q:Union[torch.Tensor, np.ndarray]):
    '''
    Transform parts of the euler-like SKEL parameters representation all to rotation matrix.

    ### Args
    - params_q: Union[torch.Tensor, np.ndarray], shape = (...B, 46) or (...B, 46)

    ### Returns
    - shape = (...B, 24, 9)  # 24 joints, each joint has a 3x3 matrix, but for some joints, the matrix is not all used.
    '''
    # Check the type and unify to torch.
    is_np = isinstance(params_q, np.ndarray)
    if is_np:
        params_q = torch.from_numpy(params_q)

    # Prepare for necessary variables.
    Bs = params_q.shape[:-1]
    ident = torch.eye(3, dtype=params_q.dtype).to(params_q.device)  # (3, 3)
    params_rot = ident.repeat(*Bs, N_JOINTS, 1, 1)  # (...B, 24, 3, 3)

    # Deal with each joints separately. Modified from the `skel_model.py` but a static version.
    sid = 0
    for jid in range(N_JOINTS):
        joint_obj = JOINTS_DEF[jid]
        eid = sid + joint_obj.nb_dof.item()
        params_rot[..., jid, :, :] = joint_obj.q_to_rot(params_q[..., sid:eid])
        sid = eid

    if is_np:
        params_rot = params_rot.detach().cpu().numpy()
    return params_rot


def params_q2rep(params_q:Union[torch.Tensor, np.ndarray]):
    '''
    Transform the euler-like SKEL parameters representation to the continuous representation.
    This function is not supposed to be used in the training process, but only for debugging.
    The function that matters actually is the inverse of this function.

    ### Args
    - params_q: Union[torch.Tensor, np.ndarray], shape = (...B, 46) or (...B, 46)

    ### Returns
    - shape = (...B, 24, 6)
        - Among 24 joints, for 3 DoF ones, all 6 values are used to represent the rotation;
          but for 1 DoF joints, only the first 2 are used. The rest will be represented as zeros.
    '''
    # Check the type and unify to torch.
    is_np = isinstance(params_q, np.ndarray)
    if is_np:
        params_q = torch.from_numpy(params_q)

    # Prepare for necessary variables.
    Bs = params_q.shape[:-1]
    params_rep = params_q.new_zeros(*Bs, N_JOINTS, 6)  # (...B, 24, 6)

    # Deal with each joints separately. Modified from the `skel_model.py` but a static version.
    sid = 0
    for jid in range(N_JOINTS):
        joint_obj = JOINTS_DEF[jid]
        dof = joint_obj.nb_dof.item()
        eid = sid + dof
        if dof == 3:
            mat = joint_obj.q_to_rot(params_q[..., sid:eid])  # (...B, 3, 3)
            params_rep[..., jid, :] = matrix_to_rotation_6d(mat)  # (...B, 6)
        elif dof == 2:
            # mat = joint_obj.q_to_rot(params_q[..., sid:eid])  # (...B, 3, 3)
            # params_rep[..., jid, :] = matrix_to_rotation_6d(mat)  # (...B, 6)
            params_rep[..., jid, :2] = params_q[..., sid:eid]
        elif dof == 1:
            cos = torch.cos(params_q[..., sid])
            sin = torch.sin(params_q[..., sid])
            params_rep[..., jid, 0] = cos
            params_rep[..., jid, 1] = -sin

        sid = eid

    if is_np:
        params_rep = params_rep.detach().cpu().numpy()
    return params_rep


# Deprecated.
def dof3_to_q(rot, axises:List, flip:List):
    '''
    Convert a rotation matrix to SKEL style rotation representation.

    ### Args
    - rot: torch.Tensor, shape (...B, 3, 3)
        - The rotation matrix.
    - axises: list
        - [[x0, y0, z0], [x1, y1, z1], [x2, y2, z2]]
        - The axis defined in the SKEL's joint_dict. Only one of xi, yi, zi is 1, the others are 0.
    - flip: list
        - [f0, f1, f2]
        - The flip value defined in the SKEL's joint_dict. fi is 1 or -1.

    ### Returns
    - shape = (...B, 3)
    '''
    convention = [axis2convention(axis) for axis in reversed(axises)]  # SKEL use euler angle in reverse order
    convention = ''.join(convention)
    q = matrix_to_euler_angles(rot[..., :, :], convention=convention)  # (...B, 3)
    q = q[..., [2, 1, 0]]  # SKEL use euler angle in reverse order
    flip = rot.new_tensor(flip)  # (3,)
    q = flip * q
    return q


### Slow version, deprecated. ###
# def params_rep2q(params_rot:Union[torch.Tensor, np.ndarray]):
#     '''
#     Transform the continuous representation back to the SKEL style euler-like representation.
#
#     ### Args
#     - params_rot: Union[torch.Tensor, np.ndarray]
#         - shape = (...B, 24, 6)
#
#     ### Returns
#     - shape = (...B, 46)
#     '''
#
#     # Check the type and unify to torch.
#     is_np = isinstance(params_rot, np.ndarray)
#     if is_np:
#         params_rot = torch.from_numpy(params_rot)
#
#     # Prepare for necessary variables.
#     Bs = params_rot.shape[:-2]
#     params_q = params_rot.new_zeros((*Bs, 46))  # (...B, 46)
#
#     for jid in range(N_JOINTS):
#         joint_obj = JOINTS_DEF[jid]
#         dof = joint_obj.nb_dof.item()
#         sid, eid = JID2QIDS[jid][0], JID2QIDS[jid][-1] + 1
#
#         if dof == 3:
#             mat = rotation_6d_to_matrix(params_rot[..., jid, :])  # (...B, 3, 3)
#             params_q[..., sid:eid] = dof3_to_q(
#                     mat,
#                     joint_obj.axis.tolist(),
#                     joint_obj.axis_flip.detach().cpu().tolist(),
#                 )
#         elif dof == 2:
#             params_q[..., sid:eid] = params_rot[..., jid, :2]
#         else:
#             params_q[..., sid:eid] = rotation_2d_to_angle(params_rot[..., jid, :2])
#
#     if is_np:
#         params_q = params_q.detach().cpu().numpy()
#     return params_q

def orient_mat2q(orient_mat:torch.Tensor):
    ''' This is a tool function for inspecting only. orient_mat ~ (...B, 3, 3)'''
    poses_rep = orient_mat.new_zeros(orient_mat.shape[:-2] + (24, 6))  # (...B, 24, 6)
    orient_rep = matrix_to_rotation_6d(orient_mat)  # (...B, 6)
    poses_rep[..., 0, :] = orient_rep
    poses_q = params_rep2q(poses_rep)  # (...B, 46)
    return poses_q[..., :3]


# Pre-grouped the joints for different conventions
CON_GROUP2JIDS  = {'YXZ': [0, 1, 6], 'YZX': [11, 12, 13], 'XZY': [14, 19], 'ZYX': [15, 20]}
CON_GROUP2FLIPS = {'YXZ': [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, -1.0, -1.0]], 'YZX': [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], 'XZY': [[1.0, -1.0, -1.0], [1.0, 1.0, 1.0]], 'ZYX': [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]}
# Faster version.
def params_rep2q(params_rot:Union[torch.Tensor, np.ndarray]):
    '''
    Transform the continuous representation back to the SKEL style euler-like representation.

    ### Args
    - params_rot: Union[torch.Tensor, np.ndarray], shape = (...B, 24, 6)

    ### Returns
    - shape = (...B, 46)
    '''

    with PM.time_monitor('params_rep2q'):
        with PM.time_monitor('preprocess'):
            params_rot, recover_type_back = to_tensor(params_rot, device=None, temporary=True)

            # Prepare for necessary variables.
            Bs = params_rot.shape[:-2]
            params_q = params_rot.new_zeros((*Bs, 46))  # (...B, 46)

        with PM.time_monitor(f'dof1&dof2'):
            params_q[..., DoF1_QIDS] = rotation_2d_to_angle(params_rot[..., DoF1_JIDS, :2]).squeeze(-1)
            params_q[..., DoF2_QIDS] = params_rot[..., DoF2_JIDS, :2].reshape(*Bs, -1)  # (...B, J2=2 * 2)

        with PM.time_monitor(f'dof3'):
            dof3_6ds = params_rot[..., DoF3_JIDS, :].reshape(*Bs, len(DoF3_JIDS), 6)  # (...B, J3=10, 3, 6)
            dof3_mats = rotation_6d_to_matrix(dof3_6ds)  # (...B, J3=10, 3, 3)

            for convention, jids in CON_GROUP2JIDS.items():
                idxs = [DoF3_JIDS.index(jid) for jid in jids]
                mats = dof3_mats[..., idxs, :, :]  # (...B, J', 3, 3)
                qs = matrix_to_euler_angles(mats, convention=convention)  # (...B, J', 3)
                qs = qs[..., [2, 1, 0]]  # SKEL use euler angle in reverse order
                flips = qs.new_tensor(CON_GROUP2FLIPS[convention])  # (J', 3)
                qs = qs * flips  # (...B, J', 3)
                qids = [qid for jid in jids for qid in JID2QIDS[jid]]
                params_q[..., qids] = qs.reshape(*Bs, -1)

        with PM.time_monitor('post_process'):
            params_q = recover_type_back(params_q)
    return params_q