From 05341fd48f07f1ccd653ac006e8d7af4fcd99d66 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Tue, 6 Jun 2023 22:17:23 +0200 Subject: [PATCH] smpl fix --- superSMPL/smpl_torch.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/superSMPL/smpl_torch.py b/superSMPL/smpl_torch.py index 5737dd8..310c1c9 100755 --- a/superSMPL/smpl_torch.py +++ b/superSMPL/smpl_torch.py @@ -167,7 +167,9 @@ class SMPL(Module): 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])]) + 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)) + J = J - J0 R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3) if verts is not None: @@ -177,7 +179,7 @@ class SMPL(Module): 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])]) + 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( @@ -211,7 +213,7 @@ class SMPL(Module): 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] weights = self.weights if remove else self.weights_all - T = torch.stack([torch.tensordot(results[i], weights[g.to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])]) + 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 )