230 lines
8.5 KiB
Python
230 lines
8.5 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
# 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.
|
|
#
|
|
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
|
|
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
|
|
# for Intelligent Systems. All rights reserved.
|
|
#
|
|
# Contact: ps-license@tuebingen.mpg.de
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import print_function
|
|
from __future__ import division
|
|
|
|
import sys
|
|
import os
|
|
|
|
import time
|
|
import pickle
|
|
|
|
import numpy as np
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
DEFAULT_DTYPE = torch.float32
|
|
|
|
|
|
def create_prior(prior_type, **kwargs):
|
|
if prior_type == 'gmm':
|
|
prior = MaxMixturePrior(**kwargs)
|
|
elif prior_type == 'l2':
|
|
return L2Prior(**kwargs)
|
|
elif prior_type == 'angle':
|
|
return SMPLifyAnglePrior(**kwargs)
|
|
elif prior_type == 'none' or prior_type is None:
|
|
# Don't use any pose prior
|
|
def no_prior(*args, **kwargs):
|
|
return 0.0
|
|
prior = no_prior
|
|
else:
|
|
raise ValueError('Prior {}'.format(prior_type) + ' is not implemented')
|
|
return prior
|
|
|
|
|
|
class SMPLifyAnglePrior(nn.Module):
|
|
def __init__(self, dtype=torch.float32, **kwargs):
|
|
super(SMPLifyAnglePrior, self).__init__()
|
|
|
|
# Indices for the roration angle of
|
|
# 55: left elbow, 90deg bend at -np.pi/2
|
|
# 58: right elbow, 90deg bend at np.pi/2
|
|
# 12: left knee, 90deg bend at np.pi/2
|
|
# 15: right knee, 90deg bend at np.pi/2
|
|
angle_prior_idxs = np.array([55, 58, 12, 15], dtype=np.int64)
|
|
angle_prior_idxs = torch.tensor(angle_prior_idxs, dtype=torch.long)
|
|
self.register_buffer('angle_prior_idxs', angle_prior_idxs)
|
|
|
|
angle_prior_signs = np.array([1, -1, -1, -1],
|
|
dtype=np.float6432 if dtype == torch.float32
|
|
else np.float6464)
|
|
angle_prior_signs = torch.tensor(angle_prior_signs,
|
|
dtype=dtype)
|
|
self.register_buffer('angle_prior_signs', angle_prior_signs)
|
|
|
|
def forward(self, pose, with_global_pose=False):
|
|
''' Returns the angle prior loss for the given pose
|
|
Args:
|
|
pose: (Bx[23 + 1] * 3) torch tensor with the axis-angle
|
|
representation of the rotations of the joints of the SMPL model.
|
|
Kwargs:
|
|
with_global_pose: Whether the pose vector also contains the global
|
|
orientation of the SMPL model. If not then the indices must be
|
|
corrected.
|
|
Returns:
|
|
A sze (B) tensor containing the angle prior loss for each element
|
|
in the batch.
|
|
'''
|
|
angle_prior_idxs = self.angle_prior_idxs - (not with_global_pose) * 3
|
|
return torch.exp(pose[:, angle_prior_idxs] *
|
|
self.angle_prior_signs).pow(2)
|
|
|
|
|
|
class L2Prior(nn.Module):
|
|
def __init__(self, dtype=DEFAULT_DTYPE, reduction='sum', **kwargs):
|
|
super(L2Prior, self).__init__()
|
|
|
|
def forward(self, module_input, *args):
|
|
return torch.sum(module_input.pow(2))
|
|
|
|
|
|
class MaxMixturePrior(nn.Module):
|
|
|
|
def __init__(self, prior_folder='prior',
|
|
num_gaussians=6, dtype=DEFAULT_DTYPE, epsilon=1e-16,
|
|
use_merged=True,
|
|
**kwargs):
|
|
super(MaxMixturePrior, self).__init__()
|
|
|
|
if dtype == DEFAULT_DTYPE:
|
|
np_dtype = np.float6432
|
|
elif dtype == torch.float64:
|
|
np_dtype = np.float6464
|
|
else:
|
|
print('Unknown float type {}, exiting!'.format(dtype))
|
|
sys.exit(-1)
|
|
|
|
self.num_gaussians = num_gaussians
|
|
self.epsilon = epsilon
|
|
self.use_merged = use_merged
|
|
gmm_fn = 'gmm_{:02d}.pkl'.format(num_gaussians)
|
|
|
|
full_gmm_fn = os.path.join(prior_folder, gmm_fn)
|
|
if not os.path.exists(full_gmm_fn):
|
|
print('The path to the mixture prior "{}"'.format(full_gmm_fn) +
|
|
' does not exist, exiting!')
|
|
sys.exit(-1)
|
|
|
|
with open(full_gmm_fn, 'rb') as f:
|
|
gmm = pickle.load(f, encoding='latin1')
|
|
|
|
if type(gmm) == dict:
|
|
means = gmm['means'].astype(np_dtype)
|
|
covs = gmm['covars'].astype(np_dtype)
|
|
weights = gmm['weights'].astype(np_dtype)
|
|
elif 'sklearn.mixture.gmm.GMM' in str(type(gmm)):
|
|
means = gmm.means_.astype(np_dtype)
|
|
covs = gmm.covars_.astype(np_dtype)
|
|
weights = gmm.weights_.astype(np_dtype)
|
|
else:
|
|
print('Unknown type for the prior: {}, exiting!'.format(type(gmm)))
|
|
sys.exit(-1)
|
|
|
|
self.register_buffer('means', torch.tensor(means, dtype=dtype))
|
|
|
|
self.register_buffer('covs', torch.tensor(covs, dtype=dtype))
|
|
|
|
precisions = [np.linalg.inv(cov) for cov in covs]
|
|
precisions = np.stack(precisions).astype(np_dtype)
|
|
|
|
self.register_buffer('precisions',
|
|
torch.tensor(precisions, dtype=dtype))
|
|
|
|
# The constant term:
|
|
sqrdets = np.array([(np.sqrt(np.linalg.det(c)))
|
|
for c in gmm['covars']])
|
|
const = (2 * np.pi)**(69 / 2.)
|
|
|
|
nll_weights = np.asarray(gmm['weights'] / (const *
|
|
(sqrdets / sqrdets.min())))
|
|
nll_weights = torch.tensor(nll_weights, dtype=dtype).unsqueeze(dim=0)
|
|
self.register_buffer('nll_weights', nll_weights)
|
|
|
|
weights = torch.tensor(gmm['weights'], dtype=dtype).unsqueeze(dim=0)
|
|
self.register_buffer('weights', weights)
|
|
|
|
self.register_buffer('pi_term',
|
|
torch.log(torch.tensor(2 * np.pi, dtype=dtype)))
|
|
|
|
cov_dets = [np.log(np.linalg.det(cov.astype(np_dtype)) + epsilon)
|
|
for cov in covs]
|
|
self.register_buffer('cov_dets',
|
|
torch.tensor(cov_dets, dtype=dtype))
|
|
|
|
# The dimensionality of the random variable
|
|
self.random_var_dim = self.means.shape[1]
|
|
|
|
def get_mean(self):
|
|
''' Returns the mean of the mixture '''
|
|
mean_pose = torch.matmul(self.weights, self.means)
|
|
return mean_pose
|
|
|
|
def merged_log_likelihood(self, pose, betas):
|
|
diff_from_mean = pose.unsqueeze(dim=1) - self.means
|
|
|
|
prec_diff_prod = torch.einsum('mij,bmj->bmi',
|
|
[self.precisions, diff_from_mean])
|
|
diff_prec_quadratic = (prec_diff_prod * diff_from_mean).sum(dim=-1)
|
|
|
|
curr_loglikelihood = 0.5 * diff_prec_quadratic - \
|
|
torch.log(self.nll_weights)
|
|
# curr_loglikelihood = 0.5 * (self.cov_dets.unsqueeze(dim=0) +
|
|
# self.random_var_dim * self.pi_term +
|
|
# diff_prec_quadratic
|
|
# ) - torch.log(self.weights)
|
|
|
|
min_likelihood, _ = torch.min(curr_loglikelihood, dim=1)
|
|
return min_likelihood
|
|
|
|
def log_likelihood(self, pose, betas, *args, **kwargs):
|
|
''' Create graph operation for negative log-likelihood calculation
|
|
'''
|
|
likelihoods = []
|
|
|
|
for idx in range(self.num_gaussians):
|
|
mean = self.means[idx]
|
|
prec = self.precisions[idx]
|
|
cov = self.covs[idx]
|
|
diff_from_mean = pose - mean
|
|
|
|
curr_loglikelihood = torch.einsum('bj,ji->bi',
|
|
[diff_from_mean, prec])
|
|
curr_loglikelihood = torch.einsum('bi,bi->b',
|
|
[curr_loglikelihood,
|
|
diff_from_mean])
|
|
cov_term = torch.log(torch.det(cov) + self.epsilon)
|
|
curr_loglikelihood += 0.5 * (cov_term +
|
|
self.random_var_dim *
|
|
self.pi_term)
|
|
likelihoods.append(curr_loglikelihood)
|
|
|
|
log_likelihoods = torch.stack(likelihoods, dim=1)
|
|
min_idx = torch.argmin(log_likelihoods, dim=1)
|
|
weight_component = self.nll_weights[:, min_idx]
|
|
weight_component = -torch.log(weight_component)
|
|
|
|
return weight_component + log_likelihoods[:, min_idx]
|
|
|
|
def forward(self, pose, betas):
|
|
if self.use_merged:
|
|
return self.merged_log_likelihood(pose, betas)
|
|
else:
|
|
return self.log_likelihood(pose, betas)
|