save

losses.py

@@ -44,7 +44,7 @@ def _nearest_neighbour(V, B):
 def _tf_nn_parallel(V, B):
 	return tf.map_fn(fn=lambda elem: tf.py_function(func=_nearest_neighbour, inp=elem, Tout=tf.int32), elems=(V,B), fn_output_signature=tf.int32, parallel_iterations=V.shape[0])
 
-# @tf.function
+@tf.function
 def collision_loss(V, B, mask, thr=.004):
     B_vn = vertex_normals(B, tf.tile(F_coll[None], [B.shape[0], 1, 1]))
    
@@ -96,6 +96,7 @@ edges_skirt = np.stack(np.nonzero(laplacian_loss.todense())).transpose()
 def get_edge_mask(masks):
     return tf.reduce_min(tf.gather(masks, edges, axis=1), axis=-1)
 
+@tf.function
 def edge_loss(x, y, mask):
     e_mask = get_edge_mask(mask)
     x_e = tf.gather(x, edges[:,0], axis=1) - tf.gather(x, edges[:,1], axis=1)
@@ -237,13 +238,14 @@ def kl_loss(mean, logvar):
     kl_loss *= -0.5
     return kl_loss
 
-
+@tf.function
 def normal_loss(pred, gt, mask, is_skirt=False, is_smpl_orig=False):
     f_norm = None
     if not is_skirt and not is_smpl_orig:
         f_norm = F_norm
     elif is_skirt:
-        f_norm = F_norm_skirt
+        # f_norm = F_norm_skirt
+        f_norm = tf.tile(F_coll_skirt[None,:,:], [pred.shape[0],1,1])
     elif is_smpl_orig:
         f_norm = F_norm_orig
     else: