backup

losses.py

@@ -25,107 +25,14 @@ from scipy.spatial import cKDTree
 from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
 from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
 
-from custom_utils.custom_utils import maskFacesBatch, _maskFaces
-from custom_utils.graph_utils import sparse_to_tensor
-
-
-def d_logistic(real_images, generator, discriminator, z_dim, conditions):
-    batch_size = tf.shape(real_images)[0]
-    z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
-
-    # forward pass
-    fake_images = generator([z, conditions], training=True)
-    real_scores = discriminator(real_images, training=True)
-    fake_scores = discriminator(fake_images, training=True)
-
-    # gan loss
-    d_loss = tf.math.softplus(fake_scores)
-    d_loss += tf.math.softplus(-real_scores)
-
-
-    return d_loss
-
-
-def d_logistic_r1_reg(real_images, generator, discriminator, z_dim, conditions):
-    batch_size = tf.shape(real_images)[0]
-    z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
-
-    # forward pass
-    fake_images = generator([z, conditions], training=True)
-    real_scores = discriminator(real_images, training=True)
-    fake_scores = discriminator(fake_images, training=True)
-
-    # gan loss
-    d_loss = tf.math.softplus(fake_scores)
-    d_loss = d_loss + tf.math.softplus(-real_scores)
-
-    # gradient penalty
-    with tf.GradientTape() as r1_tape:
-        r1_tape.watch([real_images, conditions])
-        real_scores = discriminator(real_images, training=True)
-        real_loss = tf.reduce_sum(real_scores)
-
-    real_grads = r1_tape.gradient(real_loss, real_images)
-    r1_penalty = tf.reduce_sum(tf.math.square(real_grads), axis=[1, 2, 3])
-    r1_penalty = tf.expand_dims(r1_penalty, axis=1)
-    return d_loss, r1_penalty
-
-
-def g_logistic_non_saturating(real_images, generator, discriminator, z_dim, conditions):
-    batch_size = tf.shape(real_images)[0]
-    z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
-
-    # forward pass
-    fake_images = generator([z, conditions], training=True)
-    fake_scores = discriminator(fake_images, training=True)
-
-    # gan loss
-    g_loss = tf.math.softplus(-fake_scores)
-    return fake_images, g_loss
-
-
-def g_logistic_ns_pathreg(real_images, generator, discriminator, z_dim,
-                          pl_mean, pl_minibatch_shrink, pl_denorm, pl_decay,
-                          conditions):
-    batch_size = tf.shape(real_images)[0]
-    z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
-
-    pl_minibatch = tf.maximum(1, tf.math.floordiv(batch_size, pl_minibatch_shrink))
-    pl_z = tf.random.normal(shape=[pl_minibatch, z_dim], dtype=tf.float32)
-    pl_conditions = { key: conditions[key][:pl_minibatch] for key in conditions.keys() }
-    # for key in conditions.keys():
-    #     pl_conditions[key] = conditions[key][:pl_minibatch]
-    # pl_conditions = (conditions[0][:pl_minibatch], conditions[1][:pl_minibatch], conditions[2][:pl_minibatch], conditions[3][:pl_minibatch], conditions[4][:pl_minibatch])
-
-    # forward pass
-    fake_images = generator([z, conditions], ret_w_broadcasted=False, training=True)
-    fake_scores = discriminator(fake_images, training=True)
-    g_loss = tf.math.softplus(-fake_scores)
-
-    # Evaluate the regularization term using a smaller minibatch to conserve memory.
-    with tf.GradientTape() as pl_tape:
-        pl_tape.watch([pl_z, pl_conditions])
-        pl_fake_images, pl_w_broadcasted = generator([pl_z, pl_conditions], ret_w_broadcasted=True, training=True)
-
-        pl_noise = tf.random.normal(tf.shape(pl_fake_images)) * pl_denorm
-        pl_noise_applied = tf.reduce_sum(pl_fake_images * pl_noise)
-
-    pl_grads = pl_tape.gradient(pl_noise_applied, pl_w_broadcasted)
-    pl_lengths = tf.math.sqrt(tf.reduce_mean(tf.reduce_sum(tf.math.square(pl_grads), axis=2), axis=1))
-
-    # Track exponential moving average of |J*y|.
-    pl_mean_val = pl_mean + pl_decay * (tf.reduce_mean(pl_lengths) - pl_mean)
-    pl_mean.assign(pl_mean_val)
-
-    # Calculate (|J*y|-a)^2.
-    pl_penalty = tf.square(pl_lengths - pl_mean)
-    return fake_images, g_loss, pl_penalty
 
 
 # from hbertiche/PBNS
 """ COLLISION FUNCTIONS START """
-F = np.load("custom_utils/smpl_faces.npy")
-
+F_coll = np.load("custom_utils/smpl_faces.npy")
+F_coll_skirt = np.load("custom_utils/smpl_faces_skirt.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])
 
 def _nearest_neighbour(V, B):
 	tree = cKDTree(B.numpy())
@@ -137,7 +44,7 @@ def _tf_nn_parallel(V, B):
 
 # @tf.function
 def collision_loss(V, B, mask, thr=.004):
-    B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1]))
+    B_vn = vertex_normals(B, tf.tile(F_coll[None], [B.shape[0], 1, 1]))
    
     # compute correspondences
     idx = _tf_nn_parallel(V, B)
@@ -149,9 +56,11 @@ def collision_loss(V, B, mask, thr=.004):
     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.reduce_sum((tf.minimum(dot - thr, 0)  * mask) ** 2, axis=-1) # / tf.reduce_sum(mask, axis=-1)
+    thresholded = tf.minimum(dot - thr, 0)  * mask
+    loss = tf.reduce_sum(thresholded ** 2, axis=-1)
+    ratio = tf.reduce_sum(tf.cast(thresholded < 0, dtype=tf.float32)) / tf.reduce_sum(mask)
 
-    return loss
+    return loss, ratio
 
 
 def collision_ratio(V, B, mask, thr=.004):
@@ -327,6 +236,11 @@ 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)
+    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))
 # vgg_block_1 = tf.keras.Sequential(vgg.layers[:4])
 # vgg_block_2 = tf.keras.Sequential(vgg.layers[4:7])