making ready for skirts

losses.py

@@ -142,6 +142,10 @@ def collision_loss(V, B, mask, thr=.004):
 laplacian_loss = sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo()
 edges = np.stack(np.nonzero(laplacian_loss.todense())).transpose()
 
+laplacian_loss_skirt = sio.loadmat("custom_utils/ssmpl_skirt0.mat")['L'].tocoo()
+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)
 
@@ -158,6 +162,22 @@ def edge_loss(x, y, mask):
     loss = tf.reduce_sum(d_e ** 2, axis=-1)# / tf.reduce_sum(e_mask, axis=-1)
     return loss
 
+def get_edge_mask_skirt(masks):
+    return tf.reduce_min(tf.gather(masks, edges_skirt, axis=1), axis=-1)
+
+def edge_loss_skirt(x, y, mask):
+    e_mask = get_edge_mask_skirt(mask)
+    x_e = tf.gather(x, edges_skirt[:,0], axis=1) - tf.gather(x, edges_skirt[:,1], axis=1)
+    x_e = tf.reduce_sum(x_e ** 2, -1)
+    x_e = tf.sqrt(x_e + 1e-9)
+    y_e = tf.gather(y, edges_skirt[:,0], axis=1) - tf.gather(y, edges_skirt[:,1], axis=1)
+    y_e = tf.reduce_sum(y_e ** 2, -1)
+    y_e = tf.sqrt(y_e + 1e-9)
+    d_e = y_e - x_e
+    d_e = d_e * e_mask
+    loss = tf.reduce_sum(d_e ** 2, axis=-1)# / tf.reduce_sum(e_mask, axis=-1)
+    return loss
+
 
 class PairwiseDistance(tf.keras.layers.Layer):
     def __init__(self, **kwargs):