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