finalize

Dockerfile

@@ -1,5 +1,5 @@
 # Use an existing image as the base
-FROM pytorch/pytorch:1.12.0-cuda11.3-cudnn8-devel
+FROM nvcr.io/nvidia/pytorch:23.03-py3
 
 # Add the current user to the image
 ARG USER_NAME=lhunor
@@ -8,7 +8,7 @@ ARG GROUP_NAME=guest
 ARG GROUP_ID=11000
 
 # Update the system and install required packages
-RUN apt-get update && apt-get upgrade -y && apt-get install -y sudo git
+RUN apt-get update && apt-get upgrade -y && apt-get install -y sudo git apt-utils
 
 # Create a new group with the same GID as the host
 RUN groupadd --gid $GROUP_ID $GROUP_NAME
@@ -26,8 +26,10 @@ RUN usermod -aG sudo $USER_NAME
 # Set the default user to the new user
 USER $USER_NAME
 
-WORKDIR /home/$USER_NAME
+WORKDIR /workspace
 
+ENV PATH="$PATH:/home/$USER_NAME/.local/bin"
+RUN pip install --upgrade pip
 COPY requirements.txt /tmp/
 RUN pip install --requirement /tmp/requirements.txt
 

docker-compose.yml

@@ -0,0 +1,21 @@
+services:
+  lhunor_torch:
+    image: lhunor_smpl:v2
+    container_name: lhunor_smpl
+    #command: python -c "import tensorflow as tf;tf.test.gpu_device_name()"
+    ipc: host
+    tty: true
+    ulimits:
+      memlock: -1
+      stack: 67108864
+    volumes:
+      - '~/workspace:/workspace'
+      - '/data-net/guest/lhunor/data:/data-net'
+      - '/data-fast/126-data1/lhunor/data:/data-local/126'
+      - '/data-fast/127-data1/lhunor/data:/data-local/127'
+      - '/data-fast/128-data1/lhunor/data:/data-local/128'
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu]

smpl_tf.py

@@ -141,6 +141,7 @@ class SMPL:
     def pack(self, x):
         return tf.concat((tf.zeros((x.get_shape().as_list()[0], 4, 3), dtype=tf.float32), x), axis=2)
         
+    # @tf.function
     def run_rodrigues(self, pose):
         batch_size = pose.shape[0]
         R_out = self.rodrigues(tf.reshape(pose, (-1, 1, 3)))
@@ -151,6 +152,7 @@ class SMPL:
         #R_out = tf.stack(R_out, axis=0)
         return R_out
 
+    # @tf.function
     def run(self, R, betas, remove=True):
         batch_size = betas.shape[0]
         # R = self.run_rodrigues(pose)
@@ -185,11 +187,14 @@ class SMPL:
             J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
             J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
             
+
             T = tf.tensordot(weights[g], results, axes=((1), (0)))
             rest_shape_h = tf.concat((v_posed, tf.ones((v_posed.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1)
             v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
             v = tf.reshape(v, (-1, 4))[:, :3]
 
+            # J = tf.transpose(tf.tensordot(v, self.J_regressor[g], axes=([0], [1])), [1,0])
+
             return v, tf.squeeze(J), v_posed, lrotmin
 
         gender = tf.cast(tf.greater_equal(betas[:, 10] , 0.5), tf.int32)

smpl_torch.py

@@ -134,13 +134,6 @@ class SMPL(Module):
     ret = torch.cat((zeros43, x), dim=3)
     return ret
 
-  def write_obj(self, verts, file_name):
-    with open(file_name, 'w') as fp:
-      for v in verts:
-        fp.write('v %f %f %f\n' % (v[0], v[1], v[2]))
-
-      for f in self.faces + 1:
-        fp.write('f %d %d %d\n' % (f[0], f[1], f[2]))
 
   def forward(self, betas, pose, remove=True):
     
@@ -208,7 +201,9 @@ class SMPL(Module):
           )
         )
       )
-    # Restart from here
+    
+    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])])
     rest_shape_h = torch.cat(
       (v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2
@@ -216,11 +211,7 @@ class SMPL(Module):
     v = torch.matmul(T, torch.reshape(rest_shape_h, (batch_num, -1, 4, 1)))
     v = torch.reshape(v, (batch_num, -1, 4))[:, :, :3]
     result = v
-    # estimate 3D joint locations
-    # print(result.shape)
-    # print(self.joint_regressor.shape)
-    joints = torch.stack([torch.tensordot(result[i], self.J_regressor[g.to(torch.long)], dims=([0], [1])) for i, g in enumerate(betas[:,-1])]).transpose(1, 2)
-    
+
     root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3))
     result = result - root_joint
     joints = joints - root_joint
@@ -228,37 +219,4 @@ class SMPL(Module):
     if remove:
       result = result[:,self.to_keep,:]
 
-    return result, joints, lrotmin, v_posed
-
-
-def test_gpu(gpu_id=[0], path=''):
-  if len(gpu_id) > 0 and torch.cuda.is_available():
-    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id[0])
-    device = torch.device('cuda')
-  else:
-    device = torch.device('cpu')
-  #print(device)
-
-  pose_size = 72
-  beta_size = 10
-
-  np.random.seed(9608)
-  model = SMPL(device=device, path=path)
-  for i in range(5):
-      pose = torch.from_numpy((np.random.rand(4, pose_size) - 0.5) * 0.4)\
-              .type(torch.float32).to(device)
-      betas = torch.from_numpy((np.random.rand(4, beta_size) - 0.5) * 0.06) \
-              .type(torch.float32).to(device)
-      betas = torch.concat((betas, torch.zeros((4,1), device=device)), axis=1)
-      s = time()
-      result, joints = model(betas, pose)
-      print(time() - s)
-      plot_mesh(result[0].cpu().numpy(), model.f).show()
-      
-  # outmesh_path = './smpl_torch_{}.obj'
-  # for i in range(result.shape[0]):
-      # model.write_obj(result[i], outmesh_path.format(i))
-  
-
-if __name__ == '__main__':
-  test_gpu([5], path='model_files')
+    return result, joints

test.py

@@ -1,10 +1,10 @@
 import os
 
-os.environ["CUDA_VISIBLE_DEVICES"] = '5'
+os.environ["CUDA_VISIBLE_DEVICES"] = '7'
 
 import numpy as np
-from smpl_tf import SMPL as SMPL_tf
 from smpl_torch import SMPL as SMPL_torch
+from smpl_tf import SMPL as SMPL_tf
 from utils import plot_mesh, plot_mesh_2
 from time import time
 
@@ -17,7 +17,7 @@ if __name__ == '__main__':
     smpl_tf : SMPL_tf = SMPL_tf('model_files')
     smpl_torch : SMPL_torch = SMPL_torch(device=device, path='model_files')
 
-
+    np.random.seed(4444)
     pose = ((np.random.rand(4, 72) - 0.5) * 0.4).astype(np.float32)
     betas = ((np.random.rand(4, 10) - 0.5) * 0.06).astype(np.float32)
     betas = np.concatenate((betas, np.array([0,1,0,1], dtype=np.float32)[:,None]), axis=1)
@@ -37,10 +37,10 @@ if __name__ == '__main__':
     # plot_mesh(result[0].cpu().numpy(), model.f).show()
 
     R = smpl_tf.run_rodrigues(pose)
-    result_tf, joints_torch, _, _ = smpl_tf.run(R, betas)
+    result_tf, joints_tf, _, _ = smpl_tf.run(R, betas)
     s = time()
     R = smpl_tf.run_rodrigues(pose)
-    result_tf, joints_torch, v_shaped_2, v_posed_2 = smpl_tf.run(R, betas)
+    result_tf, joints_tf, v_shaped_2, v_posed_2 = smpl_tf.run(R, betas)
     print(time() - s)
 
     result_torch = result_torch.cpu().numpy()