lmd/Pose_to_SMPL_an_230402
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Rename smpl->smplpytorch.

Browse Source
This commit is contained in:
gulvarol
2019-07-16 15:41:45 +02:00
parent 6c67b6de28
commit 6f67ac4199
17 changed files with 35 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
smplpytorch/pytorch/__init__.py Normal file
View File

85
smplpytorch/pytorch/rodrigues_layer.py Normal file
View File

@ -0,0 +1,85 @@
"""
This part reuses code from https://github.com/MandyMo/pytorch_HMR/blob/master/src/util.py
which is part of a PyTorch port of SMPL.
Thanks to Zhang Xiong (MandyMo) for making this great code available on github !
"""
import argparse
from torch.autograd import gradcheck
import torch
from torch.autograd import Variable
def quat2mat(quat):
"""Convert quaternion coefficients to rotation matrix.
Args:
quat: size = [batch_size, 4] 4 <===>(w, x, y, z)
Returns:
Rotation matrix corresponding to the quaternion -- size = [batch_size, 3, 3]
"""
norm_quat = quat
norm_quat = norm_quat / norm_quat.norm(p=2, dim=1, keepdim=True)
w, x, y, z = norm_quat[:, 0], norm_quat[:, 1], norm_quat[:,
2], norm_quat[:,
3]
batch_size = quat.size(0)
w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)
wx, wy, wz = w * x, w * y, w * z
xy, xz, yz = x * y, x * z, y * z
rotMat = torch.stack([
w2 + x2 - y2 - z2, 2 * xy - 2 * wz, 2 * wy + 2 * xz, 2 * wz + 2 * xy,
w2 - x2 + y2 - z2, 2 * yz - 2 * wx, 2 * xz - 2 * wy, 2 * wx + 2 * yz,
w2 - x2 - y2 + z2
],
dim=1).view(batch_size, 3, 3)
return rotMat
def batch_rodrigues(axisang):
#axisang N x 3
axisang_norm = torch.norm(axisang + 1e-8, p=2, dim=1)
angle = torch.unsqueeze(axisang_norm, -1)
axisang_normalized = torch.div(axisang, angle)
angle = angle * 0.5
v_cos = torch.cos(angle)
v_sin = torch.sin(angle)
quat = torch.cat([v_cos, v_sin * axisang_normalized], dim=1)
rot_mat = quat2mat(quat)
rot_mat = rot_mat.view(rot_mat.shape[0], 9)
return rot_mat
def th_get_axis_angle(vector):
angle = torch.norm(vector, 2, 1)
axes = vector / angle.unsqueeze(1)
return axes, angle
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--batch_size', default=1, type=int)
parser.add_argument('--cuda', action='store_true')
args = parser.parse_args()
n_components = 6
rot = 3
inputs = torch.rand(args.batch_size, rot)
inputs_var = Variable(inputs.double(), requires_grad=True)
if args.cuda:
inputs = inputs.cuda()
# outputs = batch_rodrigues(inputs)
test_function = gradcheck(batch_rodrigues, (inputs_var, ))
print('batch test passed !')
inputs = torch.rand(rot)
inputs_var = Variable(inputs.double(), requires_grad=True)
test_function = gradcheck(th_cv2_rod_sub_id.apply, (inputs_var, ))
print('th_cv2_rod test passed')
inputs = torch.rand(rot)
inputs_var = Variable(inputs.double(), requires_grad=True)
test_th = gradcheck(th_cv2_rod.apply, (inputs_var, ))
print('th_cv2_rod_id test passed !')

156
smplpytorch/pytorch/smpl_layer.py Normal file
View File

@ -0,0 +1,156 @@
import os
import numpy as np
import torch
from torch.nn import Module
from smplpytorch.native.webuser.serialization import ready_arguments
from smplpytorch.pytorch import rodrigues_layer
from smplpytorch.pytorch.tensutils import (th_posemap_axisang, th_with_zeros, th_pack, make_list, subtract_flat_id)
class SMPL_Layer(Module):
__constants__ = ['kintree_parents', 'gender', 'center_idx', 'num_joints']
def __init__(self,
center_idx=None,
gender='neutral',
model_root='smpl/native/models'):
"""
Args:
center_idx: index of center joint in our computations,
model_root: path to pkl files for the model
gender: 'neutral' (default) or 'female' or 'male'
"""
super().__init__()
self.center_idx = center_idx
self.gender = gender
if gender == 'neutral':
self.model_path = os.path.join(model_root, 'basicModel_neutral_lbs_10_207_0_v1.0.0.pkl')
elif gender == 'female':
self.model_path = os.path.join(model_root, 'basicModel_f_lbs_10_207_0_v1.0.0.pkl')
elif gender == 'male':
self.model_path = os.path.join(model_root, 'basicModel_m_lbs_10_207_0_v1.0.0.pkl')
smpl_data = ready_arguments(self.model_path)
self.smpl_data = smpl_data
self.register_buffer('th_betas',
torch.Tensor(smpl_data['betas'].r).unsqueeze(0))
self.register_buffer('th_shapedirs',
torch.Tensor(smpl_data['shapedirs'].r))
self.register_buffer('th_posedirs',
torch.Tensor(smpl_data['posedirs'].r))
self.register_buffer(
'th_v_template',
torch.Tensor(smpl_data['v_template'].r).unsqueeze(0))
self.register_buffer(
'th_J_regressor',
torch.Tensor(np.array(smpl_data['J_regressor'].toarray())))
self.register_buffer('th_weights',
torch.Tensor(smpl_data['weights'].r))
self.register_buffer('th_faces',
torch.Tensor(smpl_data['f'].astype(np.int32)).long())
# Kinematic chain params
self.kintree_table = smpl_data['kintree_table']
parents = list(self.kintree_table[0].tolist())
self.kintree_parents = parents
self.num_joints = len(parents) # 24
def forward(self,
th_pose_axisang,
th_betas=torch.zeros(1),
th_trans=torch.zeros(1)):
"""
Args:
th_pose_axisang (Tensor (batch_size x 72)): pose parameters in axis-angle representation
th_betas (Tensor (batch_size x 10)): if provided, uses given shape parameters
th_trans (Tensor (batch_size x 3)): if provided, applies trans to joints and vertices
"""
batch_size = th_pose_axisang.shape[0]
# Convert axis-angle representation to rotation matrix rep.
th_pose_rotmat = th_posemap_axisang(th_pose_axisang)
# Take out the first rotmat (global rotation)
root_rot = th_pose_rotmat[:, :9].view(batch_size, 3, 3)
# Take out the remaining rotmats (23 joints)
th_pose_rotmat = th_pose_rotmat[:, 9:]
th_pose_map = subtract_flat_id(th_pose_rotmat)
# Below does: v_shaped = v_template + shapedirs * betas
# If shape parameters are not provided
if th_betas is None or bool(torch.norm(th_betas) == 0):
th_v_shaped = self.th_v_template + torch.matmul(
self.th_shapedirs, self.th_betas.transpose(1, 0)).permute(2, 0, 1)
th_j = torch.matmul(self.th_J_regressor, th_v_shaped).repeat(
batch_size, 1, 1)
else:
th_v_shaped = self.th_v_template + torch.matmul(
self.th_shapedirs, th_betas.transpose(1, 0)).permute(2, 0, 1)
th_j = torch.matmul(self.th_J_regressor, th_v_shaped)
# Below does: v_posed = v_shaped + posedirs * pose_map
th_v_posed = th_v_shaped + torch.matmul(
self.th_posedirs, th_pose_map.transpose(0, 1)).permute(2, 0, 1)
# Final T pose with transformation done!
# Global rigid transformation
th_results = []
root_j = th_j[:, 0, :].contiguous().view(batch_size, 3, 1)
th_results.append(th_with_zeros(torch.cat([root_rot, root_j], 2)))
# Rotate each part
for i in range(self.num_joints - 1):
i_val = int(i + 1)
joint_rot = th_pose_rotmat[:, (i_val - 1) * 9:i_val *
9].contiguous().view(batch_size, 3, 3)
joint_j = th_j[:, i_val, :].contiguous().view(batch_size, 3, 1)
parent = make_list(self.kintree_parents)[i_val]
parent_j = th_j[:, parent, :].contiguous().view(batch_size, 3, 1)
joint_rel_transform = th_with_zeros(
torch.cat([joint_rot, joint_j - parent_j], 2))
th_results.append(
torch.matmul(th_results[parent], joint_rel_transform))
th_results_global = th_results
th_results2 = torch.zeros((batch_size, 4, 4, self.num_joints),
dtype=root_j.dtype,
device=root_j.device)
for i in range(self.num_joints):
padd_zero = torch.zeros(1, dtype=th_j.dtype, device=th_j.device)
joint_j = torch.cat(
[th_j[:, i],
padd_zero.view(1, 1).repeat(batch_size, 1)], 1)
tmp = torch.bmm(th_results[i], joint_j.unsqueeze(2))
th_results2[:, :, :, i] = th_results[i] - th_pack(tmp)
th_T = torch.matmul(th_results2, self.th_weights.transpose(0, 1))
th_rest_shape_h = torch.cat([
th_v_posed.transpose(2, 1),
torch.ones((batch_size, 1, th_v_posed.shape[1]),
dtype=th_T.dtype,
device=th_T.device),
], 1)
th_verts = (th_T * th_rest_shape_h.unsqueeze(1)).sum(2).transpose(2, 1)
th_verts = th_verts[:, :, :3]
th_jtr = torch.stack(th_results_global, dim=1)[:, :, :3, 3]
# If translation is not provided
if th_trans is None or bool(torch.norm(th_trans) == 0):
if self.center_idx is not None:
center_joint = th_jtr[:, self.center_idx].unsqueeze(1)
th_jtr = th_jtr - center_joint
th_verts = th_verts - center_joint
else:
th_jtr = th_jtr + th_trans.unsqueeze(1)
th_verts = th_verts + th_trans.unsqueeze(1)
# Vertices and joints in meters
return th_verts, th_jtr

53
smplpytorch/pytorch/tensutils.py Normal file
View File

@ -0,0 +1,53 @@
import torch
from smplpytorch.pytorch import rodrigues_layer
def th_posemap_axisang(pose_vectors):
'''
Converts axis-angle to rotmat
pose_vectors (Tensor (batch_size x 72)): pose parameters in axis-angle representation
'''
rot_nb = int(pose_vectors.shape[1] / 3)
rot_mats = []
for joint_idx in range(rot_nb):
axis_ang = pose_vectors[:, joint_idx * 3:(joint_idx + 1) * 3]
rot_mat = rodrigues_layer.batch_rodrigues(axis_ang)
rot_mats.append(rot_mat)
rot_mats = torch.cat(rot_mats, 1)
return rot_mats
def th_with_zeros(tensor):
batch_size = tensor.shape[0]
padding = tensor.new([0.0, 0.0, 0.0, 1.0])
padding.requires_grad = False
concat_list = [tensor, padding.view(1, 1, 4).repeat(batch_size, 1, 1)]
cat_res = torch.cat(concat_list, 1)
return cat_res
def th_pack(tensor):
batch_size = tensor.shape[0]
padding = tensor.new_zeros((batch_size, 4, 3))
padding.requires_grad = False
pack_list = [padding, tensor]
pack_res = torch.cat(pack_list, 2)
return pack_res
def subtract_flat_id(rot_mats):
# Subtracts identity as a flattened tensor
id_flat = torch.eye(
3, dtype=rot_mats.dtype, device=rot_mats.device).view(1, 9).repeat(
rot_mats.shape[0], 23)
# id_flat.requires_grad = False
results = rot_mats - id_flat
return results
def make_list(tensor):
# type: (List[int]) -> List[int]
return tensor