From 79a09e0f6241658e03cd6850c424715ac0aaa151 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 3 Feb 2023 16:29:06 +0000 Subject: [PATCH] implement pytorch version --- smpl_torch.py | 261 ++++++++++++++++++++++++++++++++++++++++++++++++++ utils.py | 9 +- 2 files changed, 267 insertions(+), 3 deletions(-) create mode 100644 smpl_torch.py diff --git a/smpl_torch.py b/smpl_torch.py new file mode 100644 index 0000000..e7fc7b9 --- /dev/null +++ b/smpl_torch.py @@ -0,0 +1,261 @@ +# 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 SMPLModel(Module): + def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False): + + super(SMPLModel, 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_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 + 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 write_obj(self, verts, file_name): + with open(file_name, 'w') as fp: + for v in verts: + fp.write('v %f %f %f\n' % (v[0], v[1], v[2])) + + for f in self.faces + 1: + fp.write('f %d %d %d\n' % (f[0], f[1], f[2])) + + def forward(self, betas, pose, trans, 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) + ) + ) + ) + # Restart from here + 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 + torch.reshape(trans, (batch_num, 1, 3)) + # estimate 3D joint locations + # print(result.shape) + # print(self.joint_regressor.shape) + joints = torch.stack([torch.tensordot(result[i], self.J_regressor[g.to(torch.long)], dims=([0], [1])) for i, g in enumerate(betas[:,-1])]).transpose(1, 2) + + 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 + + +def test_gpu(gpu_id=[0], path=''): + if len(gpu_id) > 0 and torch.cuda.is_available(): + os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id[0]) + device = torch.device('cuda') + else: + device = torch.device('cpu') + #print(device) + + pose_size = 72 + beta_size = 10 + + np.random.seed(9608) + model = SMPLModel(device=device, path=path) + for i in range(5): + pose = torch.from_numpy((np.random.rand(4, pose_size) - 0.5) * 0.4)\ + .type(torch.float32).to(device) + betas = torch.from_numpy((np.random.rand(4, beta_size) - 0.5) * 0.06) \ + .type(torch.float32).to(device) + betas = torch.concat((betas, torch.zeros((4,1), device=device)), axis=1) + s = time() + trans = torch.from_numpy(np.zeros((4, 3))).type(torch.float32).to(device) + result, joints = model(betas, pose, trans) + print(time() - s) + plot_mesh(result[0].cpu().numpy(), model.f).show() + + # outmesh_path = './smpl_torch_{}.obj' + # for i in range(result.shape[0]): + # model.write_obj(result[i], outmesh_path.format(i)) + + +if __name__ == '__main__': + test_gpu([5], path='model_files') \ No newline at end of file diff --git a/utils.py b/utils.py index 7b47685..08346c7 100644 --- a/utils.py +++ b/utils.py @@ -15,9 +15,12 @@ def plot_mesh(verts, faces): ),) fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1)) - helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False), - yaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False), - zaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),) + # helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False), + # yaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False), + # zaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),) + helper = dict(xaxis = dict(range=[-1,1], ), + yaxis = dict(range=[-1,1], ), + zaxis = dict(range=[-1,1], ) ) fig.update_layout(scene = helper) fig.update_layout(margin=dict(l=0, r=0, b=0, t=0, pad=0))