docker update

docker/docker-compose.yml

@@ -0,0 +1,21 @@
+name: "cvpr_2023"
+services:
+  lhunor_torch:
+    image: lhunor_torch:v1.13
+    container_name: lhunor_torch_v1.13
+    ipc: host
+    tty: true
+    ulimits:
+      memlock: -1
+      stack: 67108864
+    volumes:
+      - '/home-net/lhunor:/home-net'
+      - '/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]

docker/requirements.txt

@@ -11,3 +11,4 @@ tables
 numpy-stl
 torchmetrics
 pytorch_msssim
+importlib

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(