# based on https://github.com/CalciferZh/SMPL/blob/master/smpl_torch_batch.py import os import sys import pickle import torch import numpy as np import scipy.io as sio from time import time from torch.nn import Module from utils import plot_mesh class SMPL(Module): def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False): super(SMPL, self).__init__() self.J_regressor = [] self.v_template = [] self.shapedirs = [] self.posedirs = [] self.weights = [] self.is_skirt = is_skirt # Load SMPL models # path = os.path.abspath(os.path.dirname(__file__)) with open(os.path.join(path, 'super_model_f.pkl'), 'rb') as f: female = pickle.load(f, encoding="latin1") self.J_regressor.append(torch.tensor(np.array(female['J_regressor'].todense(), dtype=np.float32))) self.v_template.append(torch.tensor(female['v_template'], dtype=torch.float32)) self.posedirs.append(torch.tensor(female['posedirs'], dtype=torch.float32)) self.shapedirs.append(torch.tensor(female['shapedirs'], dtype=torch.float32)) self.weights.append(torch.tensor(female['weights'], dtype=torch.float32)) with open(os.path.join(path, 'super_model_m.pkl'), 'rb') as f: male = pickle.load(f, encoding="latin1") self.J_regressor.append(torch.tensor(np.array(male['J_regressor'].todense(), dtype=np.float32))) self.v_template.append(torch.tensor(male['v_template'], dtype=torch.float32)) self.posedirs.append(torch.tensor(male['posedirs'], dtype=torch.float32)) self.shapedirs.append(torch.tensor(male['shapedirs'], dtype=torch.float32)) self.weights.append(torch.tensor(male['weights'], dtype=torch.float32)) kintree_table = male['kintree_table'] id_to_col = {kintree_table[1, i]: i for i in range(kintree_table.shape[1])} self.parent = {i: id_to_col[kintree_table[0, i]] for i in range(1, kintree_table.shape[1])} self.J_regressor = torch.stack(self.J_regressor, axis=0).to(device) self.v_template = torch.stack(self.v_template, axis=0).to(device) self.shapedirs = torch.stack(self.shapedirs, axis=0).to(device) self.posedirs = torch.stack(self.posedirs, axis=0).to(device) self.weights = torch.stack(self.weights, axis=0).to(device) # Topology no head/hands/feet F = sio.loadmat(os.path.join(path, 'ssmpl_faces.mat')) self.f = np.array(F['faces_default']) self.f_skirt = np.array(F['faces_skirt']) self.nn = np.array(sio.loadmat(os.path.join(path, 'NN.mat'))) self.to_keep = sio.loadmat(os.path.join(path, 'keep_remove.mat'))['keep'].reshape(-1) self.weights_all = self.weights # self.weights = self.weights[:,self.to_keep,:] self.shape = (1,27578,3) self.device = device if orig_smpl: self.to_keep = self.to_keep[self.to_keep < 6890] self.J_regressor = self.J_regressor[:,:,:6890] self.v_template = self.v_template[:,:6890,:] self.shapedirs = self.shapedirs[:,:6890,:,:] self.posedirs = self.posedirs[:,:6890,:,:] self.weights_all = self.weights_all[:,:6890,:] self.weights = self.weights_all[:,self.to_keep,:] self.shape = (1, 6890, 3) @staticmethod def rodrigues(r): """ Rodrigues' rotation formula that turns axis-angle tensor into rotation matrix in a batch-ed manner. Parameter: ---------- r: Axis-angle rotation tensor of shape [batch_size * angle_num, 1, 3]. Return: ------- Rotation matrix of shape [batch_size * angle_num, 3, 3]. """ # eps = r.clone().normal_(std=1e-8) # theta = torch.norm(r + eps, dim=(1, 2), keepdim=True) # dim cannot be tuple theta = torch.norm(r, dim=(1, 2), keepdim=True) theta = torch.maximum(theta, torch.ones_like(theta) * np.finfo(np.float32).eps) theta_dim = theta.shape[0] r_hat = r / theta # cos = torch.cos(theta) z_stick = torch.zeros(theta_dim, dtype=torch.float32).to(r.device) m = torch.stack( (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick, -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), dim=1) m = torch.reshape(m, (-1, 3, 3)) i_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) \ + torch.zeros((theta_dim, 3, 3), dtype=torch.float32)).to(r.device) A = r_hat.permute(0, 2, 1) # dot = torch.matmul(A, r_hat) # R = cos * i_cube + (1 - cos) * dot + torch.sin(theta) * m R = i_cube + m * torch.sin(theta) + torch.matmul(m, m) * (1 - torch.cos(theta)) return R @staticmethod def with_zeros(x): """ Append a [0, 0, 0, 1] tensor to a [3, 4] tensor. Parameter: --------- x: Tensor to be appended. Return: ------ Tensor after appending of shape [4,4] """ ones = torch.tensor( [[[0.0, 0.0, 0.0, 1.0]]], dtype=torch.float32 ).expand(x.shape[0],-1,-1).to(x.device) ret = torch.cat((x, ones), dim=1) return ret @staticmethod def pack(x): """ Append zero tensors of shape [4, 3] to a batch of [4, 1] shape tensor. Parameter: ---------- x: A tensor of shape [batch_size, 4, 1] Return: ------ A tensor of shape [batch_size, 4, 4] after appending. """ zeros43 = torch.zeros( (x.shape[0], x.shape[1], 4, 3), dtype=torch.float32).to(x.device) ret = torch.cat((zeros43, x), dim=3) return ret def forward(self, betas, pose, remove=True): """ Construct a compute graph that takes in parameters and outputs a tensor as model vertices. Face indices are also returned as a numpy ndarray. 20190128: Add batch support. Parameters: --------- pose: Also known as 'theta', an [N, 24, 3] tensor indicating child joint rotation relative to parent joint. For root joint it's global orientation. Represented in a axis-angle format. betas: Parameter for model shape. A tensor of shape [N, 10] as coefficients of PCA components. Only 10 components were released by SMPL author. trans: Global translation tensor of shape [N, 3]. Return: ------ A 3-D tensor of [N * 6890 * 3] for vertices, and the corresponding [N * 19 * 3] joint positions. """ batch_num = betas.shape[0] # v_shaped = torch.tensordot(betas, self.shapedirs, dims=([1], [2])) + self.v_template v_shaped = [] v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[0], dims=[[1], [2]]) + torch.reshape(self.v_template[0], shape=self.shape)) # female v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[1], dims=[[1], [2]]) + torch.reshape(self.v_template[1], shape=self.shape)) # male v_shaped = torch.stack(v_shaped, axis=-1) # J = torch.matmul(self.J_regressor, v_shaped) J = torch.stack([torch.matmul(self.J_regressor[g.to(torch.long)], v_shaped[i, :, :, g.to(torch.long)]) for i, g in enumerate(betas[:,-1])]) R_cube_big = self.rodrigues(pose.view(-1, 1, 3)).reshape(batch_num, -1, 3, 3) R_cube = R_cube_big[:, 1:, :, :] I_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) + \ torch.zeros((batch_num, R_cube.shape[1], 3, 3), dtype=torch.float32)).to(self.device) lrotmin = (R_cube - I_cube).reshape(batch_num, -1, 1).squeeze(dim=2) v_posed = torch.stack([v_shaped[i, :, :, g.to(torch.long)] + torch.tensordot(lrotmin[i], self.posedirs[g.to(torch.long)], dims=([0], [2])) for i, g in enumerate(betas[:,-1])]) results = [] results.append( self.with_zeros(torch.cat((R_cube_big[:, 0], torch.reshape(J[:, 0, :], (-1, 3, 1))), dim=2)) ) for i in range(1, 24): results.append( torch.matmul( results[self.parent[i]], self.with_zeros( torch.cat( (R_cube_big[:, i], torch.reshape(J[:, i, :] - J[:, self.parent[i], :], (-1, 3, 1))), dim=2 ) ) ) ) stacked = torch.stack(results, dim=1) results = stacked - \ self.pack( torch.matmul( stacked, torch.reshape( torch.cat((J, torch.zeros((batch_num, 24, 1), dtype=torch.float32).to(self.device)), dim=2), (batch_num, 24, 4, 1) ) ) ) joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0] T = torch.stack([torch.tensordot(results[i], self.weights[g.to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])]) rest_shape_h = torch.cat( (v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2 ) v = torch.matmul(T, torch.reshape(rest_shape_h, (batch_num, -1, 4, 1))) v = torch.reshape(v, (batch_num, -1, 4))[:, :, :3] result = v root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3)) result = result - root_joint joints = joints - root_joint if remove: result = result[:,self.to_keep,:] return result, joints