final smpl fix

This commit is contained in:
2023-06-30 00:13:14 +02:00
parent ff9771a415
commit 837a1449b8
+8 -4
View File
@@ -163,9 +163,12 @@ 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[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
if verts is None:
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))
J = J - J0
R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3)
if verts is not None:
@@ -208,7 +211,8 @@ 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]
T = torch.stack([torch.tensordot(results[i], self.weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
weights = self.weights if verts is not None else self.weights_all
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
)