commit all

losses.py

@@ -31,8 +31,10 @@ from tensorflow_graphics.geometry.convolution.graph_convolution import edge_conv
 """ COLLISION FUNCTIONS START """
 F_coll = np.load("custom_utils/smpl_faces.npy")
 F_coll_skirt = np.load("custom_utils/smpl_faces_skirt.npy")
+F_coll_orig = np.load("custom_utils/smpl_faces_orig.npy")
 F_norm = np.tile(F_coll[np.newaxis,:,:], [4,1,1])
 F_norm_skirt = np.tile(F_coll_skirt[np.newaxis,:,:], [4,1,1])
+F_norm_orig = np.tile(F_coll_orig[np.newaxis,:,:], [4,1,1])
 
 def _nearest_neighbour(V, B):
 	tree = cKDTree(B.numpy())
@@ -236,9 +238,19 @@ def kl_loss(mean, logvar):
     return kl_loss
 
 
-def normal_loss(pred, gt, mask, is_skirt=False):
-    pred_norm = vertex_normals(pred, F_norm_skirt if is_skirt else F_norm)
-    gt_norm = vertex_normals(gt, F_norm_skirt if is_skirt else F_norm)
+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
+    elif is_smpl_orig:
+        f_norm = F_norm_orig
+    else:
+        raise NotImplementedError
+
+    pred_norm = vertex_normals(pred, f_norm)
+    gt_norm = vertex_normals(gt, f_norm)
     return tf.reduce_mean(tf.reduce_mean((tf.reduce_sum(tf.math.abs(pred_norm - gt_norm), axis=-1) * mask), axis=-1))
 
 # vgg = tf.keras.applications.vgg16.VGG16(include_top=False, input_shape=(512,256,3))