| 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 '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, 1024) |
| self.drop1_hand = nn.Dropout() |
| self.fc2_hand = nn.Linear(1024, 1024) |
| self.drop2_hand = nn.Dropout() |
|
|
| |
| self.decrhand = nn.Linear(1024, 15 * 6) |
| |
| 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_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: |
| |
| 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) |
| ) |
| if cfg.MODEL.PyMAF.PRED_VIS_H: |
| pred_vis_hands = None |
|
|
| |
| if self.smplx_mode or 'hand' in self.bhf_names: |
| if 'init_lhand' not in kwargs: |
| |
| |
| 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.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: |
| |
| 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) |
| ) |
| 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: |
| |
| 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) |
|
|
| |
| opt_lelbow = torch.bmm(pred_rotmat_body[:, 18], lelbow_twist) |
| |
| 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) |
| ) |
| 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) |
| ) |
| 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) |
| ) |
| assert pred_rotmat.shape[1] == 24 |
| else: |
| pred_rotmat = rot6d_to_rotmat( |
| pred_pose.reshape(batch_size, -1, 6) |
| ) |
| assert pred_rotmat.shape[1] == 24 |
|
|
| |
| 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) |
|
|
| |
| 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.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: |
| |
| 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.smplx_mode: |
| |
| 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_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: |
| |
| 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: |
| |
| 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': {}} |
|
|
| |
| 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 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} |
|
|
| |
| 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 |
|
|
| |
| global_feat_dim = 2048 |
| bhf_ma_feat_dim = {} |
| |
| if 'body' in bhf_names: |
| |
| 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}) |
|
|
| |
| 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_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] |
|
|
| |
| 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}) |
|
|
| |
| 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]} |
| ) |
|
|
| |
| 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 |
| ) |
| ) |
|
|
| |
| 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} |
| ) |
|
|
| |
| 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]}) |
|
|
| |
| 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 '.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 |
| ''' |
|
|
| |
| |
|
|
| |
| |
| 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 |
|
|
| |
| 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] |
| |
| 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_points = torch.transpose(self.grid_points.expand(batch_size, -1, -1), 1, 2) |
|
|
| |
| mesh_output = self.init_mesh(batch_size, J_regressor, rw_cam) |
|
|
| out_dict['mesh_out'] = [mesh_output] |
| out_dict['dp_out'] = [] |
|
|
| |
| vis_feat_list = [] |
|
|
| |
| 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) |
|
|
| |
| 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] |
|
|
| |
| 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) |
| } |
|
|
| |
| 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: |
| |
| 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: |
| |
| 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:] |
|
|
| |
| 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_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] |
|
|
| |
| 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: |
| |
| pred_smpl_verts_ds = self.mesh_sampler.downsample( |
| pred_smpl_verts |
| ) |
| 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 |
| ) |
|
|
| 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() |
|
|
| |
| 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 |
|
|
