smpl update

smpl/smpl_torch.py

@@ -8,7 +8,6 @@ import numpy as np
 import scipy.io as sio
 from time import time
 from torch.nn import Module
-from utils import plot_mesh
 
 class SMPL(Module):
   def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False):
@@ -53,7 +52,7 @@ class SMPL(Module):
     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.weights = self.weights[:,self.to_keep,:]
     
     self.shape = (1,27578,3)
     self.device = device
@@ -135,7 +134,7 @@ class SMPL(Module):
     return ret
 
 
-  def forward(self, betas, pose, remove=True):
+  def forward(self, pose, betas, verts: torch.Tensor = None, remove: bool = True):
     
     """
           Construct a compute graph that takes in parameters and outputs a tensor as
@@ -165,13 +164,16 @@ class SMPL(Module):
     
     # 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_big = self.rodrigues(pose.reshape(-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])])
+    if verts is not None:
+      v_posed = verts
+    else:
+      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(