tried new losses

losses.py

@@ -1,6 +1,23 @@
+# import os
+# os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+# import tensorflow as tf
+
+# gpus = tf.config.experimental.list_physical_devices('GPU')
+# if gpus:
+#     try:
+#         # Currently, memory growth needs to be the same across GPUs
+#         for gpu in gpus:
+#             tf.config.experimental.set_memory_growth(gpu, True)
+#         logical_gpus = tf.config.experimental.list_logical_devices('GPU')
+#         print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
+#     except RuntimeError as e:
+#         # Memory growth must be set before GPUs have been initialized
+#         print(e)
+
+import tensorflow as tf
+
 import numpy as np
 import tensorflow_graphics as tfg
-import tensorflow as tf
 import scipy as sp
 import scipy.io as sio
 
@@ -136,6 +153,25 @@ def collision_loss(V, B, mask, thr=.004):
 
     return loss
 
+
+def collision_ratio(V, B, mask, thr=.004):
+    B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1]))
+   
+    # compute correspondences
+    idx = _tf_nn_parallel(V, B)
+    N, M = V.shape[:2]
+    # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
+    _idx = tf.ones((N, M), dtype=tf.int32) * np.array(list(range(N)), np.int32)[:,None]
+    idx = tf.stack((_idx, idx), axis=-1)
+    # compute loss
+    D = V - tf.gather_nd(B, idx)
+    b_vn = tf.gather_nd(B_vn, idx)
+    dot = tf.einsum('abc,abc->ab', D, b_vn)
+    loss = tf.minimum(dot - thr, 0)  * mask # / tf.reduce_sum(mask, axis=-1)
+    ratio = tf.reduce_sum(tf.cast(loss < 0, dtype=tf.float32)) / tf.reduce_sum(mask)
+
+    return ratio
+
 """ COLLISION FUNCTIONS END """
 
 
@@ -291,9 +327,34 @@ def kl_loss(mean, logvar):
     return kl_loss
 
 
+# vgg = tf.keras.applications.vgg16.VGG16(include_top=False, input_shape=(512,256,3))
+# vgg_block_1 = tf.keras.Sequential(vgg.layers[:4])
+# vgg_block_2 = tf.keras.Sequential(vgg.layers[4:7])
+# vgg_block_3 = tf.keras.Sequential(vgg.layers[7:11])
+
+# def concept_loss(pred, gt):
+#     pred = tf.keras.applications.vgg16.preprocess_input((pred + 1) * 128)
+#     gt = tf.keras.applications.vgg16.preprocess_input((gt + 1) * 128)
+
+#     pred_features_1 = vgg_block_1(pred, training=False)
+#     pred_features_2 = vgg_block_2(pred_features_1, training=False)
+#     pred_features_3 = vgg_block_3(pred_features_2, training=False)
+#     gt_features_1 = vgg_block_1(gt, training=False)
+#     gt_features_2 = vgg_block_2(gt_features_1, training=False)
+#     gt_features_3 = vgg_block_3(gt_features_2, training=False)
+
+#     l1_1 = tf.reduce_mean(tf.math.abs(pred_features_1 - gt_features_1))
+#     l1_2 = tf.reduce_mean(tf.math.abs(pred_features_2 - gt_features_2))
+#     l1_3 = tf.reduce_mean(tf.math.abs(pred_features_3 - gt_features_3))
+
+#     return l1_1 + l1_2 + l1_3
+
+
 if __name__ == '__main__':
-    V = tf.ones((4, 14475, 3))
-    B = tf.zeros((4, 14475, 3))
+
+
+    V = tf.ones((4, 512, 256, 3))
+    B = tf.ones((4, 512, 256, 3))
     m = tf.ones((4, 14475))
-    loss_1 = collision_loss(V, B, m, thr=.004)
+    loss_1 = concept_loss(V, B)
     print()