diff --git a/smpl/smpl_torch.py b/smpl/smpl_torch.py index 1b90f78..b18dfb7 100644 --- a/smpl/smpl_torch.py +++ b/smpl/smpl_torch.py @@ -10,7 +10,7 @@ from time import time from torch.nn import Module class SMPL(Module): - def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False): + def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False, batch_size=4): super(SMPL, self).__init__() self.J_regressor = [] @@ -56,6 +56,14 @@ class SMPL(Module): self.shape = (1,27578,3) self.device = device + self.batch_size = batch_size + + self.with_zeros_ones = torch.tensor( + [[[0.0, 0.0, 0.0, 1.0]]], dtype=torch.float32 + ).expand(self.batch_size,-1,-1).to(self.device) + self.results_zeros = torch.zeros((self.batch_size, 24, 1), dtype=torch.float32).to(self.device) + self.J_ones = torch.ones((self.batch_size, 24, 1), dtype=torch.float32, device=self.device) + self.rest_shape_h_ones_full = torch.ones((self.batch_size, self.weights_all.shape[1], 1), dtype=torch.float32).to(self.device) if orig_smpl: self.to_keep = self.to_keep[self.to_keep < 6890] @@ -67,6 +75,8 @@ class SMPL(Module): self.weights = self.weights_all[:,self.to_keep,:] self.shape = (1, 6890, 3) + self.rest_shape_h_ones = torch.ones((self.batch_size, self.to_keep.shape[0], 1), dtype=torch.float32).to(self.device) + @staticmethod def rodrigues(r): """ @@ -91,7 +101,7 @@ class SMPL(Module): 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)) + m = m.view((-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) @@ -112,8 +122,8 @@ class SMPL(Module): 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) + [[[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 @@ -158,8 +168,8 @@ class SMPL(Module): # 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.append(torch.tensordot(betas[:,:10], self.shapedirs[0], dims=[[1], [2]]) + self.v_template[0].view(self.shape)) # female + v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[1], dims=[[1], [2]]) + self.v_template[1].view(self.shape)) # male v_shaped = torch.stack(v_shaped, axis=-1) # J = torch.matmul(self.J_regressor, v_shaped) @@ -167,7 +177,7 @@ class SMPL(Module): 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])]) else: J = torch.stack([torch.matmul(self.J_regressor[torch.greater(g, 0.5).to(torch.long)], v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)]) for i, g in enumerate(betas[:,-1])]) - J0 = torch.reshape(J[:,0,:], (-1, 1, 3)) + J0 = J[:,0,:].view((-1, 1, 3)) J = J - J0 R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3) @@ -177,12 +187,12 @@ class SMPL(Module): 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) + lrotmin = (R_cube - I_cube).view(batch_num, -1, 1).squeeze(dim=2) v_posed = torch.stack([v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)] + torch.tensordot(lrotmin[i], self.posedirs[torch.greater(g, 0.5).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)) + self.with_zeros(torch.cat((R_cube_big[:, 0], J[:, 0, :].view((-1, 3, 1))), dim=2)) ) for i in range(1, 24): results.append( @@ -190,7 +200,7 @@ class SMPL(Module): results[self.parent[i]], self.with_zeros( torch.cat( - (R_cube_big[:, i], torch.reshape(J[:, i, :] - J[:, self.parent[i], :], (-1, 3, 1))), + (R_cube_big[:, i], (J[:, i, :] - J[:, self.parent[i], :]).view((-1, 3, 1))), dim=2 ) ) @@ -202,25 +212,25 @@ class SMPL(Module): self.pack( torch.matmul( stacked, - torch.reshape( - torch.cat((J, torch.zeros((batch_num, 24, 1), dtype=torch.float32).to(self.device)), dim=2), + torch.cat((J, self.results_zeros), dim=2).view( (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] + joints = torch.matmul(results, torch.transpose(torch.concat([J, self.J_ones], axis=-1)[:,:,None,:],2,3))[:,:,:3,0] weights = self.weights if verts is not None else self.weights_all + rest_shape_h_ones_local = self.rest_shape_h_ones if verts is not None else self.rest_shape_h_ones_full T = torch.stack([torch.tensordot(results[i], weights[torch.greater(g, 0.5).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_posed, rest_shape_h_ones_local), 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] + v = torch.matmul(T, rest_shape_h.view((batch_num, -1, 4, 1))) + v = v.view((batch_num, -1, 4))[:, :, :3] result = v - root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3)) + root_joint = joints[:,0,:].view((-1, 1, 3)) result = result - root_joint joints = joints - root_joint