File size: 5,103 Bytes
d8dc31c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | # -*- coding: utf-8 -*-
#
# Copyright (C) 2019 Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG),
# acting on behalf of its Max Planck Institute for Intelligent Systems and the
# Max Planck Institute for Biological Cybernetics. All rights reserved.
#
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is holder of all proprietary rights
# on this computer program. You can only use this computer program if you have closed a license agreement
# with MPG or you get the right to use the computer program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and liable to prosecution.
# Contact: ps-license@tuebingen.mpg.de
#
#
# If you use this code in a research publication please consider citing the following:
#
# Expressive Body Capture: 3D Hands, Face, and Body from a Single Image <https://arxiv.org/abs/1904.05866>
#
#
# Code Developed by:
# Nima Ghorbani <https://nghorbani.github.io/>
#
# 2020.12.12
import numpy as np
import torch
# from human_body_prior.models.model_components import BatchFlatten
# from human_body_prior.tools.rotation_tools import matrot2aa
from nets import tgm_conversion as tgm
from torch import nn
from torch.nn import functional as F
class BatchFlatten(nn.Module):
def __init__(self):
super(BatchFlatten, self).__init__()
self._name = 'batch_flatten'
def forward(self, x):
return x.view(x.shape[0], -1)
class ContinousRotReprDecoder(nn.Module):
def __init__(self):
super(ContinousRotReprDecoder, self).__init__()
def forward(self, module_input):
reshaped_input = module_input.view(-1, 3, 2)
b1 = F.normalize(reshaped_input[:, :, 0], dim=1)
dot_prod = torch.sum(b1 * reshaped_input[:, :, 1], dim=1, keepdim=True)
b2 = F.normalize(reshaped_input[:, :, 1] - dot_prod * b1, dim=-1)
b3 = torch.cross(b1, b2, dim=1)
return torch.stack([b1, b2, b3], dim=-1)
class NormalDistDecoder(nn.Module):
def __init__(self, num_feat_in, latentD):
super(NormalDistDecoder, self).__init__()
self.mu = nn.Linear(num_feat_in, latentD)
self.logvar = nn.Linear(num_feat_in, latentD)
def forward(self, Xout):
return torch.distributions.normal.Normal(self.mu(Xout), F.softplus(self.logvar(Xout)))
class VPoser(nn.Module):
def __init__(self, model_ps):
super(VPoser, self).__init__()
num_neurons, self.latentD = model_ps.model_params.num_neurons, model_ps.model_params.latentD
self.num_joints = 21
n_features = self.num_joints * 3
self.encoder_net = nn.Sequential(
BatchFlatten(),
nn.BatchNorm1d(n_features),
nn.Linear(n_features, num_neurons),
nn.LeakyReLU(),
nn.BatchNorm1d(num_neurons),
nn.Dropout(0.1),
nn.Linear(num_neurons, num_neurons),
nn.Linear(num_neurons, num_neurons),
NormalDistDecoder(num_neurons, self.latentD)
)
self.decoder_net = nn.Sequential(
nn.Linear(self.latentD, num_neurons),
nn.LeakyReLU(),
nn.Dropout(0.1),
nn.Linear(num_neurons, num_neurons),
nn.LeakyReLU(),
nn.Linear(num_neurons, self.num_joints * 6),
ContinousRotReprDecoder(),
)
def encode(self, pose_body):
'''
:param Pin: Nx(numjoints*3)
:param rep_type: 'matrot'/'aa' for matrix rotations or axis-angle
:return:
'''
return self.encoder_net(pose_body)
def decode(self, Zin):
bs = Zin.shape[0]
prec = self.decoder_net(Zin)
return {
'pose_body': self.matrot2aa(prec.view(-1, 3, 3)).view(bs, -1, 3),
'pose_body_matrot': prec.view(bs, -1, 9)
}
def matrot2aa(self, pose_matrot):
'''
:param pose_matrot: Nx3x3
:return: Nx3
'''
bs = pose_matrot.size(0)
homogen_matrot = F.pad(pose_matrot, [0, 1])
pose = tgm.rotation_matrix_to_angle_axis(homogen_matrot)
return pose
def forward(self, pose_body):
'''
:param Pin: aa: Nx1xnum_jointsx3 / matrot: Nx1xnum_jointsx9
:param input_type: matrot / aa for matrix rotations or axis angles
:param output_type: matrot / aa
:return:
'''
q_z = self.encode(pose_body)
q_z_sample = q_z.rsample()
decode_results = self.decode(q_z_sample)
decode_results.update({'poZ_body_mean': q_z.mean, 'poZ_body_std': q_z.scale, 'q_z': q_z})
return q_z_sample, decode_results
def sample_poses(self, num_poses, seed=None):
np.random.seed(seed)
some_weight = [a for a in self.parameters()][0]
dtype = some_weight.dtype
device = some_weight.device
self.eval()
with torch.no_grad():
Zgen = torch.tensor(np.random.normal(0., 1., size=(num_poses, self.latentD)), dtype=dtype, device=device)
return self.decode(Zgen)
|