Spaces:
Running
on
L40S
Running
on
L40S
| 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 | |