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

@@ -84,11 +84,12 @@ class SMPL:
         self.to_keep = sio.loadmat(os.path.join(path, 'keep_remove.mat'))['keep'].reshape(-1)
 
         self.weights_all = self.weights
-        if use_one_hot:
-            self.weights = tf.gather(self.weights_one_hot, self.to_keep, axis=0)
-            self.weights = tf.tile(self.weights[tf.newaxis,:,:], [2,1,1])
-        else:
-            self.weights = tf.gather(self.weights, self.to_keep, axis=1)
+        # if use_one_hot:
+        #     self.weights = tf.gather(self.weights_one_hot, self.to_keep, axis=0)
+        #     self.weights = tf.tile(self.weights[tf.newaxis,:,:], [2,1,1])
+        # else:
+        #     self.weights = tf.gather(self.weights, self.to_keep, axis=1)
+        
 
         self.shape = (1,27578,3)
 
@@ -140,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)))
@@ -150,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)
@@ -184,26 +187,29 @@ 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]
 
-            return v, tf.squeeze(J), v_posed
+            # 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)
-        T_out, J_out, v_posed = tf.map_fn(single_run, (R, v_shaped, gender), (tf.float32, tf.float32, tf.float32))
+        T_out, J_out, v_posed,lrotmin = tf.map_fn(single_run, (R, v_shaped, gender), (tf.float32, tf.float32, tf.float32, tf.float32))
         
         J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
         J_out = tf.subtract(J_out, J0)
         T_out = tf.subtract(T_out, J0)
-        v_posed = tf.subtract(v_posed, J0)
+        # v_posed = tf.subtract(v_posed, J0)
         
         if remove:
             T_out = tf.gather(T_out, self.to_keep, axis=1)
-            v_posed = tf.gather(v_posed, self.to_keep, axis=1)
+            # v_posed = tf.gather(v_posed, self.to_keep, axis=1)
         
-        return T_out, J_out, v_posed
+        return T_out, J_out, lrotmin, v_posed
 
 
     def run_with_weights(self, R, betas, v_cloth, weights=None, filter=None):

smpl_torch.py

@@ -0,0 +1,222 @@
+# based on https://github.com/CalciferZh/SMPL/blob/master/smpl_torch_batch.py
+
+import os
+import sys
+import pickle
+import torch
+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):
+    
+    super(SMPL, self).__init__()
+    self.J_regressor = []
+    self.v_template = []
+    self.shapedirs = []
+    self.posedirs = []
+    self.weights = []
+    self.is_skirt = is_skirt
+    # Load SMPL models
+    # path = os.path.abspath(os.path.dirname(__file__))
+    with open(os.path.join(path, 'super_model_f.pkl'), 'rb') as f:
+        female = pickle.load(f, encoding="latin1")
+        self.J_regressor.append(torch.tensor(np.array(female['J_regressor'].todense(), dtype=np.float32)))
+        self.v_template.append(torch.tensor(female['v_template'], dtype=torch.float32))
+        self.posedirs.append(torch.tensor(female['posedirs'], dtype=torch.float32))
+        self.shapedirs.append(torch.tensor(female['shapedirs'], dtype=torch.float32))
+        self.weights.append(torch.tensor(female['weights'], dtype=torch.float32))
+    with open(os.path.join(path, 'super_model_m.pkl'), 'rb') as f:
+        male = pickle.load(f, encoding="latin1")
+        self.J_regressor.append(torch.tensor(np.array(male['J_regressor'].todense(), dtype=np.float32)))
+        self.v_template.append(torch.tensor(male['v_template'], dtype=torch.float32))
+        self.posedirs.append(torch.tensor(male['posedirs'], dtype=torch.float32))
+        self.shapedirs.append(torch.tensor(male['shapedirs'], dtype=torch.float32))
+        self.weights.append(torch.tensor(male['weights'], dtype=torch.float32))
+        kintree_table = male['kintree_table']
+        id_to_col = {kintree_table[1, i]: i for i in range(kintree_table.shape[1])}
+        self.parent = {i: id_to_col[kintree_table[0, i]] for i in range(1, kintree_table.shape[1])}
+    self.J_regressor = torch.stack(self.J_regressor, axis=0).to(device)
+    self.v_template = torch.stack(self.v_template, axis=0).to(device)
+    self.shapedirs = torch.stack(self.shapedirs, axis=0).to(device)
+    self.posedirs = torch.stack(self.posedirs, axis=0).to(device)
+    self.weights = torch.stack(self.weights, axis=0).to(device)
+
+    # Topology no head/hands/feet
+    F = sio.loadmat(os.path.join(path, 'ssmpl_faces.mat'))
+    self.f = np.array(F['faces_default'])
+    self.f_skirt = np.array(F['faces_skirt'])
+    self.nn = np.array(sio.loadmat(os.path.join(path, 'NN.mat')))
+    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.shape = (1,27578,3)
+    self.device = device
+
+    if orig_smpl:
+        self.to_keep = self.to_keep[self.to_keep < 6890]
+        self.J_regressor = self.J_regressor[:,:,:6890]
+        self.v_template = self.v_template[:,:6890,:]
+        self.shapedirs = self.shapedirs[:,:6890,:,:]
+        self.posedirs = self.posedirs[:,:6890,:,:]
+        self.weights_all = self.weights_all[:,:6890,:]
+        self.weights = self.weights_all[:,self.to_keep,:]
+        self.shape = (1, 6890, 3)
+
+  @staticmethod
+  def rodrigues(r):
+    """
+    Rodrigues' rotation formula that turns axis-angle tensor into rotation
+    matrix in a batch-ed manner.
+    Parameter:
+    ----------
+    r: Axis-angle rotation tensor of shape [batch_size * angle_num, 1, 3].
+    Return:
+    -------
+    Rotation matrix of shape [batch_size * angle_num, 3, 3].
+    """
+    # eps = r.clone().normal_(std=1e-8)
+    # theta = torch.norm(r + eps, dim=(1, 2), keepdim=True)  # dim cannot be tuple
+    theta = torch.norm(r, dim=(1, 2), keepdim=True)
+    theta = torch.maximum(theta, torch.ones_like(theta) * np.finfo(np.float32).eps)
+
+    theta_dim = theta.shape[0]
+    r_hat = r / theta
+    # cos = torch.cos(theta)
+    z_stick = torch.zeros(theta_dim, dtype=torch.float32).to(r.device)
+    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))
+    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)
+    # dot = torch.matmul(A, r_hat)
+    # R = cos * i_cube + (1 - cos) * dot + torch.sin(theta) * m
+    R = i_cube + m * torch.sin(theta) + torch.matmul(m, m) * (1 - torch.cos(theta))
+    return R
+
+  @staticmethod
+  def with_zeros(x):
+    """
+    Append a [0, 0, 0, 1] tensor to a [3, 4] tensor.
+    Parameter:
+    ---------
+    x: Tensor to be appended.
+    Return:
+    ------
+    Tensor after appending of shape [4,4]
+    """
+    ones = torch.tensor(
+      [[[0.0, 0.0, 0.0, 1.0]]], dtype=torch.float32
+    ).expand(x.shape[0],-1,-1).to(x.device)
+    ret = torch.cat((x, ones), dim=1)
+    return ret
+
+  @staticmethod
+  def pack(x):
+    """
+    Append zero tensors of shape [4, 3] to a batch of [4, 1] shape tensor.
+    Parameter:
+    ----------
+    x: A tensor of shape [batch_size, 4, 1]
+    Return:
+    ------
+    A tensor of shape [batch_size, 4, 4] after appending.
+    """
+    zeros43 = torch.zeros(
+      (x.shape[0], x.shape[1], 4, 3), dtype=torch.float32).to(x.device)
+    ret = torch.cat((zeros43, x), dim=3)
+    return ret
+
+
+  def forward(self, betas, pose, remove=True):
+    
+    """
+          Construct a compute graph that takes in parameters and outputs a tensor as
+          model vertices. Face indices are also returned as a numpy ndarray.
+          
+          20190128: Add batch support.
+          Parameters:
+          ---------
+          pose: Also known as 'theta', an [N, 24, 3] tensor indicating child joint rotation
+          relative to parent joint. For root joint it's global orientation.
+          Represented in a axis-angle format.
+          betas: Parameter for model shape. A tensor of shape [N, 10] as coefficients of
+          PCA components. Only 10 components were released by SMPL author.
+          trans: Global translation tensor of shape [N, 3].
+          Return:
+          ------
+          A 3-D tensor of [N * 6890 * 3] for vertices,
+          and the corresponding [N * 19 * 3] joint positions.
+    """
+    batch_num = betas.shape[0]
+
+    # 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[1], dims=[[1], [2]]) + torch.reshape(self.v_template[1], shape=self.shape)) # male
+    v_shaped = torch.stack(v_shaped, axis=-1)
+    
+    # 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 = 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(
+      self.with_zeros(torch.cat((R_cube_big[:, 0], torch.reshape(J[:, 0, :], (-1, 3, 1))), dim=2))
+    )
+    for i in range(1, 24):
+      results.append(
+        torch.matmul(
+          results[self.parent[i]],
+          self.with_zeros(
+            torch.cat(
+              (R_cube_big[:, i], torch.reshape(J[:, i, :] - J[:, self.parent[i], :], (-1, 3, 1))),
+              dim=2
+            )
+          )
+        )
+      )
+    
+    stacked = torch.stack(results, dim=1)
+    results = stacked - \
+      self.pack(
+        torch.matmul(
+          stacked,
+          torch.reshape(
+            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]
+            
+    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
+    )
+    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
+
+    root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3))
+    result = result - root_joint
+    joints = joints - root_joint
+
+    if remove:
+      result = result[:,self.to_keep,:]
+
+    return result, joints

test.py

@@ -0,0 +1,52 @@
+import os
+
+os.environ["CUDA_VISIBLE_DEVICES"] = '7'
+
+import numpy as np
+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
+
+import torch
+import tensorflow as tf
+
+if __name__ == '__main__':
+    device = torch.device('cuda')
+
+    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)
+
+    # pose = np.zeros((4, 72), dtype=np.float32)
+    # betas = np.zeros((4, 11), dtype=np.float32)
+    
+    pose_torch = torch.tensor(pose, device=device)
+    betas_torch = torch.tensor(betas, device=device)
+
+    result_torch, joints_torch, _, _ = smpl_torch(betas_torch, pose_torch)
+    s = time()
+    result_torch, joints_torch, v_shaped_1, v_posed_1 = smpl_torch(betas_torch, pose_torch)
+    print(time() - s)
+    R_1 = smpl_torch.rodrigues(pose_torch.view(-1, 1, 3)).reshape(4, -1, 3, 3)
+
+    # plot_mesh(result[0].cpu().numpy(), model.f).show()
+
+    R = smpl_tf.run_rodrigues(pose)
+    result_tf, joints_tf, _, _ = smpl_tf.run(R, betas)
+    s = time()
+    R = smpl_tf.run_rodrigues(pose)
+    result_tf, joints_tf, v_shaped_2, v_posed_2 = smpl_tf.run(R, betas)
+    print(time() - s)
+
+    result_torch = result_torch.cpu().numpy()
+    result_tf = result_tf.numpy()
+    diff = np.abs(result_torch - result_tf).sum(-1).sum(-1)
+
+    plot_mesh_2(result_torch[0], result_tf[0], smpl_tf.f).show()
+    plot_mesh_2(result_torch[1], result_tf[1], smpl_tf.f).show()
+    pass

utils.py

@@ -1,4 +1,5 @@
 import plotly.graph_objects as go
+from plotly.subplots import make_subplots
 
 
 def plot_mesh(verts, faces):
@@ -15,11 +16,53 @@ def plot_mesh(verts, faces):
         ),)
 
     fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1))
-    helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
-                yaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
-                zaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),)
+    # helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
+    #             yaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
+    #             zaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),)
+    helper = dict(xaxis = dict(range=[-1,1], ),
+                yaxis = dict(range=[-1,1], ),
+                zaxis = dict(range=[-1,1], ) )
     fig.update_layout(scene = helper)
     
     fig.update_layout(margin=dict(l=0, r=0, b=0, t=0, pad=0))
 
     return fig
+
+
+def plot_mesh_2(verts_1, verts_2, faces):
+    # fig = make_subplots(rows=1, cols=2, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(2)] ])
+    fig = go.Figure()
+    
+    fig.add_trace(go.Mesh3d(
+            x=verts_1[:,0],
+            y=verts_1[:,2],
+            z=verts_1[:,1],
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+            text="",            
+        ),)
+
+    fig.add_trace(go.Mesh3d(
+            x=verts_2[:,0],
+            y=verts_2[:,2],
+            z=verts_2[:,1],
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+            text="",            
+        ),)
+
+    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
+                    # scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1)
+                    )
+    helper = dict(xaxis = dict(range=[-1,1], ),
+                yaxis = dict(range=[-1,1], ),
+                zaxis = dict(range=[-1,1], ) )
+    fig.update_layout(scene = helper, 
+                        # scene2 = helper
+    )
+    
+    fig.update_layout(margin=dict(l=0, r=0, b=0, t=0, pad=0))
+
+    return fig