lib/pymafx/models/pymaf_net.py

import torch
import torch.nn as nn
import numpy as np
from lib.pymafx.core import constants

from lib.common.config import cfg
from lib.pymafx.utils.geometry import rot6d_to_rotmat, rotmat_to_rot6d, projection, rotation_matrix_to_angle_axis, compute_twist_rotation
from .maf_extractor import MAF_Extractor, Mesh_Sampler
from .smpl import SMPL, SMPL_MODEL_DIR, SMPL_MEAN_PARAMS, get_partial_smpl, SMPL_Family
from lib.smplx.lbs import batch_rodrigues
from .res_module import IUV_predict_layer
from .hr_module import get_hrnet_encoder
from .pose_resnet import get_resnet_encoder
from lib.pymafx.utils.imutils import j2d_processing
from lib.pymafx.utils.cam_params import homo_vector
from .attention import get_att_block

import logging

logger = logging.getLogger(__name__)

BN_MOMENTUM = 0.1


class Regressor(nn.Module):
    def __init__(
        self,
        feat_dim,
        smpl_mean_params,
        use_cam_feats=False,
        feat_dim_hand=0,
        feat_dim_face=0,
        bhf_names=['body'],
        smpl_models={}
    ):
        super().__init__()

        npose = 24 * 6
        shape_dim = 10
        cam_dim = 3
        hand_dim = 15 * 6
        face_dim = 3 * 6 + 10

        self.body_feat_dim = feat_dim

        self.smpl_mode = (cfg.MODEL.MESH_MODEL == 'smpl')
        self.smplx_mode = (cfg.MODEL.MESH_MODEL == 'smplx')
        self.use_cam_feats = use_cam_feats

        cam_feat_len = 4 if self.use_cam_feats else 0

        self.bhf_names = bhf_names
        self.hand_only_mode = (cfg.TRAIN.BHF_MODE == 'hand_only')
        self.face_only_mode = (cfg.TRAIN.BHF_MODE == 'face_only')
        self.body_hand_mode = (cfg.TRAIN.BHF_MODE == 'body_hand')
        self.full_body_mode = (cfg.TRAIN.BHF_MODE == 'full_body')

        # if self.use_cam_feats:
        #     assert cfg.MODEL.USE_IWP_CAM is False
        if 'body' in self.bhf_names:
            self.fc1 = nn.Linear(feat_dim + npose + cam_feat_len + shape_dim + cam_dim, 1024)
            self.drop1 = nn.Dropout()
            self.fc2 = nn.Linear(1024, 1024)
            self.drop2 = nn.Dropout()
            self.decpose = nn.Linear(1024, npose)
            self.decshape = nn.Linear(1024, 10)
            self.deccam = nn.Linear(1024, 3)
            nn.init.xavier_uniform_(self.decpose.weight, gain=0.01)
            nn.init.xavier_uniform_(self.decshape.weight, gain=0.01)
            nn.init.xavier_uniform_(self.deccam.weight, gain=0.01)

        if not self.smpl_mode:
            if self.hand_only_mode:
                self.part_names = ['rhand']
            elif self.face_only_mode:
                self.part_names = ['face']
            elif self.body_hand_mode:
                self.part_names = ['lhand', 'rhand']
            elif self.full_body_mode:
                self.part_names = ['lhand', 'rhand', 'face']
            else:
                self.part_names = []

            if 'rhand' in self.part_names:
                # self.fc1_hand = nn.Linear(feat_dim_hand + hand_dim + rh_orient_dim + rh_shape_dim + rh_cam_dim, 1024)
                self.fc1_hand = nn.Linear(feat_dim_hand + hand_dim, 1024)
                self.drop1_hand = nn.Dropout()
                self.fc2_hand = nn.Linear(1024, 1024)
                self.drop2_hand = nn.Dropout()

                # self.declhand = nn.Linear(1024, 15*6)
                self.decrhand = nn.Linear(1024, 15 * 6)
                # nn.init.xavier_uniform_(self.declhand.weight, gain=0.01)
                nn.init.xavier_uniform_(self.decrhand.weight, gain=0.01)

                if cfg.MODEL.MESH_MODEL == 'mano' or cfg.MODEL.PyMAF.OPT_WRIST:
                    rh_cam_dim = 3
                    rh_orient_dim = 6
                    rh_shape_dim = 10
                    self.fc3_hand = nn.Linear(
                        1024 + rh_orient_dim + rh_shape_dim + rh_cam_dim, 1024
                    )
                    self.drop3_hand = nn.Dropout()

                    self.decshape_rhand = nn.Linear(1024, 10)
                    self.decorient_rhand = nn.Linear(1024, 6)
                    self.deccam_rhand = nn.Linear(1024, 3)
                    nn.init.xavier_uniform_(self.decshape_rhand.weight, gain=0.01)
                    nn.init.xavier_uniform_(self.decorient_rhand.weight, gain=0.01)
                    nn.init.xavier_uniform_(self.deccam_rhand.weight, gain=0.01)

            if 'face' in self.part_names:
                self.fc1_face = nn.Linear(feat_dim_face + face_dim, 1024)
                self.drop1_face = nn.Dropout()
                self.fc2_face = nn.Linear(1024, 1024)
                self.drop2_face = nn.Dropout()

                self.dechead = nn.Linear(1024, 3 * 6)
                self.decexp = nn.Linear(1024, 10)
                nn.init.xavier_uniform_(self.dechead.weight, gain=0.01)
                nn.init.xavier_uniform_(self.decexp.weight, gain=0.01)

                if cfg.MODEL.MESH_MODEL == 'flame':
                    rh_cam_dim = 3
                    rh_orient_dim = 6
                    rh_shape_dim = 10
                    self.fc3_face = nn.Linear(
                        1024 + rh_orient_dim + rh_shape_dim + rh_cam_dim, 1024
                    )
                    self.drop3_face = nn.Dropout()

                    self.decshape_face = nn.Linear(1024, 10)
                    self.decorient_face = nn.Linear(1024, 6)
                    self.deccam_face = nn.Linear(1024, 3)
                    nn.init.xavier_uniform_(self.decshape_face.weight, gain=0.01)
                    nn.init.xavier_uniform_(self.decorient_face.weight, gain=0.01)
                    nn.init.xavier_uniform_(self.deccam_face.weight, gain=0.01)

            if self.smplx_mode and cfg.MODEL.PyMAF.PRED_VIS_H:
                self.fc1_vis = nn.Linear(1024 + 1024 + 1024, 1024)
                self.drop1_vis = nn.Dropout()
                self.fc2_vis = nn.Linear(1024, 1024)
                self.drop2_vis = nn.Dropout()
                self.decvis = nn.Linear(1024, 2)
                nn.init.xavier_uniform_(self.decvis.weight, gain=0.01)

        if 'body' in smpl_models:
            self.smpl = smpl_models['body']
        if 'hand' in smpl_models:
            self.mano = smpl_models['hand']
        if 'face' in smpl_models:
            self.flame = smpl_models['face']

        if cfg.MODEL.PyMAF.OPT_WRIST:
            self.body_model = SMPL(model_path=SMPL_MODEL_DIR, batch_size=64, create_transl=False)

        mean_params = np.load(smpl_mean_params)
        init_pose = torch.from_numpy(mean_params['pose'][:]).unsqueeze(0)
        init_shape = torch.from_numpy(mean_params['shape'][:].astype('float32')).unsqueeze(0)
        init_cam = torch.from_numpy(mean_params['cam']).unsqueeze(0)
        self.register_buffer('init_pose', init_pose)
        self.register_buffer('init_shape', init_shape)
        self.register_buffer('init_cam', init_cam)
        self.register_buffer('init_orient', init_pose[:, :6])

        self.flip_vector = torch.ones((1, 9), dtype=torch.float32)
        self.flip_vector[:, [1, 2, 3, 6]] *= -1
        self.flip_vector = self.flip_vector.reshape(1, 3, 3)

        if not self.smpl_mode:
            lhand_mean_rot6d = rotmat_to_rot6d(
                batch_rodrigues(self.smpl.model.model_neutral.left_hand_mean.view(-1, 3)).view(
                    [-1, 3, 3]
                )
            )
            rhand_mean_rot6d = rotmat_to_rot6d(
                batch_rodrigues(self.smpl.model.model_neutral.right_hand_mean.view(-1, 3)).view(
                    [-1, 3, 3]
                )
            )
            init_lhand = lhand_mean_rot6d.reshape(-1).unsqueeze(0)
            init_rhand = rhand_mean_rot6d.reshape(-1).unsqueeze(0)
            # init_hand = torch.cat([init_lhand, init_rhand]).unsqueeze(0)
            init_face = rotmat_to_rot6d(torch.stack([torch.eye(3)] * 3)).reshape(-1).unsqueeze(0)
            init_exp = torch.zeros(10).unsqueeze(0)

        if self.smplx_mode or 'hand' in bhf_names:
            # init_hand = torch.cat([init_lhand, init_rhand]).unsqueeze(0)
            self.register_buffer('init_lhand', init_lhand)
            self.register_buffer('init_rhand', init_rhand)
        if self.smplx_mode or 'face' in bhf_names:
            self.register_buffer('init_face', init_face)
            self.register_buffer('init_exp', init_exp)

    def forward(
        self,
        x=None,
        n_iter=1,
        J_regressor=None,
        rw_cam={},
        init_mode=False,
        global_iter=-1,
        **kwargs
    ):
        if x is not None:
            batch_size = x.shape[0]
        else:
            if 'xc_rhand' in kwargs:
                batch_size = kwargs['xc_rhand'].shape[0]
            elif 'xc_face' in kwargs:
                batch_size = kwargs['xc_face'].shape[0]

        if 'body' in self.bhf_names:
            if 'init_pose' not in kwargs:
                kwargs['init_pose'] = self.init_pose.expand(batch_size, -1)
            if 'init_shape' not in kwargs:
                kwargs['init_shape'] = self.init_shape.expand(batch_size, -1)
            if 'init_cam' not in kwargs:
                kwargs['init_cam'] = self.init_cam.expand(batch_size, -1)

            pred_cam = kwargs['init_cam']
            pred_pose = kwargs['init_pose']
            pred_shape = kwargs['init_shape']

        if self.full_body_mode or self.body_hand_mode:
            if cfg.MODEL.PyMAF.OPT_WRIST:
                pred_rotmat_body = rot6d_to_rotmat(
                    pred_pose.reshape(batch_size, -1, 6)
                )    # .view(batch_size, 24, 3, 3)
            if cfg.MODEL.PyMAF.PRED_VIS_H:
                pred_vis_hands = None

        # if self.full_body_mode or 'hand' in self.bhf_names:
        if self.smplx_mode or 'hand' in self.bhf_names:
            if 'init_lhand' not in kwargs:
                # kwargs['init_lhand'] = self.init_lhand.expand(batch_size, -1)
                # init with **right** hand pose
                kwargs['init_lhand'] = self.init_rhand.expand(batch_size, -1)
            if 'init_rhand' not in kwargs:
                kwargs['init_rhand'] = self.init_rhand.expand(batch_size, -1)

            pred_lhand, pred_rhand = kwargs['init_lhand'], kwargs['init_rhand']

            if cfg.MODEL.MESH_MODEL == 'mano' or cfg.MODEL.PyMAF.OPT_WRIST:
                if 'init_orient_rh' not in kwargs:
                    kwargs['init_orient_rh'] = self.init_orient.expand(batch_size, -1)
                if 'init_shape_rh' not in kwargs:
                    kwargs['init_shape_rh'] = self.init_shape.expand(batch_size, -1)
                if 'init_cam_rh' not in kwargs:
                    kwargs['init_cam_rh'] = self.init_cam.expand(batch_size, -1)
                pred_orient_rh = kwargs['init_orient_rh']
                pred_shape_rh = kwargs['init_shape_rh']
                pred_cam_rh = kwargs['init_cam_rh']
                if cfg.MODEL.PyMAF.OPT_WRIST:
                    if 'init_orient_lh' not in kwargs:
                        kwargs['init_orient_lh'] = self.init_orient.expand(batch_size, -1)
                    if 'init_shape_lh' not in kwargs:
                        kwargs['init_shape_lh'] = self.init_shape.expand(batch_size, -1)
                    if 'init_cam_lh' not in kwargs:
                        kwargs['init_cam_lh'] = self.init_cam.expand(batch_size, -1)
                    pred_orient_lh = kwargs['init_orient_lh']
                    pred_shape_lh = kwargs['init_shape_lh']
                    pred_cam_lh = kwargs['init_cam_lh']
                if cfg.MODEL.MESH_MODEL == 'mano':
                    pred_cam = torch.cat([pred_cam_rh[:, 0:1] * 10., pred_cam_rh[:, 1:]], dim=1)

        # if self.full_body_mode or 'face' in self.bhf_names:
        if self.smplx_mode or 'face' in self.bhf_names:
            if 'init_face' not in kwargs:
                kwargs['init_face'] = self.init_face.expand(batch_size, -1)
            if 'init_hand' not in kwargs:
                kwargs['init_exp'] = self.init_exp.expand(batch_size, -1)

            pred_face = kwargs['init_face']
            pred_exp = kwargs['init_exp']

            if cfg.MODEL.MESH_MODEL == 'flame' or cfg.MODEL.PyMAF.OPT_WRIST:
                if 'init_orient_fa' not in kwargs:
                    kwargs['init_orient_fa'] = self.init_orient.expand(batch_size, -1)
                pred_orient_fa = kwargs['init_orient_fa']
                if 'init_shape_fa' not in kwargs:
                    kwargs['init_shape_fa'] = self.init_shape.expand(batch_size, -1)
                if 'init_cam_fa' not in kwargs:
                    kwargs['init_cam_fa'] = self.init_cam.expand(batch_size, -1)
                pred_shape_fa = kwargs['init_shape_fa']
                pred_cam_fa = kwargs['init_cam_fa']
                if cfg.MODEL.MESH_MODEL == 'flame':
                    pred_cam = torch.cat([pred_cam_fa[:, 0:1] * 10., pred_cam_fa[:, 1:]], dim=1)

        if not init_mode:
            for i in range(n_iter):
                if 'body' in self.bhf_names:
                    xc = torch.cat([x, pred_pose, pred_shape, pred_cam], 1)
                    if self.use_cam_feats:
                        if cfg.MODEL.USE_IWP_CAM:
                            # for IWP camera, simply use pre-defined values
                            vfov = torch.ones((batch_size, 1)).to(xc) * 0.8
                            crop_ratio = torch.ones((batch_size, 1)).to(xc) * 0.3
                            crop_center = torch.ones((batch_size, 2)).to(xc) * 0.5
                        else:
                            vfov = rw_cam['vfov'][:, None]
                            crop_ratio = rw_cam['crop_ratio'][:, None]
                            crop_center = rw_cam['bbox_center'] / torch.cat(
                                [rw_cam['img_w'][:, None], rw_cam['img_h'][:, None]], 1
                            )
                        xc = torch.cat([xc, vfov, crop_ratio, crop_center], 1)

                    xc = self.fc1(xc)
                    xc = self.drop1(xc)
                    xc = self.fc2(xc)
                    xc = self.drop2(xc)

                    pred_cam = self.deccam(xc) + pred_cam
                    pred_pose = self.decpose(xc) + pred_pose
                    pred_shape = self.decshape(xc) + pred_shape

                if not self.smpl_mode:
                    if self.hand_only_mode:
                        xc_rhand = kwargs['xc_rhand']
                        xc_rhand = torch.cat([xc_rhand, pred_rhand], 1)
                    elif self.face_only_mode:
                        xc_face = kwargs['xc_face']
                        xc_face = torch.cat([xc_face, pred_face, pred_exp], 1)
                    elif self.body_hand_mode:
                        xc_lhand, xc_rhand = kwargs['xc_lhand'], kwargs['xc_rhand']
                        xc_lhand = torch.cat([xc_lhand, pred_lhand], 1)
                        xc_rhand = torch.cat([xc_rhand, pred_rhand], 1)
                    elif self.full_body_mode:
                        xc_lhand, xc_rhand, xc_face = kwargs['xc_lhand'], kwargs['xc_rhand'
                                                                                ], kwargs['xc_face']
                        xc_lhand = torch.cat([xc_lhand, pred_lhand], 1)
                        xc_rhand = torch.cat([xc_rhand, pred_rhand], 1)
                        xc_face = torch.cat([xc_face, pred_face, pred_exp], 1)

                    if 'lhand' in self.part_names:
                        xc_lhand = self.drop1_hand(self.fc1_hand(xc_lhand))
                        xc_lhand = self.drop2_hand(self.fc2_hand(xc_lhand))
                        pred_lhand = self.decrhand(xc_lhand) + pred_lhand

                        if cfg.MODEL.PyMAF.OPT_WRIST:
                            xc_lhand = torch.cat(
                                [xc_lhand, pred_shape_lh, pred_orient_lh, pred_cam_lh], 1
                            )
                            xc_lhand = self.drop3_hand(self.fc3_hand(xc_lhand))

                            pred_shape_lh = self.decshape_rhand(xc_lhand) + pred_shape_lh
                            pred_orient_lh = self.decorient_rhand(xc_lhand) + pred_orient_lh
                            pred_cam_lh = self.deccam_rhand(xc_lhand) + pred_cam_lh

                    if 'rhand' in self.part_names:
                        xc_rhand = self.drop1_hand(self.fc1_hand(xc_rhand))
                        xc_rhand = self.drop2_hand(self.fc2_hand(xc_rhand))
                        pred_rhand = self.decrhand(xc_rhand) + pred_rhand

                        if cfg.MODEL.MESH_MODEL == 'mano' or cfg.MODEL.PyMAF.OPT_WRIST:
                            xc_rhand = torch.cat(
                                [xc_rhand, pred_shape_rh, pred_orient_rh, pred_cam_rh], 1
                            )
                            xc_rhand = self.drop3_hand(self.fc3_hand(xc_rhand))

                            pred_shape_rh = self.decshape_rhand(xc_rhand) + pred_shape_rh
                            pred_orient_rh = self.decorient_rhand(xc_rhand) + pred_orient_rh
                            pred_cam_rh = self.deccam_rhand(xc_rhand) + pred_cam_rh

                            if cfg.MODEL.MESH_MODEL == 'mano':
                                pred_cam = torch.cat(
                                    [pred_cam_rh[:, 0:1] * 10., pred_cam_rh[:, 1:] / 10.], dim=1
                                )

                    if 'face' in self.part_names:
                        xc_face = self.drop1_face(self.fc1_face(xc_face))
                        xc_face = self.drop2_face(self.fc2_face(xc_face))
                        pred_face = self.dechead(xc_face) + pred_face
                        pred_exp = self.decexp(xc_face) + pred_exp

                        if cfg.MODEL.MESH_MODEL == 'flame':
                            xc_face = torch.cat(
                                [xc_face, pred_shape_fa, pred_orient_fa, pred_cam_fa], 1
                            )
                            xc_face = self.drop3_face(self.fc3_face(xc_face))

                            pred_shape_fa = self.decshape_face(xc_face) + pred_shape_fa
                            pred_orient_fa = self.decorient_face(xc_face) + pred_orient_fa
                            pred_cam_fa = self.deccam_face(xc_face) + pred_cam_fa

                            if cfg.MODEL.MESH_MODEL == 'flame':
                                pred_cam = torch.cat(
                                    [pred_cam_fa[:, 0:1] * 10., pred_cam_fa[:, 1:] / 10.], dim=1
                                )

                    if self.full_body_mode or self.body_hand_mode:
                        if cfg.MODEL.PyMAF.PRED_VIS_H:
                            xc_vis = torch.cat([xc, xc_lhand, xc_rhand], 1)

                            xc_vis = self.drop1_vis(self.fc1_vis(xc_vis))
                            xc_vis = self.drop2_vis(self.fc2_vis(xc_vis))
                            pred_vis_hands = self.decvis(xc_vis)

                            pred_vis_lhand = pred_vis_hands[:, 0] > cfg.MODEL.PyMAF.HAND_VIS_TH
                            pred_vis_rhand = pred_vis_hands[:, 1] > cfg.MODEL.PyMAF.HAND_VIS_TH

                        if cfg.MODEL.PyMAF.OPT_WRIST:

                            pred_rotmat_body = rot6d_to_rotmat(
                                pred_pose.reshape(batch_size, -1, 6)
                            )    # .view(batch_size, 24, 3, 3)
                            pred_lwrist = pred_rotmat_body[:, 20]
                            pred_rwrist = pred_rotmat_body[:, 21]

                            pred_gl_body, body_joints = self.body_model.get_global_rotation(
                                global_orient=pred_rotmat_body[:, 0:1],
                                body_pose=pred_rotmat_body[:, 1:]
                            )
                            pred_gl_lelbow = pred_gl_body[:, 18]
                            pred_gl_relbow = pred_gl_body[:, 19]

                            target_gl_lwrist = rot6d_to_rotmat(
                                pred_orient_lh.reshape(batch_size, -1, 6)
                            )
                            target_gl_lwrist *= self.flip_vector.to(target_gl_lwrist.device)
                            target_gl_rwrist = rot6d_to_rotmat(
                                pred_orient_rh.reshape(batch_size, -1, 6)
                            )

                            opt_lwrist = torch.bmm(pred_gl_lelbow.transpose(1, 2), target_gl_lwrist)
                            opt_rwrist = torch.bmm(pred_gl_relbow.transpose(1, 2), target_gl_rwrist)

                            if cfg.MODEL.PyMAF.ADAPT_INTEGR:
                                # if cfg.MODEL.PyMAF.ADAPT_INTEGR and global_iter == (cfg.MODEL.PyMAF.N_ITER - 1):
                                tpose_joints = self.smpl.get_tpose(betas=pred_shape)
                                lelbow_twist_axis = nn.functional.normalize(
                                    tpose_joints[:, 20] - tpose_joints[:, 18], dim=1
                                )
                                relbow_twist_axis = nn.functional.normalize(
                                    tpose_joints[:, 21] - tpose_joints[:, 19], dim=1
                                )

                                lelbow_twist, lelbow_twist_angle = compute_twist_rotation(
                                    opt_lwrist, lelbow_twist_axis
                                )
                                relbow_twist, relbow_twist_angle = compute_twist_rotation(
                                    opt_rwrist, relbow_twist_axis
                                )

                                min_angle = -0.4 * float(np.pi)
                                max_angle = 0.4 * float(np.pi)

                                lelbow_twist_angle[lelbow_twist_angle == torch.
                                                   clamp(lelbow_twist_angle, min_angle, max_angle)
                                                  ] = 0
                                relbow_twist_angle[relbow_twist_angle == torch.
                                                   clamp(relbow_twist_angle, min_angle, max_angle)
                                                  ] = 0
                                lelbow_twist_angle[lelbow_twist_angle > max_angle] -= max_angle
                                lelbow_twist_angle[lelbow_twist_angle < min_angle] -= min_angle
                                relbow_twist_angle[relbow_twist_angle > max_angle] -= max_angle
                                relbow_twist_angle[relbow_twist_angle < min_angle] -= min_angle

                                lelbow_twist = batch_rodrigues(
                                    lelbow_twist_axis * lelbow_twist_angle
                                )
                                relbow_twist = batch_rodrigues(
                                    relbow_twist_axis * relbow_twist_angle
                                )

                                opt_lwrist = torch.bmm(lelbow_twist.transpose(1, 2), opt_lwrist)
                                opt_rwrist = torch.bmm(relbow_twist.transpose(1, 2), opt_rwrist)

                                # left elbow: 18
                                opt_lelbow = torch.bmm(pred_rotmat_body[:, 18], lelbow_twist)
                                # right elbow: 19
                                opt_relbow = torch.bmm(pred_rotmat_body[:, 19], relbow_twist)

                                if cfg.MODEL.PyMAF.PRED_VIS_H and global_iter == (
                                    cfg.MODEL.PyMAF.N_ITER - 1
                                ):
                                    opt_lwrist_filtered = [
                                        opt_lwrist[_i]
                                        if pred_vis_lhand[_i] else pred_rotmat_body[_i, 20]
                                        for _i in range(batch_size)
                                    ]
                                    opt_rwrist_filtered = [
                                        opt_rwrist[_i]
                                        if pred_vis_rhand[_i] else pred_rotmat_body[_i, 21]
                                        for _i in range(batch_size)
                                    ]
                                    opt_lelbow_filtered = [
                                        opt_lelbow[_i]
                                        if pred_vis_lhand[_i] else pred_rotmat_body[_i, 18]
                                        for _i in range(batch_size)
                                    ]
                                    opt_relbow_filtered = [
                                        opt_relbow[_i]
                                        if pred_vis_rhand[_i] else pred_rotmat_body[_i, 19]
                                        for _i in range(batch_size)
                                    ]

                                    opt_lwrist = torch.stack(opt_lwrist_filtered)
                                    opt_rwrist = torch.stack(opt_rwrist_filtered)
                                    opt_lelbow = torch.stack(opt_lelbow_filtered)
                                    opt_relbow = torch.stack(opt_relbow_filtered)

                                pred_rotmat_body = torch.cat(
                                    [
                                        pred_rotmat_body[:, :18],
                                        opt_lelbow.unsqueeze(1),
                                        opt_relbow.unsqueeze(1),
                                        opt_lwrist.unsqueeze(1),
                                        opt_rwrist.unsqueeze(1), pred_rotmat_body[:, 22:]
                                    ], 1
                                )
                            else:
                                if cfg.MODEL.PyMAF.PRED_VIS_H and global_iter == (
                                    cfg.MODEL.PyMAF.N_ITER - 1
                                ):
                                    opt_lwrist_filtered = [
                                        opt_lwrist[_i]
                                        if pred_vis_lhand[_i] else pred_rotmat_body[_i, 20]
                                        for _i in range(batch_size)
                                    ]
                                    opt_rwrist_filtered = [
                                        opt_rwrist[_i]
                                        if pred_vis_rhand[_i] else pred_rotmat_body[_i, 21]
                                        for _i in range(batch_size)
                                    ]

                                    opt_lwrist = torch.stack(opt_lwrist_filtered)
                                    opt_rwrist = torch.stack(opt_rwrist_filtered)

                                pred_rotmat_body = torch.cat(
                                    [
                                        pred_rotmat_body[:, :20],
                                        opt_lwrist.unsqueeze(1),
                                        opt_rwrist.unsqueeze(1), pred_rotmat_body[:, 22:]
                                    ], 1
                                )

        if self.hand_only_mode:
            pred_rotmat_rh = rot6d_to_rotmat(
                torch.cat([pred_orient_rh, pred_rhand], dim=1).reshape(batch_size, -1, 6)
            )    # .view(batch_size, 16, 3, 3)
            assert pred_rotmat_rh.shape[1] == 1 + 15
        elif self.face_only_mode:
            pred_rotmat_fa = rot6d_to_rotmat(
                torch.cat([pred_orient_fa, pred_face], dim=1).reshape(batch_size, -1, 6)
            )    # .view(batch_size, 16, 3, 3)
            assert pred_rotmat_fa.shape[1] == 1 + 3
        elif self.full_body_mode or self.body_hand_mode:
            if cfg.MODEL.PyMAF.OPT_WRIST:
                pred_rotmat = pred_rotmat_body
            else:
                pred_rotmat = rot6d_to_rotmat(
                    pred_pose.reshape(batch_size, -1, 6)
                )    # .view(batch_size, 24, 3, 3)
            assert pred_rotmat.shape[1] == 24
        else:
            pred_rotmat = rot6d_to_rotmat(
                pred_pose.reshape(batch_size, -1, 6)
            )    # .view(batch_size, 24, 3, 3)
            assert pred_rotmat.shape[1] == 24

        # if self.full_body_mode:
        if self.smplx_mode:
            if cfg.MODEL.PyMAF.PRED_VIS_H and global_iter == (cfg.MODEL.PyMAF.N_ITER - 1):
                pred_lhand_filtered = [
                    pred_lhand[_i] if pred_vis_lhand[_i] else self.init_rhand[0]
                    for _i in range(batch_size)
                ]
                pred_rhand_filtered = [
                    pred_rhand[_i] if pred_vis_rhand[_i] else self.init_rhand[0]
                    for _i in range(batch_size)
                ]
                pred_lhand_filtered = torch.stack(pred_lhand_filtered)
                pred_rhand_filtered = torch.stack(pred_rhand_filtered)
                pred_hf6d = torch.cat([pred_lhand_filtered, pred_rhand_filtered, pred_face],
                                      dim=1).reshape(batch_size, -1, 6)
            else:
                pred_hf6d = torch.cat([pred_lhand, pred_rhand, pred_face],
                                      dim=1).reshape(batch_size, -1, 6)
            pred_hfrotmat = rot6d_to_rotmat(pred_hf6d)
            assert pred_hfrotmat.shape[1] == (15 * 2 + 3)

            # flip left hand pose
            pred_lhand_rotmat = pred_hfrotmat[:, :15] * self.flip_vector.to(pred_hfrotmat.device
                                                                           ).unsqueeze(0)
            pred_rhand_rotmat = pred_hfrotmat[:, 15:30]
            pred_face_rotmat = pred_hfrotmat[:, 30:]

        if self.hand_only_mode:
            pred_output = self.mano(
                betas=pred_shape_rh,
                right_hand_pose=pred_rotmat_rh[:, 1:],
                global_orient=pred_rotmat_rh[:, 0].unsqueeze(1),
                pose2rot=False,
            )
        elif self.face_only_mode:
            pred_output = self.flame(
                betas=pred_shape_fa,
                global_orient=pred_rotmat_fa[:, 0].unsqueeze(1),
                jaw_pose=pred_rotmat_fa[:, 1:2],
                leye_pose=pred_rotmat_fa[:, 2:3],
                reye_pose=pred_rotmat_fa[:, 3:4],
                expression=pred_exp,
                pose2rot=False,
            )
        else:
            smplx_kwargs = {}
            # if self.full_body_mode:
            if self.smplx_mode:
                smplx_kwargs['left_hand_pose'] = pred_lhand_rotmat
                smplx_kwargs['right_hand_pose'] = pred_rhand_rotmat
                smplx_kwargs['jaw_pose'] = pred_face_rotmat[:, 0:1]
                smplx_kwargs['leye_pose'] = pred_face_rotmat[:, 1:2]
                smplx_kwargs['reye_pose'] = pred_face_rotmat[:, 2:3]
                smplx_kwargs['expression'] = pred_exp

            pred_output = self.smpl(
                betas=pred_shape,
                body_pose=pred_rotmat[:, 1:],
                global_orient=pred_rotmat[:, 0].unsqueeze(1),
                pose2rot=False,
                **smplx_kwargs,
            )

            pred_vertices = pred_output.vertices
            pred_joints = pred_output.joints

        if self.hand_only_mode:
            pred_joints_full = pred_output.rhand_joints
        elif self.face_only_mode:
            pred_joints_full = pred_output.face_joints
        elif self.smplx_mode:
            pred_joints_full = torch.cat(
                [
                    pred_joints, pred_output.lhand_joints, pred_output.rhand_joints,
                    pred_output.face_joints, pred_output.lfoot_joints, pred_output.rfoot_joints
                ],
                dim=1
            )
        else:
            pred_joints_full = pred_joints
        pred_keypoints_2d = projection(
            pred_joints_full, {
                **rw_cam, 'cam_sxy': pred_cam
            }, iwp_mode=cfg.MODEL.USE_IWP_CAM
        )
        if cfg.MODEL.USE_IWP_CAM:
            # Normalize keypoints to [-1,1]
            pred_keypoints_2d = pred_keypoints_2d / (224. / 2.)
        else:
            pred_keypoints_2d = j2d_processing(pred_keypoints_2d, rw_cam['kps_transf'])

        len_b_kp = len(constants.JOINT_NAMES)
        output = {}
        if self.smpl_mode or self.smplx_mode:
            if J_regressor is not None:
                kp_3d = torch.matmul(J_regressor, pred_vertices)
                pred_pelvis = kp_3d[:, [0], :].clone()
                kp_3d = kp_3d[:, constants.H36M_TO_J14, :]
                kp_3d = kp_3d - pred_pelvis
            else:
                kp_3d = pred_joints
            pose = rotation_matrix_to_angle_axis(pred_rotmat.reshape(-1, 3, 3)).reshape(-1, 72)
            output.update(
                {
                    'theta': torch.cat([pred_cam, pred_shape, pose], dim=1),
                    'verts': pred_vertices,
                    'kp_2d': pred_keypoints_2d[:, :len_b_kp],
                    'kp_3d': kp_3d,
                    'pred_joints': pred_joints,
                    'smpl_kp_3d': pred_output.smpl_joints,
                    'rotmat': pred_rotmat,
                    'pred_cam': pred_cam,
                    'pred_shape': pred_shape,
                    'pred_pose': pred_pose,
                }
            )
            # if self.full_body_mode:
            if self.smplx_mode:
                # assert pred_keypoints_2d.shape[1] == 144
                len_h_kp = len(constants.HAND_NAMES)
                len_f_kp = len(constants.FACIAL_LANDMARKS)
                len_feet_kp = 2 * len(constants.FOOT_NAMES)
                output.update(
                    {
                        'smplx_verts':
                            pred_output.smplx_vertices if cfg.MODEL.EVAL_MODE else None,
                        'pred_lhand':
                            pred_lhand,
                        'pred_rhand':
                            pred_rhand,
                        'pred_face':
                            pred_face,
                        'pred_exp':
                            pred_exp,
                        'verts_lh':
                            pred_output.lhand_vertices,
                        'verts_rh':
                            pred_output.rhand_vertices,
                # 'pred_arm_rotmat': pred_arm_rotmat,
                # 'pred_hfrotmat': pred_hfrotmat,
                        'pred_lhand_rotmat':
                            pred_lhand_rotmat,
                        'pred_rhand_rotmat':
                            pred_rhand_rotmat,
                        'pred_face_rotmat':
                            pred_face_rotmat,
                        'pred_lhand_kp3d':
                            pred_output.lhand_joints,
                        'pred_rhand_kp3d':
                            pred_output.rhand_joints,
                        'pred_face_kp3d':
                            pred_output.face_joints,
                        'pred_lhand_kp2d':
                            pred_keypoints_2d[:, len_b_kp:len_b_kp + len_h_kp],
                        'pred_rhand_kp2d':
                            pred_keypoints_2d[:, len_b_kp + len_h_kp:len_b_kp + len_h_kp * 2],
                        'pred_face_kp2d':
                            pred_keypoints_2d[:, len_b_kp + len_h_kp * 2:len_b_kp + len_h_kp * 2 +
                                              len_f_kp],
                        'pred_feet_kp2d':
                            pred_keypoints_2d[:, len_b_kp + len_h_kp * 2 + len_f_kp:len_b_kp +
                                              len_h_kp * 2 + len_f_kp + len_feet_kp],
                    }
                )
                if cfg.MODEL.PyMAF.OPT_WRIST:
                    output.update(
                        {
                            'pred_orient_lh': pred_orient_lh,
                            'pred_shape_lh': pred_shape_lh,
                            'pred_orient_rh': pred_orient_rh,
                            'pred_shape_rh': pred_shape_rh,
                            'pred_cam_fa': pred_cam_fa,
                            'pred_cam_lh': pred_cam_lh,
                            'pred_cam_rh': pred_cam_rh,
                        }
                    )
                if cfg.MODEL.PyMAF.PRED_VIS_H:
                    output.update({'pred_vis_hands': pred_vis_hands})
        elif self.hand_only_mode:
            # hand mesh out
            assert pred_keypoints_2d.shape[1] == 21
            output.update(
                {
                    'theta': pred_cam,
                    'pred_cam': pred_cam,
                    'pred_rhand': pred_rhand,
                    'pred_rhand_rotmat': pred_rotmat_rh[:, 1:],
                    'pred_orient_rh': pred_orient_rh,
                    'pred_orient_rh_rotmat': pred_rotmat_rh[:, 0],
                    'verts_rh': pred_output.rhand_vertices,
                    'pred_cam_rh': pred_cam_rh,
                    'pred_shape_rh': pred_shape_rh,
                    'pred_rhand_kp3d': pred_output.rhand_joints,
                    'pred_rhand_kp2d': pred_keypoints_2d,
                }
            )
        elif self.face_only_mode:
            # face mesh out
            assert pred_keypoints_2d.shape[1] == 68
            output.update(
                {
                    'theta': pred_cam,
                    'pred_cam': pred_cam,
                    'pred_face': pred_face,
                    'pred_exp': pred_exp,
                    'pred_face_rotmat': pred_rotmat_fa[:, 1:],
                    'pred_orient_fa': pred_orient_fa,
                    'pred_orient_fa_rotmat': pred_rotmat_fa[:, 0],
                    'verts_fa': pred_output.flame_vertices,
                    'pred_cam_fa': pred_cam_fa,
                    'pred_shape_fa': pred_shape_fa,
                    'pred_face_kp3d': pred_output.face_joints,
                    'pred_face_kp2d': pred_keypoints_2d,
                }
            )
        return output


def get_attention_modules(
    module_keys, img_feature_dim_list, hidden_feat_dim, n_iter, num_attention_heads=1
):

    align_attention = nn.ModuleDict()
    for k in module_keys:
        align_attention[k] = nn.ModuleList()
        for i in range(n_iter):
            align_attention[k].append(
                get_att_block(
                    img_feature_dim=img_feature_dim_list[k][i],
                    hidden_feat_dim=hidden_feat_dim,
                    num_attention_heads=num_attention_heads
                )
            )

    return align_attention


def get_fusion_modules(module_keys, ma_feat_dim, grid_feat_dim, n_iter, out_feat_len):

    feat_fusion = nn.ModuleDict()
    for k in module_keys:
        feat_fusion[k] = nn.ModuleList()
        for i in range(n_iter):
            feat_fusion[k].append(nn.Linear(grid_feat_dim + ma_feat_dim[k], out_feat_len[k]))

    return feat_fusion


class PyMAF(nn.Module):
    """ PyMAF based Regression Network for Human Mesh Recovery / Full-body Mesh Recovery
    PyMAF: 3D Human Pose and Shape Regression with Pyramidal Mesh Alignment Feedback Loop, in ICCV, 2021
    PyMAF-X: Towards Well-aligned Full-body Model Regression from Monocular Images, arXiv:2207.06400, 2022
    """
    def __init__(
        self, smpl_mean_params=SMPL_MEAN_PARAMS, pretrained=True, device=torch.device('cuda')
    ):
        super().__init__()

        self.device = device

        self.smpl_mode = (cfg.MODEL.MESH_MODEL == 'smpl')
        self.smplx_mode = (cfg.MODEL.MESH_MODEL == 'smplx')

        assert cfg.TRAIN.BHF_MODE in [
            'body_only', 'hand_only', 'face_only', 'body_hand', 'full_body'
        ]
        self.hand_only_mode = (cfg.TRAIN.BHF_MODE == 'hand_only')
        self.face_only_mode = (cfg.TRAIN.BHF_MODE == 'face_only')
        self.body_hand_mode = (cfg.TRAIN.BHF_MODE == 'body_hand')
        self.full_body_mode = (cfg.TRAIN.BHF_MODE == 'full_body')

        bhf_names = []
        if cfg.TRAIN.BHF_MODE in ['body_only', 'body_hand', 'full_body']:
            bhf_names.append('body')
        if cfg.TRAIN.BHF_MODE in ['hand_only', 'body_hand', 'full_body']:
            bhf_names.append('hand')
        if cfg.TRAIN.BHF_MODE in ['face_only', 'full_body']:
            bhf_names.append('face')
        self.bhf_names = bhf_names

        self.part_module_names = {'body': {}, 'hand': {}, 'face': {}, 'link': {}}

        # the limb parts need to be handled
        if self.hand_only_mode:
            self.part_names = ['rhand']
        elif self.face_only_mode:
            self.part_names = ['face']
        elif self.body_hand_mode:
            self.part_names = ['lhand', 'rhand']
        elif self.full_body_mode:
            self.part_names = ['lhand', 'rhand', 'face']
        else:
            self.part_names = []

        # joint index info
        if not self.smpl_mode:
            h_root_idx = constants.HAND_NAMES.index('wrist')
            h_idx = constants.HAND_NAMES.index('middle1')
            f_idx = constants.FACIAL_LANDMARKS.index('nose_middle')
            self.hf_center_idx = {'lhand': h_idx, 'rhand': h_idx, 'face': f_idx}
            self.hf_root_idx = {'lhand': h_root_idx, 'rhand': h_root_idx, 'face': f_idx}

            lh_idx_coco = constants.COCO_KEYPOINTS.index('left_wrist')
            rh_idx_coco = constants.COCO_KEYPOINTS.index('right_wrist')
            f_idx_coco = constants.COCO_KEYPOINTS.index('nose')
            self.hf_root_idx_coco = {'lhand': lh_idx_coco, 'rhand': rh_idx_coco, 'face': f_idx_coco}

        # create parametric mesh models
        self.smpl_family = {}
        if self.hand_only_mode and cfg.MODEL.MESH_MODEL == 'mano':
            self.smpl_family['hand'] = SMPL_Family(model_type='mano')
            self.smpl_family['body'] = SMPL_Family(model_type='smplx')
        elif self.face_only_mode and cfg.MODEL.MESH_MODEL == 'flame':
            self.smpl_family['face'] = SMPL_Family(model_type='flame')
            self.smpl_family['body'] = SMPL_Family(model_type='smplx')
        else:
            self.smpl_family['body'] = SMPL_Family(
                model_type=cfg.MODEL.MESH_MODEL, all_gender=cfg.MODEL.ALL_GENDER
            )

        self.init_mesh_output = None
        self.batch_size = 1

        self.encoders = nn.ModuleDict()
        self.global_mode = not cfg.MODEL.PyMAF.MAF_ON

        # build encoders
        global_feat_dim = 2048
        bhf_ma_feat_dim = {}
        # encoder for the body part
        if 'body' in bhf_names:
            # if self.smplx_mode or 'hr' in cfg.MODEL.PyMAF.BACKBONE:
            if cfg.MODEL.PyMAF.BACKBONE == 'res50':
                body_encoder = get_resnet_encoder(
                    cfg, init_weight=(not cfg.MODEL.EVAL_MODE), global_mode=self.global_mode
                )
                body_sfeat_dim = list(cfg.POSE_RES_MODEL.EXTRA.NUM_DECONV_FILTERS)
            elif cfg.MODEL.PyMAF.BACKBONE == 'hr48':
                body_encoder = get_hrnet_encoder(
                    cfg, init_weight=(not cfg.MODEL.EVAL_MODE), global_mode=self.global_mode
                )
                body_sfeat_dim = list(cfg.HR_MODEL.EXTRA.STAGE4.NUM_CHANNELS)
                body_sfeat_dim.reverse()
                body_sfeat_dim = body_sfeat_dim[1:]
            else:
                raise NotImplementedError
            self.encoders['body'] = body_encoder
            self.part_module_names['body'].update({'encoders.body': self.encoders['body']})

            self.mesh_sampler = Mesh_Sampler(type='smpl')
            self.part_module_names['body'].update({'mesh_sampler': self.mesh_sampler})

            if not cfg.MODEL.PyMAF.GRID_FEAT:
                ma_feat_dim = self.mesh_sampler.Dmap.shape[0] * cfg.MODEL.PyMAF.MLP_DIM[-1]
            else:
                ma_feat_dim = 0
            bhf_ma_feat_dim['body'] = ma_feat_dim

            dp_feat_dim = body_sfeat_dim[-1]
            self.with_uv = cfg.LOSS.POINT_REGRESSION_WEIGHTS > 0
            if cfg.MODEL.PyMAF.AUX_SUPV_ON:
                assert cfg.MODEL.PyMAF.MAF_ON
                self.dp_head = IUV_predict_layer(feat_dim=dp_feat_dim)
                self.part_module_names['body'].update({'dp_head': self.dp_head})

        # encoders for the hand / face parts
        if 'hand' in self.bhf_names or 'face' in self.bhf_names:
            for hf in ['hand', 'face']:
                if hf in bhf_names:
                    if cfg.MODEL.PyMAF.HF_BACKBONE == 'res50':
                        self.encoders[hf] = get_resnet_encoder(
                            cfg,
                            init_weight=(not cfg.MODEL.EVAL_MODE),
                            global_mode=self.global_mode
                        )
                        self.part_module_names[hf].update({f'encoders.{hf}': self.encoders[hf]})
                        hf_sfeat_dim = list(cfg.POSE_RES_MODEL.EXTRA.NUM_DECONV_FILTERS)
                    else:
                        raise NotImplementedError

            if cfg.MODEL.PyMAF.HF_AUX_SUPV_ON:
                assert cfg.MODEL.PyMAF.MAF_ON
                self.dp_head_hf = nn.ModuleDict()
                if 'hand' in bhf_names:
                    self.dp_head_hf['hand'] = IUV_predict_layer(
                        feat_dim=hf_sfeat_dim[-1], mode='pncc'
                    )
                    self.part_module_names['hand'].update(
                        {'dp_head_hf.hand': self.dp_head_hf['hand']}
                    )
                if 'face' in bhf_names:
                    self.dp_head_hf['face'] = IUV_predict_layer(
                        feat_dim=hf_sfeat_dim[-1], mode='pncc'
                    )
                    self.part_module_names['face'].update(
                        {'dp_head_hf.face': self.dp_head_hf['face']}
                    )

            smpl2limb_vert_faces = get_partial_smpl()

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

        # grid points for grid feature extraction
        grid_size = 21
        xv, yv = torch.meshgrid(
            [torch.linspace(-1, 1, grid_size),
             torch.linspace(-1, 1, grid_size)]
        )
        grid_points = torch.stack([xv.reshape(-1), yv.reshape(-1)]).unsqueeze(0)
        self.register_buffer('grid_points', grid_points)
        grid_feat_dim = grid_size * grid_size * cfg.MODEL.PyMAF.MLP_DIM[-1]

        # the fusion of grid and mesh-aligned features
        self.fuse_grid_align = cfg.MODEL.PyMAF.GRID_ALIGN.USE_ATT or cfg.MODEL.PyMAF.GRID_ALIGN.USE_FC
        assert not (cfg.MODEL.PyMAF.GRID_ALIGN.USE_ATT and cfg.MODEL.PyMAF.GRID_ALIGN.USE_FC)

        if self.fuse_grid_align:
            self.att_starts = cfg.MODEL.PyMAF.GRID_ALIGN.ATT_STARTS
            n_iter_att = cfg.MODEL.PyMAF.N_ITER - self.att_starts
            att_feat_dim_idx = -cfg.MODEL.PyMAF.GRID_ALIGN.ATT_FEAT_IDX
            num_att_heads = cfg.MODEL.PyMAF.GRID_ALIGN.ATT_HEAD
            hidden_feat_dim = cfg.MODEL.PyMAF.MLP_DIM[att_feat_dim_idx]
            bhf_att_feat_dim = {'body': 2048}

        if 'hand' in self.bhf_names:
            self.mano_sampler = Mesh_Sampler(type='mano', level=1)
            self.mano_ds_len = self.mano_sampler.Dmap.shape[0]
            self.part_module_names['hand'].update({'mano_sampler': self.mano_sampler})

            bhf_ma_feat_dim.update({'hand': self.mano_ds_len * cfg.MODEL.PyMAF.HF_MLP_DIM[-1]})

            if self.fuse_grid_align:
                bhf_att_feat_dim.update({'hand': 1024})

        if 'face' in self.bhf_names:
            bhf_ma_feat_dim.update(
                {'face': len(constants.FACIAL_LANDMARKS) * cfg.MODEL.PyMAF.HF_MLP_DIM[-1]}
            )
            if self.fuse_grid_align:
                bhf_att_feat_dim.update({'face': 1024})

        # spatial alignment attention
        if cfg.MODEL.PyMAF.GRID_ALIGN.USE_ATT:
            hfimg_feat_dim_list = {}
            if 'body' in bhf_names:
                hfimg_feat_dim_list['body'] = body_sfeat_dim[-n_iter_att:]

            if 'hand' in self.bhf_names or 'face' in self.bhf_names:
                if 'hand' in bhf_names:
                    hfimg_feat_dim_list['hand'] = hf_sfeat_dim[-n_iter_att:]
                if 'face' in bhf_names:
                    hfimg_feat_dim_list['face'] = hf_sfeat_dim[-n_iter_att:]

            self.align_attention = get_attention_modules(
                bhf_names,
                hfimg_feat_dim_list,
                hidden_feat_dim,
                n_iter=n_iter_att,
                num_attention_heads=num_att_heads
            )

            for part in bhf_names:
                self.part_module_names[part].update(
                    {f'align_attention.{part}': self.align_attention[part]}
                )

        if self.fuse_grid_align:
            self.att_feat_reduce = get_fusion_modules(
                bhf_names,
                bhf_ma_feat_dim,
                grid_feat_dim,
                n_iter=n_iter_att,
                out_feat_len=bhf_att_feat_dim
            )
            for part in bhf_names:
                self.part_module_names[part].update(
                    {f'att_feat_reduce.{part}': self.att_feat_reduce[part]}
                )

        # build regressor for parameter prediction
        self.regressor = nn.ModuleList()
        for i in range(cfg.MODEL.PyMAF.N_ITER):
            ref_infeat_dim = 0
            if 'body' in self.bhf_names:
                if cfg.MODEL.PyMAF.MAF_ON:
                    if self.fuse_grid_align:
                        if i >= self.att_starts:
                            ref_infeat_dim = bhf_att_feat_dim['body']
                        elif i == 0 or cfg.MODEL.PyMAF.GRID_FEAT:
                            ref_infeat_dim = grid_feat_dim
                        else:
                            ref_infeat_dim = ma_feat_dim
                    else:
                        if i == 0 or cfg.MODEL.PyMAF.GRID_FEAT:
                            ref_infeat_dim = grid_feat_dim
                        else:
                            ref_infeat_dim = ma_feat_dim
                else:
                    ref_infeat_dim = global_feat_dim

            if self.smpl_mode:
                self.regressor.append(
                    Regressor(
                        feat_dim=ref_infeat_dim,
                        smpl_mean_params=smpl_mean_params,
                        use_cam_feats=cfg.MODEL.PyMAF.USE_CAM_FEAT,
                        smpl_models=self.smpl_family
                    )
                )
            else:
                if cfg.MODEL.PyMAF.MAF_ON:
                    if 'hand' in self.bhf_names or 'face' in self.bhf_names:
                        if i == 0:
                            feat_dim_hand = grid_feat_dim if 'hand' in self.bhf_names else None
                            feat_dim_face = grid_feat_dim if 'face' in self.bhf_names else None
                        else:
                            if self.fuse_grid_align:
                                feat_dim_hand = bhf_att_feat_dim[
                                    'hand'] if 'hand' in self.bhf_names else None
                                feat_dim_face = bhf_att_feat_dim[
                                    'face'] if 'face' in self.bhf_names else None
                            else:
                                feat_dim_hand = bhf_ma_feat_dim[
                                    'hand'] if 'hand' in self.bhf_names else None
                                feat_dim_face = bhf_ma_feat_dim[
                                    'face'] if 'face' in self.bhf_names else None
                    else:
                        feat_dim_hand = ref_infeat_dim
                        feat_dim_face = ref_infeat_dim
                else:
                    ref_infeat_dim = global_feat_dim
                    feat_dim_hand = global_feat_dim
                    feat_dim_face = global_feat_dim

                self.regressor.append(
                    Regressor(
                        feat_dim=ref_infeat_dim,
                        smpl_mean_params=smpl_mean_params,
                        use_cam_feats=cfg.MODEL.PyMAF.USE_CAM_FEAT,
                        feat_dim_hand=feat_dim_hand,
                        feat_dim_face=feat_dim_face,
                        bhf_names=bhf_names,
                        smpl_models=self.smpl_family
                    )
                )

            # assign sub-regressor to each part
            for dec_name, dec_module in self.regressor[-1].named_children():
                if 'hand' in dec_name:
                    self.part_module_names['hand'].update(
                        {'regressor.{}.{}.'.format(len(self.regressor) - 1, dec_name): dec_module}
                    )
                elif 'face' in dec_name or 'head' in dec_name or 'exp' in dec_name:
                    self.part_module_names['face'].update(
                        {'regressor.{}.{}.'.format(len(self.regressor) - 1, dec_name): dec_module}
                    )
                elif 'res' in dec_name or 'vis' in dec_name:
                    self.part_module_names['link'].update(
                        {'regressor.{}.{}.'.format(len(self.regressor) - 1, dec_name): dec_module}
                    )
                elif 'body' in self.part_module_names:
                    self.part_module_names['body'].update(
                        {'regressor.{}.{}.'.format(len(self.regressor) - 1, dec_name): dec_module}
                    )

        # mesh-aligned feature extractor
        self.maf_extractor = nn.ModuleDict()
        for part in bhf_names:
            self.maf_extractor[part] = nn.ModuleList()
            filter_channels_default = cfg.MODEL.PyMAF.MLP_DIM if part == 'body' else cfg.MODEL.PyMAF.HF_MLP_DIM
            sfeat_dim = body_sfeat_dim if part == 'body' else hf_sfeat_dim
            for i in range(cfg.MODEL.PyMAF.N_ITER):
                for f_i, f_dim in enumerate(filter_channels_default):
                    if sfeat_dim[i] > f_dim:
                        filter_start = f_i
                        break
                filter_channels = [sfeat_dim[i]] + filter_channels_default[filter_start:]

                if cfg.MODEL.PyMAF.GRID_ALIGN.USE_ATT and i >= self.att_starts:
                    self.maf_extractor[part].append(
                        MAF_Extractor(
                            filter_channels=filter_channels_default[att_feat_dim_idx:],
                            iwp_cam_mode=cfg.MODEL.USE_IWP_CAM
                        )
                    )
                else:
                    self.maf_extractor[part].append(
                        MAF_Extractor(
                            filter_channels=filter_channels, iwp_cam_mode=cfg.MODEL.USE_IWP_CAM
                        )
                    )
            self.part_module_names[part].update({f'maf_extractor.{part}': self.maf_extractor[part]})

        # check all modules have been added to part_module_names
        model_dict_all = dict.fromkeys(self.state_dict().keys())
        for key in self.part_module_names.keys():
            for name in list(model_dict_all.keys()):
                for k in self.part_module_names[key].keys():
                    if name.startswith(k):
                        del model_dict_all[name]
                # if name.startswith('regressor.') and '.smpl.' in name:
                #     del model_dict_all[name]
                # if name.startswith('regressor.') and '.mano.' in name:
                #     del model_dict_all[name]
                if name.startswith('regressor.') and '.init_' in name:
                    del model_dict_all[name]
                if name == 'grid_points':
                    del model_dict_all[name]
        assert (len(model_dict_all.keys()) == 0)

    def init_mesh(self, batch_size, J_regressor=None, rw_cam={}):
        """ initialize the mesh model with default poses and shapes
        """
        if self.init_mesh_output is None or self.batch_size != batch_size:
            self.init_mesh_output = self.regressor[0](
                torch.zeros(batch_size), J_regressor=J_regressor, rw_cam=rw_cam, init_mode=True
            )
            self.batch_size = batch_size
        return self.init_mesh_output

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(
                    self.inplanes,
                    planes * block.expansion,
                    kernel_size=1,
                    stride=stride,
                    bias=False
                ),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def _make_deconv_layer(self, num_layers, num_filters, num_kernels):
        """
        Deconv_layer used in Simple Baselines:
        Xiao et al. Simple Baselines for Human Pose Estimation and Tracking
        https://github.com/microsoft/human-pose-estimation.pytorch
        """
        assert num_layers == len(num_filters), \
            'ERROR: num_deconv_layers is different len(num_deconv_filters)'
        assert num_layers == len(num_kernels), \
            'ERROR: num_deconv_layers is different len(num_deconv_filters)'

        def _get_deconv_cfg(deconv_kernel, index):
            if deconv_kernel == 4:
                padding = 1
                output_padding = 0
            elif deconv_kernel == 3:
                padding = 1
                output_padding = 1
            elif deconv_kernel == 2:
                padding = 0
                output_padding = 0

            return deconv_kernel, padding, output_padding

        layers = []
        for i in range(num_layers):
            kernel, padding, output_padding = _get_deconv_cfg(num_kernels[i], i)

            planes = num_filters[i]
            layers.append(
                nn.ConvTranspose2d(
                    in_channels=self.inplanes,
                    out_channels=planes,
                    kernel_size=kernel,
                    stride=2,
                    padding=padding,
                    output_padding=output_padding,
                    bias=self.deconv_with_bias
                )
            )
            layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM))
            layers.append(nn.ReLU(inplace=True))
            self.inplanes = planes

        return nn.Sequential(*layers)

    def to(self, *args, **kwargs):
        super().to(*args, **kwargs)
        for m in ['body', 'hand', 'face']:
            if m in self.smpl_family:
                self.smpl_family[m].model.to(*args, **kwargs)
        return self

    def cuda(self, *args, **kwargs):
        super().cuda(*args, **kwargs)
        for m in ['body', 'hand', 'face']:
            if m in self.smpl_family:
                self.smpl_family[m].model.cuda(*args, **kwargs)
        return self

    def forward(self, batch={}, J_regressor=None, rw_cam={}):
        '''
        Args:
            batch: input dictionary, including 
                   images: 'img_{part}', for part in body, hand, and face if applicable
                   inversed affine transformation for the cropping of hand/face images: '{part}_theta_inv' for part in lhand, rhand, and face if applicable
            J_regressor: joint regression matrix
            rw_cam: real-world camera information, applied when cfg.MODEL.USE_IWP_CAM is False
        Returns:
            out_dict: the list containing the predicted parameters
            vis_feat_list: the list containing features for visualization
        '''

        # batch keys: ['img_body', 'orig_height', 'orig_width', 'person_id', 'img_lhand',
        # 'lhand_theta_inv', 'img_rhand', 'rhand_theta_inv', 'img_face', 'face_theta_inv']

        # extract spatial features or global features
        # run encoder for body
        if 'body' in self.bhf_names:
            img_body = batch['img_body']
            batch_size = img_body.shape[0]
            s_feat_body, g_feat = self.encoders['body'](batch['img_body'])
            if cfg.MODEL.PyMAF.MAF_ON:
                assert len(s_feat_body) == cfg.MODEL.PyMAF.N_ITER

        # run encoders for hand / face
        if 'hand' in self.bhf_names or 'face' in self.bhf_names:
            limb_feat_dict = {}
            limb_gfeat_dict = {}
            if 'face' in self.bhf_names:
                img_face = batch['img_face']
                batch_size = img_face.shape[0]
                limb_feat_dict['face'], limb_gfeat_dict['face'] = self.encoders['face'](img_face)

            if 'hand' in self.bhf_names:
                if 'lhand' in self.part_names:
                    img_rhand = batch['img_rhand']
                    batch_size = img_rhand.shape[0]
                    # flip left hand images
                    img_lhand = torch.flip(batch['img_lhand'], [3])
                    img_hands = torch.cat([img_rhand, img_lhand])
                    s_feat_hands, g_feat_hands = self.encoders['hand'](img_hands)
                    limb_feat_dict['rhand'] = [feat[:batch_size] for feat in s_feat_hands]
                    limb_feat_dict['lhand'] = [feat[batch_size:] for feat in s_feat_hands]
                    if g_feat_hands is not None:
                        limb_gfeat_dict['rhand'] = g_feat_hands[:batch_size]
                        limb_gfeat_dict['lhand'] = g_feat_hands[batch_size:]
                else:
                    img_rhand = batch['img_rhand']
                    batch_size = img_rhand.shape[0]
                    limb_feat_dict['rhand'], limb_gfeat_dict['rhand'] = self.encoders['hand'](
                        img_rhand
                    )

            if cfg.MODEL.PyMAF.MAF_ON:
                for k in limb_feat_dict.keys():
                    assert len(limb_feat_dict[k]) == cfg.MODEL.PyMAF.N_ITER

        out_dict = {}

        # grid-pattern points
        grid_points = torch.transpose(self.grid_points.expand(batch_size, -1, -1), 1, 2)

        # initial parameters
        mesh_output = self.init_mesh(batch_size, J_regressor, rw_cam)

        out_dict['mesh_out'] = [mesh_output]
        out_dict['dp_out'] = []

        # for visulization
        vis_feat_list = []

        # dense prediction during training
        if not cfg.MODEL.EVAL_MODE:
            if 'body' in self.bhf_names:
                if cfg.MODEL.PyMAF.AUX_SUPV_ON:
                    iuv_out_dict = self.dp_head(s_feat_body[-1])
                    out_dict['dp_out'].append(iuv_out_dict)
            elif self.hand_only_mode:
                if cfg.MODEL.PyMAF.HF_AUX_SUPV_ON:
                    out_dict['rhand_dpout'] = []
                    dphand_out_dict = self.dp_head_hf['hand'](limb_feat_dict['rhand'][-1])
                    out_dict['rhand_dpout'].append(dphand_out_dict)
            elif self.face_only_mode:
                if cfg.MODEL.PyMAF.HF_AUX_SUPV_ON:
                    out_dict['face_dpout'] = []
                    dpface_out_dict = self.dp_head_hf['face'](limb_feat_dict['face'][-1])
                    out_dict['face_dpout'].append(dpface_out_dict)

        # parameter predictions
        for rf_i in range(cfg.MODEL.PyMAF.N_ITER):
            current_states = {}
            if 'body' in self.bhf_names:
                pred_cam = mesh_output['pred_cam'].detach()
                pred_shape = mesh_output['pred_shape'].detach()
                pred_pose = mesh_output['pred_pose'].detach()

                current_states['init_cam'] = pred_cam
                current_states['init_shape'] = pred_shape
                current_states['init_pose'] = pred_pose

                pred_smpl_verts = mesh_output['verts'].detach()

                if cfg.MODEL.PyMAF.MAF_ON:
                    s_feat_i = s_feat_body[rf_i]

            # re-project mesh on the image plane
            if self.hand_only_mode:
                pred_cam = mesh_output['pred_cam'].detach()
                pred_rhand_v = self.mano_sampler(mesh_output['verts_rh'])
                pred_rhand_proj = projection(
                    pred_rhand_v, {
                        **rw_cam, 'cam_sxy': pred_cam
                    }, iwp_mode=cfg.MODEL.USE_IWP_CAM
                )
                if cfg.MODEL.USE_IWP_CAM:
                    pred_rhand_proj = pred_rhand_proj / (224. / 2.)
                else:
                    pred_rhand_proj = j2d_processing(pred_rhand_proj, rw_cam['kps_transf'])
                proj_hf_center = {
                    'rhand':
                        mesh_output['pred_rhand_kp2d'][:, self.hf_root_idx['rhand']].unsqueeze(1)
                }
                proj_hf_pts = {
                    'rhand': torch.cat([proj_hf_center['rhand'], pred_rhand_proj], dim=1)
                }
            elif self.face_only_mode:
                pred_cam = mesh_output['pred_cam'].detach()
                pred_face_v = mesh_output['pred_face_kp3d']
                pred_face_proj = projection(
                    pred_face_v, {
                        **rw_cam, 'cam_sxy': pred_cam
                    }, iwp_mode=cfg.MODEL.USE_IWP_CAM
                )
                if cfg.MODEL.USE_IWP_CAM:
                    pred_face_proj = pred_face_proj / (224. / 2.)
                else:
                    pred_face_proj = j2d_processing(pred_face_proj, rw_cam['kps_transf'])
                proj_hf_center = {
                    'face': mesh_output['pred_face_kp2d'][:, self.hf_root_idx['face']].unsqueeze(1)
                }
                proj_hf_pts = {'face': torch.cat([proj_hf_center['face'], pred_face_proj], dim=1)}
            elif self.body_hand_mode:
                pred_lhand_v = self.mano_sampler(pred_smpl_verts[:, self.smpl2lhand])
                pred_rhand_v = self.mano_sampler(pred_smpl_verts[:, self.smpl2rhand])
                pred_hand_v = torch.cat([pred_lhand_v, pred_rhand_v], dim=1)
                pred_hand_proj = projection(
                    pred_hand_v, {
                        **rw_cam, 'cam_sxy': pred_cam
                    }, iwp_mode=cfg.MODEL.USE_IWP_CAM
                )
                if cfg.MODEL.USE_IWP_CAM:
                    pred_hand_proj = pred_hand_proj / (224. / 2.)
                else:
                    pred_hand_proj = j2d_processing(pred_hand_proj, rw_cam['kps_transf'])

                proj_hf_center = {
                    'lhand':
                        mesh_output['pred_lhand_kp2d'][:, self.hf_root_idx['lhand']].unsqueeze(1),
                    'rhand':
                        mesh_output['pred_rhand_kp2d'][:, self.hf_root_idx['rhand']].unsqueeze(1),
                }
                proj_hf_pts = {
                    'lhand':
                        torch.cat(
                            [proj_hf_center['lhand'], pred_hand_proj[:, :self.mano_ds_len]], dim=1
                        ),
                    'rhand':
                        torch.cat(
                            [proj_hf_center['rhand'], pred_hand_proj[:, self.mano_ds_len:]], dim=1
                        ),
                }
            elif self.full_body_mode:
                pred_lhand_v = self.mano_sampler(pred_smpl_verts[:, self.smpl2lhand])
                pred_rhand_v = self.mano_sampler(pred_smpl_verts[:, self.smpl2rhand])
                pred_hand_v = torch.cat([pred_lhand_v, pred_rhand_v], dim=1)
                pred_hand_proj = projection(
                    pred_hand_v, {
                        **rw_cam, 'cam_sxy': pred_cam
                    }, iwp_mode=cfg.MODEL.USE_IWP_CAM
                )
                if cfg.MODEL.USE_IWP_CAM:
                    pred_hand_proj = pred_hand_proj / (224. / 2.)
                else:
                    pred_hand_proj = j2d_processing(pred_hand_proj, rw_cam['kps_transf'])

                proj_hf_center = {
                    'lhand':
                        mesh_output['pred_lhand_kp2d'][:, self.hf_root_idx['lhand']].unsqueeze(1),
                    'rhand':
                        mesh_output['pred_rhand_kp2d'][:, self.hf_root_idx['rhand']].unsqueeze(1),
                    'face':
                        mesh_output['pred_face_kp2d'][:, self.hf_root_idx['face']].unsqueeze(1)
                }
                proj_hf_pts = {
                    'lhand':
                        torch.cat(
                            [proj_hf_center['lhand'], pred_hand_proj[:, :self.mano_ds_len]], dim=1
                        ),
                    'rhand':
                        torch.cat(
                            [proj_hf_center['rhand'], pred_hand_proj[:, self.mano_ds_len:]], dim=1
                        ),
                    'face':
                        torch.cat([proj_hf_center['face'], mesh_output['pred_face_kp2d']], dim=1)
                }

            # extract mesh-aligned features for the hand / face part
            if 'hand' in self.bhf_names or 'face' in self.bhf_names:
                limb_rf_i = rf_i
                hand_face_feat = {}

                for hf_i, part_name in enumerate(self.part_names):
                    if 'hand' in part_name:
                        hf_key = 'hand'
                    elif 'face' in part_name:
                        hf_key = 'face'

                    if cfg.MODEL.PyMAF.MAF_ON:
                        if cfg.MODEL.PyMAF.HF_BACKBONE == 'res50':
                            limb_feat_i = limb_feat_dict[part_name][limb_rf_i]
                        else:
                            raise NotImplementedError

                        limb_reduce_dim = (not self.fuse_grid_align) or (rf_i < self.att_starts)

                        if limb_rf_i == 0 or cfg.MODEL.PyMAF.GRID_FEAT:
                            limb_ref_feat_ctd = self.maf_extractor[hf_key][limb_rf_i].sampling(
                                grid_points, im_feat=limb_feat_i, reduce_dim=limb_reduce_dim
                            )
                        else:
                            if self.hand_only_mode or self.face_only_mode:
                                proj_hf_pts_crop = proj_hf_pts[part_name][:, :, :2]

                                proj_hf_v_center = proj_hf_pts_crop[:, 0].unsqueeze(1)

                                if cfg.MODEL.PyMAF.HF_BOX_CENTER:
                                    part_box_ul = torch.min(proj_hf_pts_crop, dim=1)[0].unsqueeze(1)
                                    part_box_br = torch.max(proj_hf_pts_crop, dim=1)[0].unsqueeze(1)
                                    part_box_center = (part_box_ul + part_box_br) / 2.
                                    proj_hf_pts_crop_ctd = proj_hf_pts_crop[:, 1:] - part_box_center
                                else:
                                    proj_hf_pts_crop_ctd = proj_hf_pts_crop[:, 1:]

                            elif self.full_body_mode or self.body_hand_mode:
                                # convert projection points to the space of cropped hand/face images
                                theta_i_inv = batch[f'{part_name}_theta_inv']
                                proj_hf_pts_crop = torch.bmm(
                                    theta_i_inv,
                                    homo_vector(proj_hf_pts[part_name][:, :, :2]).permute(0, 2, 1)
                                ).permute(0, 2, 1)

                                if part_name == 'lhand':
                                    flip_x = torch.tensor([-1, 1])[None,
                                                                   None, :].to(proj_hf_pts_crop)
                                    proj_hf_pts_crop *= flip_x

                                if cfg.MODEL.PyMAF.HF_BOX_CENTER:
                                    # align projection points with the cropped image center
                                    part_box_ul = torch.min(proj_hf_pts_crop, dim=1)[0].unsqueeze(1)
                                    part_box_br = torch.max(proj_hf_pts_crop, dim=1)[0].unsqueeze(1)
                                    part_box_center = (part_box_ul + part_box_br) / 2.
                                    proj_hf_pts_crop_ctd = proj_hf_pts_crop[:, 1:] - part_box_center
                                else:
                                    proj_hf_pts_crop_ctd = proj_hf_pts_crop[:, 1:]

                                # 0 is the root point
                                proj_hf_v_center = proj_hf_pts_crop[:, 0].unsqueeze(1)

                            limb_ref_feat_ctd = self.maf_extractor[hf_key][limb_rf_i].sampling(
                                proj_hf_pts_crop_ctd.detach(),
                                im_feat=limb_feat_i,
                                reduce_dim=limb_reduce_dim
                            )

                        if self.fuse_grid_align and limb_rf_i >= self.att_starts:

                            limb_grid_feature_ctd = self.maf_extractor[hf_key][limb_rf_i].sampling(
                                grid_points, im_feat=limb_feat_i, reduce_dim=limb_reduce_dim
                            )
                            limb_grid_ref_feat_ctd = torch.cat(
                                [limb_grid_feature_ctd, limb_ref_feat_ctd], dim=-1
                            ).permute(0, 2, 1)

                            if cfg.MODEL.PyMAF.GRID_ALIGN.USE_ATT:
                                att_ref_feat_ctd = self.align_attention[hf_key][
                                    limb_rf_i - self.att_starts](limb_grid_ref_feat_ctd)[0]
                            elif cfg.MODEL.PyMAF.GRID_ALIGN.USE_FC:
                                att_ref_feat_ctd = limb_grid_ref_feat_ctd

                            att_ref_feat_ctd = self.maf_extractor[hf_key][limb_rf_i].reduce_dim(
                                att_ref_feat_ctd.permute(0, 2, 1)
                            ).view(batch_size, -1)
                            limb_ref_feat_ctd = self.att_feat_reduce[hf_key][
                                limb_rf_i - self.att_starts](att_ref_feat_ctd)

                        else:
                            # limb_ref_feat = limb_ref_feat.view(batch_size, -1)
                            limb_ref_feat_ctd = limb_ref_feat_ctd.view(batch_size, -1)
                        hand_face_feat[part_name] = limb_ref_feat_ctd
                    else:
                        hand_face_feat[part_name] = limb_gfeat_dict[part_name]

            # extract mesh-aligned features for the body part
            if 'body' in self.bhf_names:
                if cfg.MODEL.PyMAF.MAF_ON:
                    reduce_dim = (not self.fuse_grid_align) or (rf_i < self.att_starts)
                    if rf_i == 0 or cfg.MODEL.PyMAF.GRID_FEAT:
                        ref_feature = self.maf_extractor['body'][rf_i].sampling(
                            grid_points, im_feat=s_feat_i, reduce_dim=reduce_dim
                        )
                    else:
                        # TODO: use a more sparse SMPL implementation (with 431 vertices) for acceleration
                        pred_smpl_verts_ds = self.mesh_sampler.downsample(
                            pred_smpl_verts
                        )    # [B, 431, 3]
                        ref_feature = self.maf_extractor['body'][rf_i](
                            pred_smpl_verts_ds,
                            im_feat=s_feat_i,
                            cam={
                                **rw_cam, 'cam_sxy': pred_cam
                            },
                            add_att=True,
                            reduce_dim=reduce_dim
                        )    # [B, 431 * n_feat]

                    if self.fuse_grid_align and rf_i >= self.att_starts:
                        if rf_i > 0 and not cfg.MODEL.PyMAF.GRID_FEAT:
                            grid_feature = self.maf_extractor['body'][rf_i].sampling(
                                grid_points, im_feat=s_feat_i, reduce_dim=reduce_dim
                            )
                            grid_ref_feat = torch.cat([grid_feature, ref_feature], dim=-1)
                        else:
                            grid_ref_feat = ref_feature
                        grid_ref_feat = grid_ref_feat.permute(0, 2, 1)

                        if cfg.MODEL.PyMAF.GRID_ALIGN.USE_ATT:
                            att_ref_feat = self.align_attention['body'][
                                rf_i - self.att_starts](grid_ref_feat)[0]
                        elif cfg.MODEL.PyMAF.GRID_ALIGN.USE_FC:
                            att_ref_feat = grid_ref_feat

                        att_ref_feat = self.maf_extractor['body'][rf_i].reduce_dim(
                            att_ref_feat.permute(0, 2, 1)
                        )
                        att_ref_feat = att_ref_feat.view(batch_size, -1)

                        ref_feature = self.att_feat_reduce['body'][rf_i -
                                                                   self.att_starts](att_ref_feat)
                    else:
                        ref_feature = ref_feature.view(batch_size, -1)
                else:
                    ref_feature = g_feat
            else:
                ref_feature = None

            if not self.smpl_mode:
                if self.hand_only_mode:
                    current_states['xc_rhand'] = hand_face_feat['rhand']
                elif self.face_only_mode:
                    current_states['xc_face'] = hand_face_feat['face']
                elif self.body_hand_mode:
                    current_states['xc_lhand'] = hand_face_feat['lhand']
                    current_states['xc_rhand'] = hand_face_feat['rhand']
                elif self.full_body_mode:
                    current_states['xc_lhand'] = hand_face_feat['lhand']
                    current_states['xc_rhand'] = hand_face_feat['rhand']
                    current_states['xc_face'] = hand_face_feat['face']

                if rf_i > 0:
                    for part in self.part_names:
                        current_states[f'init_{part}'] = mesh_output[f'pred_{part}'].detach()
                        if part == 'face':
                            current_states['init_exp'] = mesh_output['pred_exp'].detach()
                    if self.hand_only_mode:
                        current_states['init_shape_rh'] = mesh_output['pred_shape_rh'].detach()
                        current_states['init_orient_rh'] = mesh_output['pred_orient_rh'].detach()
                        current_states['init_cam_rh'] = mesh_output['pred_cam_rh'].detach()
                    elif self.face_only_mode:
                        current_states['init_shape_fa'] = mesh_output['pred_shape_fa'].detach()
                        current_states['init_orient_fa'] = mesh_output['pred_orient_fa'].detach()
                        current_states['init_cam_fa'] = mesh_output['pred_cam_fa'].detach()
                    elif self.full_body_mode or self.body_hand_mode:
                        if cfg.MODEL.PyMAF.OPT_WRIST:
                            current_states['init_shape_lh'] = mesh_output['pred_shape_lh'].detach()
                            current_states['init_orient_lh'] = mesh_output['pred_orient_lh'].detach(
                            )
                            current_states['init_cam_lh'] = mesh_output['pred_cam_lh'].detach()

                            current_states['init_shape_rh'] = mesh_output['pred_shape_rh'].detach()
                            current_states['init_orient_rh'] = mesh_output['pred_orient_rh'].detach(
                            )
                            current_states['init_cam_rh'] = mesh_output['pred_cam_rh'].detach()

            # update mesh parameters
            mesh_output = self.regressor[rf_i](
                ref_feature,
                n_iter=1,
                J_regressor=J_regressor,
                rw_cam=rw_cam,
                global_iter=rf_i,
                **current_states
            )

            out_dict['mesh_out'].append(mesh_output)

        return out_dict, vis_feat_list


def pymaf_net(smpl_mean_params, pretrained=True, device=torch.device('cuda')):
    """ Constructs an PyMAF model with ResNet50 backbone.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = PyMAF(smpl_mean_params, pretrained, device)
    return model