lib/pymafx/models/smpl.py

# This script is extended based on https://github.com/nkolot/SPIN/blob/master/models/smpl.py

import json
import os
import pickle
from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from lib.pymafx.core import constants, path_config
from lib.smplx import SMPL as _SMPL
from lib.smplx import FLAMELayer, MANOLayer, SMPLXLayer
from lib.smplx.body_models import SMPLXOutput
from lib.smplx.lbs import (
    batch_rodrigues,
    blend_shapes,
    transform_mat,
    vertices2joints,
)

SMPL_MEAN_PARAMS = path_config.SMPL_MEAN_PARAMS
SMPL_MODEL_DIR = path_config.SMPL_MODEL_DIR


@dataclass
class ModelOutput(SMPLXOutput):
    smpl_joints: Optional[torch.Tensor] = None
    joints_J19: Optional[torch.Tensor] = None
    smplx_vertices: Optional[torch.Tensor] = None
    flame_vertices: Optional[torch.Tensor] = None
    lhand_vertices: Optional[torch.Tensor] = None
    rhand_vertices: Optional[torch.Tensor] = None
    lhand_joints: Optional[torch.Tensor] = None
    rhand_joints: Optional[torch.Tensor] = None
    face_joints: Optional[torch.Tensor] = None
    lfoot_joints: Optional[torch.Tensor] = None
    rfoot_joints: Optional[torch.Tensor] = None


class SMPL(_SMPL):
    """ Extension of the official SMPL implementation to support more joints """
    def __init__(
        self,
        create_betas=False,
        create_global_orient=False,
        create_body_pose=False,
        create_transl=False,
        *args,
        **kwargs
    ):
        super().__init__(
            create_betas=create_betas,
            create_global_orient=create_global_orient,
            create_body_pose=create_body_pose,
            create_transl=create_transl,
            *args,
            **kwargs
        )
        joints = [constants.JOINT_MAP[i] for i in constants.JOINT_NAMES]
        J_regressor_extra = np.load(path_config.JOINT_REGRESSOR_TRAIN_EXTRA)
        self.register_buffer(
            'J_regressor_extra', torch.tensor(J_regressor_extra, dtype=torch.float32)
        )
        self.joint_map = torch.tensor(joints, dtype=torch.long)
        # self.ModelOutput = namedtuple('ModelOutput_', ModelOutput._fields + ('smpl_joints', 'joints_J19',))
        # self.ModelOutput.__new__.__defaults__ = (None,) * len(self.ModelOutput._fields)

        tpose_joints = vertices2joints(self.J_regressor, self.v_template.unsqueeze(0))
        self.register_buffer('tpose_joints', tpose_joints)

    def forward(self, *args, **kwargs):
        kwargs['get_skin'] = True
        smpl_output = super().forward(*args, **kwargs)
        extra_joints = vertices2joints(self.J_regressor_extra, smpl_output.vertices)
        # smpl_output.joints: [B, 45, 3]  extra_joints: [B, 9, 3]
        vertices = smpl_output.vertices
        joints = torch.cat([smpl_output.joints, extra_joints], dim=1)
        smpl_joints = smpl_output.joints[:, :24]
        joints = joints[:, self.joint_map, :]    # [B, 49, 3]
        joints_J24 = joints[:, -24:, :]
        joints_J19 = joints_J24[:, constants.J24_TO_J19, :]
        output = ModelOutput(
            vertices=vertices,
            global_orient=smpl_output.global_orient,
            body_pose=smpl_output.body_pose,
            joints=joints,
            joints_J19=joints_J19,
            smpl_joints=smpl_joints,
            betas=smpl_output.betas,
            full_pose=smpl_output.full_pose
        )
        return output

    def get_global_rotation(
        self,
        global_orient: Optional[torch.Tensor] = None,
        body_pose: Optional[torch.Tensor] = None,
        **kwargs
    ):
        '''
        Forward pass for the SMPLX model

            Parameters
            ----------
            global_orient: torch.tensor, optional, shape Bx3x3
                If given, ignore the member variable and use it as the global
                rotation of the body. Useful if someone wishes to predicts this
                with an external model. It is expected to be in rotation matrix
                format. (default=None)
            body_pose: torch.tensor, optional, shape BxJx3x3
                If given, ignore the member variable `body_pose` and use it
                instead. For example, it can used if someone predicts the
                pose of the body joints are predicted from some external model.
                It should be a tensor that contains joint rotations in
                rotation matrix format. (default=None)
            Returns
            -------
                output: Global rotation matrix
        '''
        device, dtype = self.shapedirs.device, self.shapedirs.dtype

        model_vars = [global_orient, body_pose]
        batch_size = 1
        for var in model_vars:
            if var is None:
                continue
            batch_size = max(batch_size, len(var))

        if global_orient is None:
            global_orient = torch.eye(3, device=device,
                                      dtype=dtype).view(1, 1, 3, 3).expand(batch_size, -1, -1,
                                                                           -1).contiguous()
        if body_pose is None:
            body_pose = torch.eye(3, device=device, dtype=dtype).view(1, 1, 3, 3).expand(
                batch_size, self.NUM_BODY_JOINTS, -1, -1
            ).contiguous()

        # Concatenate all pose vectors
        full_pose = torch.cat([
            global_orient.reshape(-1, 1, 3, 3),
            body_pose.reshape(-1, self.NUM_BODY_JOINTS, 3, 3)
        ],
                              dim=1)

        rot_mats = full_pose.view(batch_size, -1, 3, 3)

        # Get the joints
        # NxJx3 array
        # joints = vertices2joints(self.J_regressor, self.v_template.unsqueeze(0).expand(batch_size, -1, -1))
        # joints = torch.unsqueeze(joints, dim=-1)

        joints = self.tpose_joints.expand(batch_size, -1, -1).unsqueeze(-1)

        rel_joints = joints.clone()
        rel_joints[:, 1:] -= joints[:, self.parents[1:]]

        transforms_mat = transform_mat(rot_mats.reshape(-1, 3, 3),
                                       rel_joints.reshape(-1, 3,
                                                          1)).reshape(-1, joints.shape[1], 4, 4)

        transform_chain = [transforms_mat[:, 0]]
        for i in range(1, self.parents.shape[0]):
            # Subtract the joint location at the rest pose
            # No need for rotation, since it's identity when at rest
            curr_res = torch.matmul(transform_chain[self.parents[i]], transforms_mat[:, i])
            transform_chain.append(curr_res)

        transforms = torch.stack(transform_chain, dim=1)

        global_rotmat = transforms[:, :, :3, :3]

        # The last column of the transformations contains the posed joints
        posed_joints = transforms[:, :, :3, 3]

        return global_rotmat, posed_joints


class SMPLX(SMPLXLayer):
    """ Extension of the official SMPLX implementation to support more functions """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def get_global_rotation(
        self,
        global_orient: Optional[torch.Tensor] = None,
        body_pose: Optional[torch.Tensor] = None,
        left_hand_pose: Optional[torch.Tensor] = None,
        right_hand_pose: Optional[torch.Tensor] = None,
        jaw_pose: Optional[torch.Tensor] = None,
        leye_pose: Optional[torch.Tensor] = None,
        reye_pose: Optional[torch.Tensor] = None,
        **kwargs
    ):
        '''
        Forward pass for the SMPLX model

            Parameters
            ----------
            global_orient: torch.tensor, optional, shape Bx3x3
                If given, ignore the member variable and use it as the global
                rotation of the body. Useful if someone wishes to predicts this
                with an external model. It is expected to be in rotation matrix
                format. (default=None)
            betas: torch.tensor, optional, shape BxN_b
                If given, ignore the member variable `betas` and use it
                instead. For example, it can used if shape parameters
                `betas` are predicted from some external model.
                (default=None)
            expression: torch.tensor, optional, shape BxN_e
                Expression coefficients.
                For example, it can used if expression parameters
                `expression` are predicted from some external model.
            body_pose: torch.tensor, optional, shape BxJx3x3
                If given, ignore the member variable `body_pose` and use it
                instead. For example, it can used if someone predicts the
                pose of the body joints are predicted from some external model.
                It should be a tensor that contains joint rotations in
                rotation matrix format. (default=None)
            left_hand_pose: torch.tensor, optional, shape Bx15x3x3
                If given, contains the pose of the left hand.
                It should be a tensor that contains joint rotations in
                rotation matrix format. (default=None)
            right_hand_pose: torch.tensor, optional, shape Bx15x3x3
                If given, contains the pose of the right hand.
                It should be a tensor that contains joint rotations in
                rotation matrix format. (default=None)
            jaw_pose: torch.tensor, optional, shape Bx3x3
                Jaw pose. It should either joint rotations in
                rotation matrix format.
            transl: torch.tensor, optional, shape Bx3
                Translation vector of the body.
                For example, it can used if the translation
                `transl` is predicted from some external model.
                (default=None)
            return_verts: bool, optional
                Return the vertices. (default=True)
            return_full_pose: bool, optional
                Returns the full pose vector (default=False)
            Returns
            -------
                output: ModelOutput
                A data class that contains the posed vertices and joints
        '''
        device, dtype = self.shapedirs.device, self.shapedirs.dtype

        model_vars = [global_orient, body_pose, left_hand_pose, right_hand_pose, jaw_pose]
        batch_size = 1
        for var in model_vars:
            if var is None:
                continue
            batch_size = max(batch_size, len(var))

        if global_orient is None:
            global_orient = torch.eye(3, device=device,
                                      dtype=dtype).view(1, 1, 3, 3).expand(batch_size, -1, -1,
                                                                           -1).contiguous()
        if body_pose is None:
            body_pose = torch.eye(3, device=device, dtype=dtype).view(1, 1, 3, 3).expand(
                batch_size, self.NUM_BODY_JOINTS, -1, -1
            ).contiguous()
        if left_hand_pose is None:
            left_hand_pose = torch.eye(3, device=device,
                                       dtype=dtype).view(1, 1, 3, 3).expand(batch_size, 15, -1,
                                                                            -1).contiguous()
        if right_hand_pose is None:
            right_hand_pose = torch.eye(3, device=device,
                                        dtype=dtype).view(1, 1, 3,
                                                          3).expand(batch_size, 15, -1,
                                                                    -1).contiguous()
        if jaw_pose is None:
            jaw_pose = torch.eye(3, device=device,
                                 dtype=dtype).view(1, 1, 3, 3).expand(batch_size, -1, -1,
                                                                      -1).contiguous()
        if leye_pose is None:
            leye_pose = torch.eye(3, device=device,
                                  dtype=dtype).view(1, 1, 3, 3).expand(batch_size, -1, -1,
                                                                       -1).contiguous()
        if reye_pose is None:
            reye_pose = torch.eye(3, device=device,
                                  dtype=dtype).view(1, 1, 3, 3).expand(batch_size, -1, -1,
                                                                       -1).contiguous()

        # Concatenate all pose vectors
        full_pose = torch.cat([
            global_orient.reshape(-1, 1, 3, 3),
            body_pose.reshape(-1, self.NUM_BODY_JOINTS, 3, 3),
            jaw_pose.reshape(-1, 1, 3, 3),
            leye_pose.reshape(-1, 1, 3, 3),
            reye_pose.reshape(-1, 1, 3, 3),
            left_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3, 3),
            right_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3, 3)
        ],
                              dim=1)

        rot_mats = full_pose.view(batch_size, -1, 3, 3)

        # Get the joints
        # NxJx3 array
        joints = vertices2joints(
            self.J_regressor,
            self.v_template.unsqueeze(0).expand(batch_size, -1, -1)
        )

        joints = torch.unsqueeze(joints, dim=-1)

        rel_joints = joints.clone()
        rel_joints[:, 1:] -= joints[:, self.parents[1:]]

        transforms_mat = transform_mat(rot_mats.reshape(-1, 3, 3),
                                       rel_joints.reshape(-1, 3,
                                                          1)).reshape(-1, joints.shape[1], 4, 4)

        transform_chain = [transforms_mat[:, 0]]
        for i in range(1, self.parents.shape[0]):
            # Subtract the joint location at the rest pose
            # No need for rotation, since it's identity when at rest
            curr_res = torch.matmul(transform_chain[self.parents[i]], transforms_mat[:, i])
            transform_chain.append(curr_res)

        transforms = torch.stack(transform_chain, dim=1)

        global_rotmat = transforms[:, :, :3, :3]

        # The last column of the transformations contains the posed joints
        posed_joints = transforms[:, :, :3, 3]

        return global_rotmat, posed_joints


class SMPLX_ALL(nn.Module):
    """ Extension of the official SMPLX implementation to support more joints """
    def __init__(self, batch_size=1, use_face_contour=True, all_gender=False, **kwargs):
        super().__init__()
        numBetas = 10
        self.use_face_contour = use_face_contour
        if all_gender:
            self.genders = ['male', 'female', 'neutral']
        else:
            self.genders = ['neutral']
        for gender in self.genders:
            assert gender in ['male', 'female', 'neutral']
        self.model_dict = nn.ModuleDict({
            gender: SMPLX(
                path_config.SMPL_MODEL_DIR,
                gender=gender,
                ext='npz',
                num_betas=numBetas,
                use_pca=False,
                batch_size=batch_size,
                use_face_contour=use_face_contour,
                num_pca_comps=45,
                **kwargs
            )
            for gender in self.genders
        })
        self.model_neutral = self.model_dict['neutral']
        joints = [constants.JOINT_MAP[i] for i in constants.JOINT_NAMES]
        J_regressor_extra = np.load(path_config.JOINT_REGRESSOR_TRAIN_EXTRA)
        self.register_buffer(
            'J_regressor_extra', torch.tensor(J_regressor_extra, dtype=torch.float32)
        )
        self.joint_map = torch.tensor(joints, dtype=torch.long)
        # smplx_to_smpl.pkl, file source: https://smpl-x.is.tue.mpg.de
        smplx_to_smpl = pickle.load(
            open(os.path.join(SMPL_MODEL_DIR, 'model_transfer/smplx_to_smpl.pkl'), 'rb')
        )
        self.register_buffer(
            'smplx2smpl', torch.tensor(smplx_to_smpl['matrix'][None], dtype=torch.float32)
        )

        smpl2limb_vert_faces = get_partial_smpl('smpl')
        self.smpl2lhand = torch.from_numpy(smpl2limb_vert_faces['lhand']['vids']).long()
        self.smpl2rhand = torch.from_numpy(smpl2limb_vert_faces['rhand']['vids']).long()

        # left and right hand joint mapping
        smplx2lhand_joints = [
            constants.SMPLX_JOINT_IDS['left_{}'.format(name)] for name in constants.HAND_NAMES
        ]
        smplx2rhand_joints = [
            constants.SMPLX_JOINT_IDS['right_{}'.format(name)] for name in constants.HAND_NAMES
        ]
        self.smplx2lh_joint_map = torch.tensor(smplx2lhand_joints, dtype=torch.long)
        self.smplx2rh_joint_map = torch.tensor(smplx2rhand_joints, dtype=torch.long)

        # left and right foot joint mapping
        smplx2lfoot_joints = [
            constants.SMPLX_JOINT_IDS['left_{}'.format(name)] for name in constants.FOOT_NAMES
        ]
        smplx2rfoot_joints = [
            constants.SMPLX_JOINT_IDS['right_{}'.format(name)] for name in constants.FOOT_NAMES
        ]
        self.smplx2lf_joint_map = torch.tensor(smplx2lfoot_joints, dtype=torch.long)
        self.smplx2rf_joint_map = torch.tensor(smplx2rfoot_joints, dtype=torch.long)

        for g in self.genders:
            J_template = torch.einsum(
                'ji,ik->jk', [self.model_dict[g].J_regressor[:24], self.model_dict[g].v_template]
            )
            J_dirs = torch.einsum(
                'ji,ikl->jkl', [self.model_dict[g].J_regressor[:24], self.model_dict[g].shapedirs]
            )

            self.register_buffer(f'{g}_J_template', J_template)
            self.register_buffer(f'{g}_J_dirs', J_dirs)

    def forward(self, *args, **kwargs):
        batch_size = kwargs['body_pose'].shape[0]
        kwargs['get_skin'] = True
        if 'pose2rot' not in kwargs:
            kwargs['pose2rot'] = True
        if 'gender' not in kwargs:
            kwargs['gender'] = 2 * torch.ones(batch_size).to(kwargs['body_pose'].device)

        # pose for 55 joints: 1, 21, 15, 15, 1, 1, 1
        pose_keys = [
            'global_orient', 'body_pose', 'left_hand_pose', 'right_hand_pose', 'jaw_pose',
            'leye_pose', 'reye_pose'
        ]
        param_keys = ['betas'] + pose_keys
        if kwargs['pose2rot']:
            for key in pose_keys:
                if key in kwargs:
                    # if key == 'left_hand_pose':
                    #     kwargs[key] += self.model_neutral.left_hand_mean
                    # elif key == 'right_hand_pose':
                    #     kwargs[key] += self.model_neutral.right_hand_mean
                    kwargs[key] = batch_rodrigues(kwargs[key].contiguous().view(-1, 3)).view([
                        batch_size, -1, 3, 3
                    ])
        if kwargs['body_pose'].shape[1] == 23:
            # remove hand pose in the body_pose
            kwargs['body_pose'] = kwargs['body_pose'][:, :21]
        gender_idx_list = []
        smplx_vertices, smplx_joints = [], []
        for gi, g in enumerate(['male', 'female', 'neutral']):
            gender_idx = ((kwargs['gender'] == gi).nonzero(as_tuple=True)[0])
            if len(gender_idx) == 0:
                continue
            gender_idx_list.extend([int(idx) for idx in gender_idx])
            gender_kwargs = {'get_skin': kwargs['get_skin'], 'pose2rot': kwargs['pose2rot']}
            gender_kwargs.update({k: kwargs[k][gender_idx] for k in param_keys if k in kwargs})
            gender_smplx_output = self.model_dict[g].forward(*args, **gender_kwargs)
            smplx_vertices.append(gender_smplx_output.vertices)
            smplx_joints.append(gender_smplx_output.joints)

        idx_rearrange = [gender_idx_list.index(i) for i in range(len(list(gender_idx_list)))]
        idx_rearrange = torch.tensor(idx_rearrange).long().to(kwargs['body_pose'].device)

        smplx_vertices = torch.cat(smplx_vertices)[idx_rearrange]
        smplx_joints = torch.cat(smplx_joints)[idx_rearrange]

        # constants.HAND_NAMES
        lhand_joints = smplx_joints[:, self.smplx2lh_joint_map]
        rhand_joints = smplx_joints[:, self.smplx2rh_joint_map]
        # constants.FACIAL_LANDMARKS
        face_joints = smplx_joints[:, -68:] if self.use_face_contour else smplx_joints[:, -51:]
        # constants.FOOT_NAMES
        lfoot_joints = smplx_joints[:, self.smplx2lf_joint_map]
        rfoot_joints = smplx_joints[:, self.smplx2rf_joint_map]

        smpl_vertices = torch.bmm(self.smplx2smpl.expand(batch_size, -1, -1), smplx_vertices)
        lhand_vertices = smpl_vertices[:, self.smpl2lhand]
        rhand_vertices = smpl_vertices[:, self.smpl2rhand]
        extra_joints = vertices2joints(self.J_regressor_extra, smpl_vertices)
        # smpl_output.joints: [B, 45, 3]  extra_joints: [B, 9, 3]
        smplx_j45 = smplx_joints[:, constants.SMPLX2SMPL_J45]
        joints = torch.cat([smplx_j45, extra_joints], dim=1)
        smpl_joints = smplx_j45[:, :24]
        joints = joints[:, self.joint_map, :]    # [B, 49, 3]
        joints_J24 = joints[:, -24:, :]
        joints_J19 = joints_J24[:, constants.J24_TO_J19, :]
        output = ModelOutput(
            vertices=smpl_vertices,
            smplx_vertices=smplx_vertices,
            lhand_vertices=lhand_vertices,
            rhand_vertices=rhand_vertices,
        # global_orient=smplx_output.global_orient,
        # body_pose=smplx_output.body_pose,
            joints=joints,
            joints_J19=joints_J19,
            smpl_joints=smpl_joints,
        # betas=smplx_output.betas,
        # full_pose=smplx_output.full_pose,
            lhand_joints=lhand_joints,
            rhand_joints=rhand_joints,
            lfoot_joints=lfoot_joints,
            rfoot_joints=rfoot_joints,
            face_joints=face_joints,
        )
        return output

    # def make_hand_regressor(self):
    #     # borrowed from https://github.com/mks0601/Hand4Whole_RELEASE/blob/main/common/utils/human_models.py
    #     regressor = self.model_neutral.J_regressor.numpy()
    #     vertex_num = self.model_neutral.J_regressor.shape[-1]
    #     lhand_regressor = np.concatenate((regressor[[20,37,38,39],:],
    #                                         np.eye(vertex_num)[5361,None],
    #                                             regressor[[25,26,27],:],
    #                                             np.eye(vertex_num)[4933,None],
    #                                             regressor[[28,29,30],:],
    #                                             np.eye(vertex_num)[5058,None],
    #                                             regressor[[34,35,36],:],
    #                                             np.eye(vertex_num)[5169,None],
    #                                             regressor[[31,32,33],:],
    #                                             np.eye(vertex_num)[5286,None]))
    #     rhand_regressor = np.concatenate((regressor[[21,52,53,54],:],
    #                                         np.eye(vertex_num)[8079,None],
    #                                             regressor[[40,41,42],:],
    #                                             np.eye(vertex_num)[7669,None],
    #                                             regressor[[43,44,45],:],
    #                                             np.eye(vertex_num)[7794,None],
    #                                             regressor[[49,50,51],:],
    #                                             np.eye(vertex_num)[7905,None],
    #                                             regressor[[46,47,48],:],
    #                                             np.eye(vertex_num)[8022,None]))
    #     return torch.from_numpy(lhand_regressor).float(), torch.from_numpy(rhand_regressor).float()

    def get_tpose(self, betas=None, gender=None):
        kwargs = {}
        if betas is None:
            betas = torch.zeros(1, 10).to(self.J_regressor_extra.device)
        kwargs['betas'] = betas

        batch_size = kwargs['betas'].shape[0]
        device = kwargs['betas'].device

        if gender is None:
            kwargs['gender'] = 2 * torch.ones(batch_size).to(device)
        else:
            kwargs['gender'] = gender

        param_keys = ['betas']

        gender_idx_list = []
        smplx_joints = []
        for gi, g in enumerate(['male', 'female', 'neutral']):
            gender_idx = ((kwargs['gender'] == gi).nonzero(as_tuple=True)[0])
            if len(gender_idx) == 0:
                continue
            gender_idx_list.extend([int(idx) for idx in gender_idx])
            gender_kwargs = {}
            gender_kwargs.update({k: kwargs[k][gender_idx] for k in param_keys if k in kwargs})

            J = getattr(self, f'{g}_J_template').unsqueeze(0) + blend_shapes(
                gender_kwargs['betas'], getattr(self, f'{g}_J_dirs')
            )

            smplx_joints.append(J)

        idx_rearrange = [gender_idx_list.index(i) for i in range(len(list(gender_idx_list)))]
        idx_rearrange = torch.tensor(idx_rearrange).long().to(device)

        smplx_joints = torch.cat(smplx_joints)[idx_rearrange]

        return smplx_joints


class MANO(MANOLayer):
    """ Extension of the official MANO implementation to support more joints """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def forward(self, *args, **kwargs):
        if 'pose2rot' not in kwargs:
            kwargs['pose2rot'] = True
        pose_keys = ['global_orient', 'right_hand_pose']
        batch_size = kwargs['global_orient'].shape[0]
        if kwargs['pose2rot']:
            for key in pose_keys:
                if key in kwargs:
                    kwargs[key] = batch_rodrigues(kwargs[key].contiguous().view(-1, 3)).view([
                        batch_size, -1, 3, 3
                    ])
        kwargs['hand_pose'] = kwargs.pop('right_hand_pose')
        mano_output = super().forward(*args, **kwargs)
        th_verts = mano_output.vertices
        th_jtr = mano_output.joints
        # https://github.com/hassony2/manopth/blob/master/manopth/manolayer.py#L248-L260
        # In addition to MANO reference joints we sample vertices on each finger
        # to serve as finger tips
        tips = th_verts[:, [745, 317, 445, 556, 673]]
        th_jtr = torch.cat([th_jtr, tips], 1)
        # Reorder joints to match visualization utilities
        th_jtr = th_jtr[:,
                        [0, 13, 14, 15, 16, 1, 2, 3, 17, 4, 5, 6, 18, 10, 11, 12, 19, 7, 8, 9, 20]]
        output = ModelOutput(
            rhand_vertices=th_verts,
            rhand_joints=th_jtr,
        )
        return output


class FLAME(FLAMELayer):
    """ Extension of the official FLAME implementation to support more joints """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def forward(self, *args, **kwargs):
        if 'pose2rot' not in kwargs:
            kwargs['pose2rot'] = True
        pose_keys = ['global_orient', 'jaw_pose', 'leye_pose', 'reye_pose']
        batch_size = kwargs['global_orient'].shape[0]
        if kwargs['pose2rot']:
            for key in pose_keys:
                if key in kwargs:
                    kwargs[key] = batch_rodrigues(kwargs[key].contiguous().view(-1, 3)).view([
                        batch_size, -1, 3, 3
                    ])
        flame_output = super().forward(*args, **kwargs)
        output = ModelOutput(
            flame_vertices=flame_output.vertices,
            face_joints=flame_output.joints[:, 5:],
        )
        return output


class SMPL_Family():
    def __init__(self, model_type='smpl', *args, **kwargs):
        if model_type == 'smpl':
            self.model = SMPL(model_path=SMPL_MODEL_DIR, *args, **kwargs)
        elif model_type == 'smplx':
            self.model = SMPLX_ALL(*args, **kwargs)
        elif model_type == 'mano':
            self.model = MANO(
                model_path=SMPL_MODEL_DIR, is_rhand=True, use_pca=False, *args, **kwargs
            )
        elif model_type == 'flame':
            self.model = FLAME(model_path=SMPL_MODEL_DIR, use_face_contour=True, *args, **kwargs)

    def __call__(self, *args, **kwargs):
        return self.model(*args, **kwargs)

    def get_tpose(self, *args, **kwargs):
        return self.model.get_tpose(*args, **kwargs)

    # def to(self, device):
    #     self.model.to(device)

    # def cuda(self, device=None):
    #     if device is None:
    #         self.model.cuda()
    #     else:
    #         self.model.cuda(device)


def get_smpl_faces():
    smpl = SMPL(model_path=SMPL_MODEL_DIR, batch_size=1)
    return smpl.faces


def get_smplx_faces():
    smplx = SMPLX(SMPL_MODEL_DIR, batch_size=1)
    return smplx.faces


def get_mano_faces(hand_type='right'):
    assert hand_type in ['right', 'left']
    is_rhand = True if hand_type == 'right' else False
    mano = MANO(SMPL_MODEL_DIR, batch_size=1, is_rhand=is_rhand)

    return mano.faces


def get_flame_faces():
    flame = FLAME(SMPL_MODEL_DIR, batch_size=1)

    return flame.faces


def get_model_faces(type='smpl'):
    if type == 'smpl':
        return get_smpl_faces()
    elif type == 'smplx':
        return get_smplx_faces()
    elif type == 'mano':
        return get_mano_faces()
    elif type == 'flame':
        return get_flame_faces()


def get_model_tpose(type='smpl'):
    if type == 'smpl':
        return get_smpl_tpose()
    elif type == 'smplx':
        return get_smplx_tpose()
    elif type == 'mano':
        return get_mano_tpose()
    elif type == 'flame':
        return get_flame_tpose()


def get_smpl_tpose():
    smpl = SMPL(
        create_betas=True,
        create_global_orient=True,
        create_body_pose=True,
        model_path=SMPL_MODEL_DIR,
        batch_size=1
    )
    vertices = smpl().vertices[0]
    return vertices.detach()


def get_smpl_tpose_joint():
    smpl = SMPL(
        create_betas=True,
        create_global_orient=True,
        create_body_pose=True,
        model_path=SMPL_MODEL_DIR,
        batch_size=1
    )
    tpose_joint = smpl().smpl_joints[0]
    return tpose_joint.detach()


def get_smplx_tpose():
    smplx = SMPLXLayer(SMPL_MODEL_DIR, batch_size=1)
    vertices = smplx().vertices[0]
    return vertices


def get_smplx_tpose_joint():
    smplx = SMPLXLayer(SMPL_MODEL_DIR, batch_size=1)
    tpose_joint = smplx().joints[0]
    return tpose_joint


def get_mano_tpose():
    mano = MANO(SMPL_MODEL_DIR, batch_size=1, is_rhand=True)
    vertices = mano(global_orient=torch.zeros(1, 3),
                    right_hand_pose=torch.zeros(1, 15 * 3)).rhand_vertices[0]
    return vertices


def get_flame_tpose():
    flame = FLAME(SMPL_MODEL_DIR, batch_size=1)
    vertices = flame(global_orient=torch.zeros(1, 3)).flame_vertices[0]
    return vertices


def get_part_joints(smpl_joints):
    batch_size = smpl_joints.shape[0]

    # part_joints = torch.zeros().to(smpl_joints.device)

    one_seg_pairs = [(0, 1), (0, 2), (0, 3), (3, 6), (9, 12), (9, 13), (9, 14), (12, 15), (13, 16),
                     (14, 17)]
    two_seg_pairs = [(1, 4), (2, 5), (4, 7), (5, 8), (16, 18), (17, 19), (18, 20), (19, 21)]

    one_seg_pairs.extend(two_seg_pairs)

    single_joints = [(10), (11), (15), (22), (23)]

    part_joints = []

    for j_p in one_seg_pairs:
        new_joint = torch.mean(smpl_joints[:, j_p], dim=1, keepdim=True)
        part_joints.append(new_joint)

    for j_p in single_joints:
        part_joints.append(smpl_joints[:, j_p:j_p + 1])

    part_joints = torch.cat(part_joints, dim=1)

    return part_joints


def get_partial_smpl(body_model='smpl', device=torch.device('cuda')):

    body_model_faces = get_model_faces(body_model)
    body_model_num_verts = len(get_model_tpose(body_model))

    part_vert_faces = {}

    for part in ['lhand', 'rhand', 'face', 'arm', 'forearm', 'larm', 'rarm', 'lwrist', 'rwrist']:
        part_vid_fname = '{}/{}_{}_vids.npz'.format(path_config.PARTIAL_MESH_DIR, body_model, part)
        if os.path.exists(part_vid_fname):
            part_vids = np.load(part_vid_fname)
            part_vert_faces[part] = {'vids': part_vids['vids'], 'faces': part_vids['faces']}
        else:
            if part in ['lhand', 'rhand']:
                with open(
                    os.path.join(SMPL_MODEL_DIR, 'model_transfer/MANO_SMPLX_vertex_ids.pkl'), 'rb'
                ) as json_file:
                    smplx_mano_id = pickle.load(json_file)
                with open(
                    os.path.join(SMPL_MODEL_DIR, 'model_transfer/smplx_to_smpl.pkl'), 'rb'
                ) as json_file:
                    smplx_smpl_id = pickle.load(json_file)

                smplx_tpose = get_smplx_tpose()
                smpl_tpose = np.matmul(smplx_smpl_id['matrix'], smplx_tpose)

                if part == 'lhand':
                    mano_vert = smplx_tpose[smplx_mano_id['left_hand']]
                elif part == 'rhand':
                    mano_vert = smplx_tpose[smplx_mano_id['right_hand']]

                smpl2mano_id = []
                for vert in mano_vert:
                    v_diff = smpl_tpose - vert
                    v_diff = torch.sum(v_diff * v_diff, dim=1)
                    v_closest = torch.argmin(v_diff)
                    smpl2mano_id.append(int(v_closest))

                smpl2mano_vids = np.array(smpl2mano_id).astype(np.long)
                mano_faces = get_mano_faces(hand_type='right' if part == 'rhand' else 'left'
                                           ).astype(np.long)

                np.savez(part_vid_fname, vids=smpl2mano_vids, faces=mano_faces)
                part_vert_faces[part] = {'vids': smpl2mano_vids, 'faces': mano_faces}

            elif part in ['face', 'arm', 'forearm', 'larm', 'rarm']:
                with open(
                    os.path.join(SMPL_MODEL_DIR, '{}_vert_segmentation.json'.format(body_model)),
                    'rb'
                ) as json_file:
                    smplx_part_id = json.load(json_file)

                # main_body_part = list(smplx_part_id.keys())
                # print('main_body_part', main_body_part)

                if part == 'face':
                    selected_body_part = ['head']
                elif part == 'arm':
                    selected_body_part = [
                        'rightHand',
                        'leftArm',
                        'leftShoulder',
                        'rightShoulder',
                        'rightArm',
                        'leftHandIndex1',
                        'rightHandIndex1',
                        'leftForeArm',
                        'rightForeArm',
                        'leftHand',
                    ]
                    # selected_body_part = ['rightHand', 'leftArm', 'rightArm', 'leftHandIndex1', 'rightHandIndex1', 'leftForeArm', 'rightForeArm', 'leftHand',]
                elif part == 'forearm':
                    selected_body_part = [
                        'rightHand',
                        'leftHandIndex1',
                        'rightHandIndex1',
                        'leftForeArm',
                        'rightForeArm',
                        'leftHand',
                    ]
                elif part == 'arm_eval':
                    selected_body_part = ['leftArm', 'rightArm', 'leftForeArm', 'rightForeArm']
                elif part == 'larm':
                    # selected_body_part = ['leftArm', 'leftForeArm']
                    selected_body_part = ['leftForeArm']
                elif part == 'rarm':
                    # selected_body_part = ['rightArm', 'rightForeArm']
                    selected_body_part = ['rightForeArm']

                part_body_idx = []
                for k in selected_body_part:
                    part_body_idx.extend(smplx_part_id[k])

                part_body_fid = []
                for f_id, face in enumerate(body_model_faces):
                    if any(f in part_body_idx for f in face):
                        part_body_fid.append(f_id)

                smpl2head_vids = np.unique(body_model_faces[part_body_fid]).astype(np.long)

                mesh_vid_raw = np.arange(body_model_num_verts)
                head_vid_new = np.arange(len(smpl2head_vids))
                mesh_vid_raw[smpl2head_vids] = head_vid_new

                head_faces = body_model_faces[part_body_fid]
                head_faces = mesh_vid_raw[head_faces].astype(np.long)

                np.savez(part_vid_fname, vids=smpl2head_vids, faces=head_faces)
                part_vert_faces[part] = {'vids': smpl2head_vids, 'faces': head_faces}

            elif part in ['lwrist', 'rwrist']:

                if body_model == 'smplx':
                    body_model_verts = get_smplx_tpose()
                    tpose_joint = get_smplx_tpose_joint()
                elif body_model == 'smpl':
                    body_model_verts = get_smpl_tpose()
                    tpose_joint = get_smpl_tpose_joint()

                wrist_joint = tpose_joint[20] if part == 'lwrist' else tpose_joint[21]

                dist = 0.005
                wrist_vids = []
                for vid, vt in enumerate(body_model_verts):

                    v_j_dist = torch.sum((vt - wrist_joint)**2)

                    if v_j_dist < dist:
                        wrist_vids.append(vid)

                wrist_vids = np.array(wrist_vids)

                part_body_fid = []
                for f_id, face in enumerate(body_model_faces):
                    if any(f in wrist_vids for f in face):
                        part_body_fid.append(f_id)

                smpl2part_vids = np.unique(body_model_faces[part_body_fid]).astype(np.long)

                mesh_vid_raw = np.arange(body_model_num_verts)
                part_vid_new = np.arange(len(smpl2part_vids))
                mesh_vid_raw[smpl2part_vids] = part_vid_new

                part_faces = body_model_faces[part_body_fid]
                part_faces = mesh_vid_raw[part_faces].astype(np.long)

                np.savez(part_vid_fname, vids=smpl2part_vids, faces=part_faces)
                part_vert_faces[part] = {'vids': smpl2part_vids, 'faces': part_faces}

                # import trimesh
                # mesh = trimesh.Trimesh(vertices=body_model_verts[smpl2part_vids], faces=part_faces, process=False)
                # mesh.export(f'results/smplx_{part}.obj')

                # mesh = trimesh.Trimesh(vertices=body_model_verts, faces=body_model_faces, process=False)
                # mesh.export(f'results/smplx_model.obj')

    return part_vert_faces