luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

smpl fix

Browse Source
This commit is contained in:
luni
2023-05-22 23:19:39 +02:00
parent b443ff29cd
commit 14e78cf495
1 changed files with 10 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
smpl/smpl_torch.py
+10 -8
View File
@@ -163,7 +163,9 @@ class SMPL(Module):
v_shaped = torch.stack(v_shaped, axis=-1) v_shaped = torch.stack(v_shaped, axis=-1)
# J = torch.matmul(self.J_regressor, v_shaped) # 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) R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3)
if verts is not None: if verts is not None:
@@ -173,7 +175,7 @@ class SMPL(Module):
I_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) + \ 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) 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).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 = []
results.append( results.append(
@@ -204,9 +206,9 @@ 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] 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])]) 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])])
rest_shape_h = torch.cat( 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, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2
) )
@@ -214,9 +216,9 @@ class SMPL(Module):
v = torch.reshape(v, (batch_num, -1, 4))[:, :, :3] v = torch.reshape(v, (batch_num, -1, 4))[:, :, :3]
result = v result = v
# root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3)) root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3))
# result = result - root_joint result = result - root_joint
# joints = joints - root_joint joints = joints - root_joint
# # plut result as a plotly 3d mesh # # plut result as a plotly 3d mesh
# import plotly.graph_objects as go # import plotly.graph_objects as go
@@ -228,4 +230,4 @@ class SMPL(Module):
if remove and verts is None: if remove and verts is None:
result = result[:,self.to_keep,:] result = result[:,self.to_keep,:]
return result, None # joints return result, joints, v_shaped, v_posed