optimize smpl

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
+        [[[0.0, 0.0, 0.0, 1.0]]], dtype=torch.float32
-    ).expand(x.shape[0],-1,-1).to(x.device)
+      ).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[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]]) + torch.reshape(self.v_template[1], shape=self.shape)) # male
+    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, self.results_zeros), dim=2).view(
-            torch.cat((J, torch.zeros((batch_num, 24, 1), dtype=torch.float32).to(self.device)), dim=2),
             (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.matmul(T, rest_shape_h.view((batch_num, -1, 4, 1)))
-    v = torch.reshape(v, (batch_num, -1, 4))[:, :, :3]
+    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