Code to get Twist phi about hybrik

[Jeff-sjtu]

Hi @Realdr4g0n, I think the code below may help you. The lbs function would return twist_angle. You can replace the original lbs function in SMPL with ours.

def lbs(betas, pose, v_template, shapedirs, posedirs, J_regressor, J_regressor_h36m, parents,
        lbs_weights, pose2rot=True, dtype=torch.float32):
    ''' Performs Linear Blend Skinning with the given shape and pose parameters

        Parameters
        ----------
        betas : torch.tensor BxNB
            The tensor of shape parameters
        pose : torch.tensor Bx(J + 1) * 3
            The pose parameters in axis-angle format
        v_template torch.tensor BxVx3
            The template mesh that will be deformed
        shapedirs : torch.tensor 1xNB
            The tensor of PCA shape displacements
        posedirs : torch.tensor Px(V * 3)
            The pose PCA coefficients
        J_regressor : torch.tensor JxV
            The regressor array that is used to calculate the joints from
            the position of the vertices
        parents: torch.tensor J
            The array that describes the kinematic tree for the model
        lbs_weights: torch.tensor N x V x (J + 1)
            The linear blend skinning weights that represent how much the
            rotation matrix of each part affects each vertex
        pose2rot: bool, optional
            Flag on whether to convert the input pose tensor to rotation
            matrices. The default value is True. If False, then the pose tensor
            should already contain rotation matrices and have a size of
            Bx(J + 1)x9
        dtype: torch.dtype, optional

        Returns
        -------
        verts: torch.tensor BxVx3
            The vertices of the mesh after applying the shape and pose
            displacements.
        joints: torch.tensor BxJx3
            The joints of the model
        rot_mats: torch.tensor BxJx3x3
            The rotation matrics of each joints
    '''
    batch_size = max(betas.shape[0], pose.shape[0])
    device = betas.device

    # Add shape contribution
    v_shaped = v_template + blend_shapes(betas, shapedirs)

    # Get the joints
    # NxJx3 array
    J = vertices2joints(J_regressor, v_shaped)

    # 3. Add pose blend shapes
    # N x J x 3 x 3
    ident = torch.eye(3, dtype=dtype, device=device)
    if pose2rot:
        if pose.numel() == batch_size * 24 * 4:
            rot_mats = quat_to_rotmat(pose.reshape(batch_size * 24, 4)).reshape(batch_size, 24, 3, 3)
        else:
            rot_mats = batch_rodrigues(
                pose.view(-1, 3), dtype=dtype).view([batch_size, -1, 3, 3])

        pose_feature = (rot_mats[:, 1:, :, :] - ident).view([batch_size, -1])
        # (N x P) x (P, V * 3) -> N x V x 3
        pose_offsets = torch.matmul(pose_feature, posedirs) \
            .view(batch_size, -1, 3)
    else:
        pose_feature = pose[:, 1:].view(batch_size, -1, 3, 3) - ident
        rot_mats = pose.view(batch_size, -1, 3, 3)

        pose_offsets = torch.matmul(pose_feature.view(batch_size, -1),
                                    posedirs).view(batch_size, -1, 3)

    v_posed = pose_offsets + v_shaped
    # 4. Get the global joint location
    J_transformed, A = batch_rigid_transform(rot_mats, J.clone(), parents, dtype=dtype)

    twist_angle = get_twist(rot_mats, J.clone(), parents)

    # 5. Do skinning:
    # W is N x V x (J + 1)
    W = lbs_weights.unsqueeze(dim=0).expand([batch_size, -1, -1])
    # (N x V x (J + 1)) x (N x (J + 1) x 16)
    num_joints = J_regressor.shape[0]
    T = torch.matmul(W, A.view(batch_size, num_joints, 16)) \
        .view(batch_size, -1, 4, 4)

    homogen_coord = torch.ones([batch_size, v_posed.shape[1], 1],
                               dtype=dtype, device=device)
    v_posed_homo = torch.cat([v_posed, homogen_coord], dim=2)
    v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, dim=-1))

    verts = v_homo[:, :, :3, 0]

    J_from_verts = vertices2joints(J_regressor_h36m, verts)

    return verts, J_transformed, rot_mats, J_from_verts, twist_angle

def get_twist(rot_mats, joints, parents):
    joints = torch.unsqueeze(joints, dim=-1)
    rel_joints = joints.clone()
    rel_joints[:, 1:] -= joints[:, parents[1:]].clone()

    # modified by xuchao
    childs = -torch.ones((parents.shape[0]), dtype=parents.dtype, device=parents.device)
    for i in range(1, parents.shape[0]):
        childs[parents[i]] = i

    dtype = rot_mats.dtype
    batch_size = rot_mats.shape[0]
    device = rot_mats.device

    angle_twist = []
    error = False
    for i in range(1, parents.shape[0]):
        # modified by xuchao
        if childs[i] < 0:
            angle_twist.append(torch.zeros((batch_size, 1), dtype=rot_mats.dtype, device=rot_mats.device))
            continue

        u = rel_joints[:, childs[i]]
        rot = rot_mats[:, i]

        v = torch.matmul(rot, u)

        u_norm = torch.norm(u, dim=1, keepdim=True)
        v_norm = torch.norm(v, dim=1, keepdim=True)

        axis = torch.cross(u, v, dim=1)
        axis_norm = torch.norm(axis, dim=1, keepdim=True)

        # (B, 1, 1)
        cos = torch.sum(u * v, dim=1, keepdim=True) / (u_norm * v_norm + 1e-8)
        sin = axis_norm / (u_norm * v_norm + 1e-8)

        # (B, 3, 1)
        axis = axis / (axis_norm + 1e-8)

        # Convert location revolve to rot_mat by rodrigues
        # (B, 1, 1)
        rx, ry, rz = torch.split(axis, 1, dim=1)
        zeros = torch.zeros((batch_size, 1, 1), dtype=dtype, device=device)

        K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
            .view((batch_size, 3, 3))
        ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
        rot_mat_pivot = ident + sin * K + (1 - cos) * torch.bmm(K, K)

        rot_mat_twist = torch.matmul(rot_mat_pivot.transpose(1, 2), rot)
        _, axis, angle = rotmat_to_aa(rot_mat_twist)

        axis = axis / torch.norm(axis, dim=1, keepdim=True)
        spin_axis = u / u_norm
        spin_axis = spin_axis.squeeze(-1)

        pos = torch.norm(spin_axis - axis, dim=1)
        neg = torch.norm(spin_axis + axis, dim=1)
        if float(neg) < float(pos):
            try:
                assert float(pos) > 1.9, (pos, neg)
                angle_twist.append(-1 * angle)
            except AssertionError:
                angle_twist.append(torch.ones_like(angle) * -999)
                error = True
        else:
            try:
                assert float(neg) > 1.9, (pos, neg, axis, angle, rot_mat_twist)
                angle_twist.append(angle)
            except AssertionError:
                angle_twist.append(torch.ones_like(angle) * -999)
                error = True

    angle_twist = torch.stack(angle_twist, dim=1)
    if error:
        angle_twist[:] = -999

    return angle_twist

def rotmat_to_aa(rotmat):
    batch_size = rotmat.shape[0]

    r11 = rotmat[:, 0, 0]
    r22 = rotmat[:, 1, 1]
    r33 = rotmat[:, 2, 2]

    r12 = rotmat[:, 0, 1]
    r21 = rotmat[:, 1, 0]
    r13 = rotmat[:, 0, 2]
    r31 = rotmat[:, 2, 0]
    r23 = rotmat[:, 1, 2]
    r32 = rotmat[:, 2, 1]
    
    angle = torch.acos((r11 + r22 + r33 - 1) / 2).unsqueeze(dim=1)
    '''
    if -1e-6 < (r11 + r22 + r33 - 1) / 2 + 1 < 1e-6:
        angle = torch.acos(-torch.ones_like(r11)).unsqueeze(dim=1)
    else:
        angle = torch.acos((r11 + r22 + r33 - 1) / 2).unsqueeze(dim=1)
    '''
    axis = torch.zeros((batch_size, 3))
    axis[:, 0] = r32 - r23
    axis[:, 1] = r13 - r31
    axis[:, 2] = r21 - r12
    axis = axis / (2 * torch.sin(angle) + 1e-8)

    aa = axis * angle
    return aa, axis, angle

[heathentw]

@dldaisy @Jeff-sjtu is the get_twist support batch style? the float(neg) < float(pos) seems to cause problem......

[Realdr4g0n]

This code is exactly what I want !!

Thank you very much!

[wangzheallen]

I use the code above and only got 18 twist angle while the GT kinematic rotation is with 24 rotation!
may I ask why? do you suppose the 3 bones connected with root joint and 3 leaf joints are without twist rotation?
additionally, are the 3 bones connected with the root joint have the same rotation as the global orientation?

[Jeff-sjtu]

In the SMPL format, the ith rotation indicates the relative rotation of the children of the ith joint w.r.t the ith joint. For example, the first rotation is the global rotation, and the first joint is the root joint. There are three joints connecting to the root joint. Therefore, the first rotation matrix indicates the rotation of the 3 children w.r.t to the root joint. For the leaf joints, since they have no child, the leaf rotation won't affect any joint positions. And thus we can't calculate the twist angle according to the joint positions.

[wangzheallen]

Thanks for the clarification!
I read here

HybrIK/hybrik/models/layers/smpl/lbs.py

Line 351 in ff6e61d

leaf_number = [411, 2445, 5905, 3216, 6617]

you concat another 5 ‘distal vertices’ to the leaf joints, does this operation makes too much performance (MPJPE or IK process) difference compared with only 24 joints setting?

[Jeff-sjtu]

Hi, @wangzheallen

We follow the standard settings for setting, i.e. only evaluate the 14 LSP joints. The extra leaf joints are only used to estimate the leaf rotations.

[wangzheallen]

Thanks for the clarification, it helps to understand!

[dldaisy]

Hi, @Jeff-sjtu
I have some questions about your reply #37 (comment), which is also relevant to several doubts regarding your IK code in models/layers/smpl/lbs.py.
As you mentioned:

In the SMPL format, the ith rotation indicates the relative rotation of the children of the ith joint w.r.t the ith joint. For example, the first rotation is the global rotation, and the first joint is the root joint. There are three joints connecting to the root joint. Therefore, the first rotation matrix indicates the rotation of the 3 children w.r.t to the root joint.

However, according to the smpl code, the ith rotation should indicates the rotation of the ith joint w.r.t its parent, except for the 0(root) joint. Thus, according to my understanding, the first rotation matrix, i.e., global orient, should mean the direction of pelvis instead of the rotation of the 3 children w.r.t to the root joint. Besides, in your hybrik function in models/layers/smpl/lbs.py, you also seemed to define rot_mat the same way: when a joint i has 3 children, the rot_mat[i] is the rotation matrix of its 3 children w.r.t to i solved by SVD. This way of defining rot_mat using its children instead of its parents also make me confused about your explanation of the leaf joints:

For the leaf joints, since they have no child, the leaf rotation won't affect any joint positions. And thus we can't calculate the twist angle according to the joint positions.

However, if the rot_mat means rotation w.r.t parent, the swing of a leaf node means the swing of the tip bone of a skeleton, which can be calculated through the position of the leaf node and its parent. In your models/layers/smpl/lbs.py code, you also directly predict the rot_mat of leaf nodes through neural networks instead of analytically calculate them.
May I know the reason why you define rot_mat using children instead of parent?
Thanks a lot!

[pangyyyyy]

Hi @Jeff-sjtu , thanks for your great work! I have some questions following this issue #37

I adopted your get_twist function and modified the SMPL_layer to get angle_twist from "ground-truth" betas and thetas (obtained from parsed h36m train annotations file downloaded from your GDrive link). However, my outputs are quite different from the "ground-truth" angle_twist provided in your annotations file.

Would you mind confirming if the g-t angle_twist is obtained using the above get_twist function? If that's not the method, could you also provide the code to derive g-t angle_twist?

Hope to get your advice on this, thank you so much!

[Jeff-sjtu]

Hi, @Jeff-sjtu
I have some questions about your reply #37 (comment), which is also relevant to several doubts regarding your IK code in models/layers/smpl/lbs.py.
As you mentioned:

In the SMPL format, the ith rotation indicates the relative rotation of the children of the ith joint w.r.t the ith joint. For example, the first rotation is the global rotation, and the first joint is the root joint. There are three joints connecting to the root joint. Therefore, the first rotation matrix indicates the rotation of the 3 children w.r.t to the root joint.

However, according to the smpl code, the ith rotation should indicates the rotation of the ith joint w.r.t its parent, except for the 0(root) joint. Thus, according to my understanding, the first rotation matrix, i.e., global orient, should mean the direction of pelvis instead of the rotation of the 3 children w.r.t to the root joint. Besides, in your hybrik function in models/layers/smpl/lbs.py, you also seemed to define rot_mat the same way: when a joint i has 3 children, the rot_mat[i] is the rotation matrix of its 3 children w.r.t to i solved by SVD. This way of defining rot_mat using its children instead of its parents also make me confused about your explanation of the leaf joints:

For the leaf joints, since they have no child, the leaf rotation won't affect any joint positions. And thus we can't calculate the twist angle according to the joint positions.

However, if the rot_mat means rotation w.r.t parent, the swing of a leaf node means the swing of the tip bone of a skeleton, which can be calculated through the position of the leaf node and its parent. In your models/layers/smpl/lbs.py code, you also directly predict the rot_mat of leaf nodes through neural networks instead of analytically calculate them.
May I know the reason why you define rot_mat using children instead of parent?
Thanks a lot!

Hi @dldaisy. When I saying the ith joints, it includes the root joint (the pelvis). So in this context, the SMPL model has 24 joints in total. And in the official implementation of SMPL, the i rotation is used to represent the rotation of the child of the ith joint w.r.t the ith joint. You can refer to the implementation of smplx.

[dldaisy]

Hi @Jeff-sjtu, Thanks for reply!
I just looked at the smplx implementation but is still a bit confused. If the i rotation is used to represent the rotation of the child of the ith joint w.r.t the ith joint, which child is used to calculate the rot_mat of SPINE_3 in smpl? I don't find the code in smplx that exclusively deal with joints with multiple children like yours. Besides, if the rot_mats means rotation of the child of the a joint w.r.t it, what does rot_mat physically mean to leaf nodes in smpl ?

[Jeff-sjtu]

Hi @pangyyyyy,

The code pasted before has some bugs. You need to replace the get_twist function with the correct one.

def get_twist(rot_mats, joints, parents):
    joints = torch.unsqueeze(joints, dim=-1)
    rel_joints = joints.clone()
    rel_joints[:, 1:] -= joints[:, parents[1:]].clone()

    # modified by xuchao
    childs = -torch.ones((parents.shape[0]), dtype=parents.dtype, device=parents.device)
    for i in range(1, parents.shape[0]):
        childs[parents[i]] = i

    dtype = rot_mats.dtype
    batch_size = rot_mats.shape[0]
    device = rot_mats.device

    angle_twist = []
    error = False
    for i in range(1, parents.shape[0]):
        # modified by xuchao
        if childs[i] < 0:
            angle_twist.append(torch.zeros((batch_size, 1), dtype=rot_mats.dtype, device=rot_mats.device))
            continue

        u = rel_joints[:, childs[i]]
        rot = rot_mats[:, i]

        v = torch.matmul(rot, u)

        u_norm = torch.norm(u, dim=1, keepdim=True)
        v_norm = torch.norm(v, dim=1, keepdim=True)

        axis = torch.cross(u, v, dim=1)
        axis_norm = torch.norm(axis, dim=1, keepdim=True)

        # (B, 1, 1)
        cos = torch.sum(u * v, dim=1, keepdim=True) / (u_norm * v_norm + 1e-8)
        sin = axis_norm / (u_norm * v_norm + 1e-8)

        # (B, 3, 1)
        axis = axis / (axis_norm + 1e-8)

        # Convert location revolve to rot_mat by rodrigues
        # (B, 1, 1)
        rx, ry, rz = torch.split(axis, 1, dim=1)
        zeros = torch.zeros((batch_size, 1, 1), dtype=dtype, device=device)

        K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
            .view((batch_size, 3, 3))
        ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
        rot_mat_pivot = ident + sin * K + (1 - cos) * torch.bmm(K, K)

        rot_mat_twist = torch.matmul(rot_mat_pivot.transpose(1, 2), rot)
        _, axis, angle = rotmat_to_aa(rot_mat_twist)

        axis = axis / torch.norm(axis, dim=1, keepdim=True)
        spin_axis = u / u_norm
        spin_axis = spin_axis.squeeze(-1)

        pos = torch.norm(spin_axis - axis, dim=1)
        neg = torch.norm(spin_axis + axis, dim=1)
        if float(neg) < float(pos):
            try:
                assert float(pos) > 1.9, (pos, neg)
                angle_twist.append(-1 * angle)
            except AssertionError:
                angle_twist.append(torch.ones_like(angle) * -999)
                error = True
        else:
            try:
                assert float(neg) > 1.9, (pos, neg, axis, angle, rot_mat_twist)
                angle_twist.append(angle)
            except AssertionError:
                angle_twist.append(torch.ones_like(angle) * -999)
                error = True

    angle_twist = torch.stack(angle_twist, dim=1)
    if error:
        angle_twist[:] = -999

    return angle_twist

[Jeff-sjtu]

Hi @Jeff-sjtu, Thanks for reply!
I just looked at the smplx implementation but is still a bit confused. If the i rotation is used to represent the rotation of the child of the ith joint w.r.t the ith joint, which child is used to calculate the rot_mat of SPINE_3 in smpl? I don't find the code in smplx that exclusively deal with joints with multiple children like yours. Besides, if the rot_mats means rotation of the child of the a joint w.r.t it, what does rot_mat physically mean to leaf nodes in smpl ?

In forward kinematics, we don't need to calculate rotmat based on the child. In inverse kinematics, we have two options. The first one is using three children and solve the rotmat by SVD. The second one is to pick one child to calculate.

For your second question, the rotmat of the leaf joints won't affect other joints, but will affect the distal vertices in the body mesh.

[pangyyyyy]

@Jeff-sjtu the fix works! I am now able to obtain the correct angle_twist as the ones provided in your g-t annotation files. Thank you so much for your help!

[Jeff-sjtu]

You're welcome ;) @dldaisy

[dldaisy]

Hi @Jeff-sjtu,
Thanks again for your patient reply!
Still a question:
May I ask what rot_mat[SPINE_3] in smpl FK means? If rot_mat means the rotation of SPINE_3's children w.r.t SPINE_3, does it indicate that the 3 children of SPINE_3 all have the same rotation w.r.t SPINE_3? But that's strange. This confuses me.

[Jeff-sjtu]

Hi @Jeff-sjtu,
Thanks again for your patient reply!
Still a question:
May I ask what rot_mat[SPINE_3] in smpl FK means? If rot_mat means the rotation of SPINE_3's children w.r.t SPINE_3, does it indicate that the 3 children of SPINE_3 all have the same rotation w.r.t SPINE_3? But that's strange. This confuses me.

Yes, the 3 children of SPINE_3 share the same rotation. In SMPL, the body part with multiple children is treated as a rigid body. Similar to SPINE_3, the 3 children of the pelvis joint also form a rigid body.

[dldaisy]

Yes, the 3 children of SPINE_3 share the same rotation. In SMPL, the body part with multiple children is treated as a rigid body. Similar to SPINE_3, the 3 children of the pelvis joint also form a rigid body.

I think you're right after reading the smplx code carefully. Thanks a lot!

[dldaisy]

Hi @Jeff-sjtu ,
I have another question regarding the code get_twist. It seems that this code assumes there's no twist of pelvis' 3 children regards to pelvis, but the twist of spine_3's children w.r.t spine_3 can be calculated. Does this go against the assumption that spine_3 and its 3 children forms a rigid body? Thanks!

[Jeff-sjtu]

Hi @Jeff-sjtu ,
I have another question regarding the code get_twist. It seems that this code assumes there's no twist of pelvis' 3 children regards to pelvis, but the twist of spine_3's children w.r.t spine_3 can be calculated. Does this go against the assumption that spine_3 and its 3 children forms a rigid body? Thanks!

The twist rotation depends on the child you choose. Even the 3 children share the same rotation, their twist angle will be different, since the swing rotations are different.

[dldaisy]

I see. Thanks a lot!

[dldaisy]

Hi @Jeff-sjtu,
I have a question on this line of the function get_twist: assert float(pos) > 1.9, (pos, neg). Why should this assertion be made to decide whether a twist angle is wrong? Thanks in advance!

[dldaisy]

@heathentw yes, I remember I made some modifications on this part.

[EveningLin]

I use the code above and only got 18 twist angle while the GT kinematic rotation is with 24 rotation! may I ask why? do you suppose the 3 bones connected with root joint and 3 leaf joints are without twist rotation? additionally, are the 3 bones connected with the root joint have the same rotation as the global orientation?

@Jeff-sjtu @wangzheallen 是不是可以理解成6个为0都是因为叶节点不能计算twist？至于为什么只得到 (24, 3)个（包含叶节点），就是因为骨盆节点也不能进行计算？
Can it be understood that 6 are 0 because the leaf node cannot calculate twist? As for why only (24, 3) (including leaf nodes) can be obtained, because the pelvic nodes can not be calculated?
twist_angle tensor([[[-0.2089], [ 0.0958], [ 0.0059], [-0.0613], [ 0.0255], [-0.0196], [-0.2302], [ 0.1953], [-0.0551], [ 0.0000], [ 0.0000], [ 0.0477], [-0.0429], [ 0.0555], [ 0.0000], [-0.1610], [ 0.1090], [-0.0854], [ 0.0906], [-0.0711], [-0.0739], [ 0.0000], [ 0.0000]]])