Spaces:
Running
on
L40S
Running
on
L40S
| import torch | |
| import torch.nn as nn | |
| import numpy as np | |
| from lib.pymaf.utils.geometry import rot6d_to_rotmat, projection, rotation_matrix_to_angle_axis | |
| from .maf_extractor import MAF_Extractor | |
| from .smpl import SMPL, SMPL_MODEL_DIR, SMPL_MEAN_PARAMS, H36M_TO_J14 | |
| from .hmr import ResNet_Backbone | |
| from .res_module import IUV_predict_layer | |
| from lib.common.config import cfg | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| BN_MOMENTUM = 0.1 | |
| class Regressor(nn.Module): | |
| def __init__(self, feat_dim, smpl_mean_params): | |
| super().__init__() | |
| npose = 24 * 6 | |
| self.fc1 = nn.Linear(feat_dim + npose + 13, 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) | |
| self.smpl = SMPL(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) | |
| def forward(self, | |
| x, | |
| init_pose=None, | |
| init_shape=None, | |
| init_cam=None, | |
| n_iter=1, | |
| J_regressor=None): | |
| batch_size = x.shape[0] | |
| if init_pose is None: | |
| init_pose = self.init_pose.expand(batch_size, -1) | |
| if init_shape is None: | |
| init_shape = self.init_shape.expand(batch_size, -1) | |
| if init_cam is None: | |
| init_cam = self.init_cam.expand(batch_size, -1) | |
| pred_pose = init_pose | |
| pred_shape = init_shape | |
| pred_cam = init_cam | |
| for i in range(n_iter): | |
| xc = torch.cat([x, pred_pose, pred_shape, pred_cam], 1) | |
| xc = self.fc1(xc) | |
| xc = self.drop1(xc) | |
| xc = self.fc2(xc) | |
| xc = self.drop2(xc) | |
| pred_pose = self.decpose(xc) + pred_pose | |
| pred_shape = self.decshape(xc) + pred_shape | |
| pred_cam = self.deccam(xc) + pred_cam | |
| pred_rotmat = rot6d_to_rotmat(pred_pose).view(batch_size, 24, 3, 3) | |
| pred_output = self.smpl(betas=pred_shape, | |
| body_pose=pred_rotmat[:, 1:], | |
| global_orient=pred_rotmat[:, 0].unsqueeze(1), | |
| pose2rot=False) | |
| pred_vertices = pred_output.vertices | |
| pred_joints = pred_output.joints | |
| pred_smpl_joints = pred_output.smpl_joints | |
| pred_keypoints_2d = projection(pred_joints, pred_cam) | |
| pose = rotation_matrix_to_angle_axis(pred_rotmat.reshape(-1, 3, | |
| 3)).reshape( | |
| -1, 72) | |
| if J_regressor is not None: | |
| pred_joints = torch.matmul(J_regressor, pred_vertices) | |
| pred_pelvis = pred_joints[:, [0], :].clone() | |
| pred_joints = pred_joints[:, H36M_TO_J14, :] | |
| pred_joints = pred_joints - pred_pelvis | |
| output = { | |
| 'theta': torch.cat([pred_cam, pred_shape, pose], dim=1), | |
| 'verts': pred_vertices, | |
| 'kp_2d': pred_keypoints_2d, | |
| 'kp_3d': pred_joints, | |
| 'smpl_kp_3d': pred_smpl_joints, | |
| 'rotmat': pred_rotmat, | |
| 'pred_cam': pred_cam, | |
| 'pred_shape': pred_shape, | |
| 'pred_pose': pred_pose, | |
| } | |
| return output | |
| def forward_init(self, | |
| x, | |
| init_pose=None, | |
| init_shape=None, | |
| init_cam=None, | |
| n_iter=1, | |
| J_regressor=None): | |
| batch_size = x.shape[0] | |
| if init_pose is None: | |
| init_pose = self.init_pose.expand(batch_size, -1) | |
| if init_shape is None: | |
| init_shape = self.init_shape.expand(batch_size, -1) | |
| if init_cam is None: | |
| init_cam = self.init_cam.expand(batch_size, -1) | |
| pred_pose = init_pose | |
| pred_shape = init_shape | |
| pred_cam = init_cam | |
| pred_rotmat = rot6d_to_rotmat(pred_pose.contiguous()).view( | |
| batch_size, 24, 3, 3) | |
| pred_output = self.smpl(betas=pred_shape, | |
| body_pose=pred_rotmat[:, 1:], | |
| global_orient=pred_rotmat[:, 0].unsqueeze(1), | |
| pose2rot=False) | |
| pred_vertices = pred_output.vertices | |
| pred_joints = pred_output.joints | |
| pred_smpl_joints = pred_output.smpl_joints | |
| pred_keypoints_2d = projection(pred_joints, pred_cam) | |
| pose = rotation_matrix_to_angle_axis(pred_rotmat.reshape(-1, 3, | |
| 3)).reshape( | |
| -1, 72) | |
| if J_regressor is not None: | |
| pred_joints = torch.matmul(J_regressor, pred_vertices) | |
| pred_pelvis = pred_joints[:, [0], :].clone() | |
| pred_joints = pred_joints[:, H36M_TO_J14, :] | |
| pred_joints = pred_joints - pred_pelvis | |
| output = { | |
| 'theta': torch.cat([pred_cam, pred_shape, pose], dim=1), | |
| 'verts': pred_vertices, | |
| 'kp_2d': pred_keypoints_2d, | |
| 'kp_3d': pred_joints, | |
| 'smpl_kp_3d': pred_smpl_joints, | |
| 'rotmat': pred_rotmat, | |
| 'pred_cam': pred_cam, | |
| 'pred_shape': pred_shape, | |
| 'pred_pose': pred_pose, | |
| } | |
| return output | |
| class PyMAF(nn.Module): | |
| """ PyMAF based Deep Regressor for Human Mesh Recovery | |
| PyMAF: 3D Human Pose and Shape Regression with Pyramidal Mesh Alignment Feedback Loop, in ICCV, 2021 | |
| """ | |
| def __init__(self, smpl_mean_params=SMPL_MEAN_PARAMS, pretrained=True): | |
| super().__init__() | |
| self.feature_extractor = ResNet_Backbone( | |
| model=cfg.MODEL.PyMAF.BACKBONE, pretrained=pretrained) | |
| # deconv layers | |
| self.inplanes = self.feature_extractor.inplanes | |
| self.deconv_with_bias = cfg.RES_MODEL.DECONV_WITH_BIAS | |
| self.deconv_layers = self._make_deconv_layer( | |
| cfg.RES_MODEL.NUM_DECONV_LAYERS, | |
| cfg.RES_MODEL.NUM_DECONV_FILTERS, | |
| cfg.RES_MODEL.NUM_DECONV_KERNELS, | |
| ) | |
| self.maf_extractor = nn.ModuleList() | |
| for _ in range(cfg.MODEL.PyMAF.N_ITER): | |
| self.maf_extractor.append(MAF_Extractor()) | |
| ma_feat_len = self.maf_extractor[-1].Dmap.shape[ | |
| 0] * cfg.MODEL.PyMAF.MLP_DIM[-1] | |
| grid_size = 21 | |
| xv, yv = torch.meshgrid([ | |
| torch.linspace(-1, 1, grid_size), | |
| torch.linspace(-1, 1, grid_size) | |
| ]) | |
| points_grid = torch.stack([xv.reshape(-1), | |
| yv.reshape(-1)]).unsqueeze(0) | |
| self.register_buffer('points_grid', points_grid) | |
| grid_feat_len = grid_size * grid_size * cfg.MODEL.PyMAF.MLP_DIM[-1] | |
| self.regressor = nn.ModuleList() | |
| for i in range(cfg.MODEL.PyMAF.N_ITER): | |
| if i == 0: | |
| ref_infeat_dim = grid_feat_len | |
| else: | |
| ref_infeat_dim = ma_feat_len | |
| self.regressor.append( | |
| Regressor(feat_dim=ref_infeat_dim, | |
| smpl_mean_params=smpl_mean_params)) | |
| dp_feat_dim = 256 | |
| self.with_uv = cfg.LOSS.POINT_REGRESSION_WEIGHTS > 0 | |
| if cfg.MODEL.PyMAF.AUX_SUPV_ON: | |
| self.dp_head = IUV_predict_layer(feat_dim=dp_feat_dim) | |
| 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 forward(self, x, J_regressor=None): | |
| batch_size = x.shape[0] | |
| # spatial features and global features | |
| s_feat, g_feat = self.feature_extractor(x) | |
| assert cfg.MODEL.PyMAF.N_ITER >= 0 and cfg.MODEL.PyMAF.N_ITER <= 3 | |
| if cfg.MODEL.PyMAF.N_ITER == 1: | |
| deconv_blocks = [self.deconv_layers] | |
| elif cfg.MODEL.PyMAF.N_ITER == 2: | |
| deconv_blocks = [self.deconv_layers[0:6], self.deconv_layers[6:9]] | |
| elif cfg.MODEL.PyMAF.N_ITER == 3: | |
| deconv_blocks = [ | |
| self.deconv_layers[0:3], self.deconv_layers[3:6], | |
| self.deconv_layers[6:9] | |
| ] | |
| out_list = {} | |
| # initial parameters | |
| # TODO: remove the initial mesh generation during forward to reduce runtime | |
| # by generating initial mesh the beforehand: smpl_output = self.init_smpl | |
| smpl_output = self.regressor[0].forward_init(g_feat, | |
| J_regressor=J_regressor) | |
| out_list['smpl_out'] = [smpl_output] | |
| out_list['dp_out'] = [] | |
| # for visulization | |
| vis_feat_list = [s_feat.detach()] | |
| # parameter predictions | |
| for rf_i in range(cfg.MODEL.PyMAF.N_ITER): | |
| pred_cam = smpl_output['pred_cam'] | |
| pred_shape = smpl_output['pred_shape'] | |
| pred_pose = smpl_output['pred_pose'] | |
| pred_cam = pred_cam.detach() | |
| pred_shape = pred_shape.detach() | |
| pred_pose = pred_pose.detach() | |
| s_feat_i = deconv_blocks[rf_i](s_feat) | |
| s_feat = s_feat_i | |
| vis_feat_list.append(s_feat_i.detach()) | |
| self.maf_extractor[rf_i].im_feat = s_feat_i | |
| self.maf_extractor[rf_i].cam = pred_cam | |
| if rf_i == 0: | |
| sample_points = torch.transpose( | |
| self.points_grid.expand(batch_size, -1, -1), 1, 2) | |
| ref_feature = self.maf_extractor[rf_i].sampling(sample_points) | |
| else: | |
| pred_smpl_verts = smpl_output['verts'].detach() | |
| # TODO: use a more sparse SMPL implementation (with 431 vertices) for acceleration | |
| pred_smpl_verts_ds = torch.matmul( | |
| self.maf_extractor[rf_i].Dmap.unsqueeze(0), | |
| pred_smpl_verts) # [B, 431, 3] | |
| ref_feature = self.maf_extractor[rf_i]( | |
| pred_smpl_verts_ds) # [B, 431 * n_feat] | |
| smpl_output = self.regressor[rf_i](ref_feature, | |
| pred_pose, | |
| pred_shape, | |
| pred_cam, | |
| n_iter=1, | |
| J_regressor=J_regressor) | |
| out_list['smpl_out'].append(smpl_output) | |
| if self.training and cfg.MODEL.PyMAF.AUX_SUPV_ON: | |
| iuv_out_dict = self.dp_head(s_feat) | |
| out_list['dp_out'].append(iuv_out_dict) | |
| return out_list | |
| def pymaf_net(smpl_mean_params, pretrained=True): | |
| """ 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) | |
| return model | |