working weight learning

custom_utils/custom_utils.py

@@ -2,6 +2,7 @@ import json
 import numpy as np
 import tensorflow as tf
 import scipy as sp
+import cmapy as cm
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
 
@@ -216,6 +217,35 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
     return fig
 
 
+def plot_weights(weights, body, faces):
+    fig = go.Figure()
+    colors = np.array([cm.color('jet', 1 - w) for w in weights.numpy()])
+
+    # if flip_xy:
+    #     mesh_gt = flip_axis(mesh_gt)
+    #     mesh_pred_unposed = flip_axis(mesh_pred_unposed)
+    #     collision = flip_axis(collision, ind=14475)
+
+    fig.add_trace(go.Mesh3d(
+                x=body[:,0],
+                y=body[:,2],
+                z=body[:,1],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+                vertexcolor=colors
+            ))
+    
+
+    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),)
+    helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
+                yaxis = dict(nticks=4, range=[-1,1],),
+                zaxis = dict(nticks=4, range=[-1,1],),)
+    fig.update_layout(scene = helper,)
+
+    return fig
+
+
 def plot_3d_collision(gt, collision_1, collision_2, collision_faces, collision_color, flip_xy=False):
     fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])
 

custom_utils/smpl.py

@@ -406,7 +406,8 @@ class SMPL:
         return T_out, J_out
 
 
-    def run_garment_with_weights(self, R, betas, v_cloth, weights, J=None):
+    def run_garment_with_weights(self, pose, betas, v_cloth, weights, J=None):
+        R = self.run_rodrigues(pose)
         
         if J is None:
             J = self.get_joints_template(betas)
@@ -430,13 +431,12 @@ class SMPL:
             J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
             
             T = tf.tensordot(weights, results, axes=((1), (0)))
-            rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1)
+            rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.shape[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)
 
-        # gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
         T_out, J_out = tf.map_fn(single_run, (R, J, v_cloth, weights), (tf.float32, tf.float32))
         
         J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))