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