updated pipeline, added uv_mask_loss

losses.py

@@ -187,6 +187,17 @@ def collision_penalty(va, vb, mask, eps=1e-3):
 
     return tf.reduce_sum(interpenetration**3) / batch_size
 
+whole_garment_uv_mask = np.load('custom_utils/mask_256x512.npy')
+
+def uv_mask_loss(pred, uv_mask):
+    inv_mask = (1 - uv_mask)  * whole_garment_uv_mask
+    pred = tf.reduce_sum(pred, axis=1)
+    invalid = pred * inv_mask
+    invalid = invalid > 0.01
+    invalid_count = tf.math.count_nonzero(invalid, dtype=tf.float32)
+    total_count = tf.math.count_nonzero(inv_mask, dtype=tf.float32)
+    return invalid_count / total_count
+
 
 if __name__ == '__main__':
     V = tf.ones((4, 14475, 3))