Upload app.py

app.py

@@ -202,17 +202,19 @@ def blue_loss(images):
     
     return -variance
 
+import torch
+
 def ymca_loss(images, weights=(1.0, 1.0, 1.0, 1.0)):
     """
     Computes the YMCA loss for a batch of images.
     
     The YMCA loss is a custom loss function combining the mean value of the Y (luminance) channel,
     the mean value of the M (magenta) channel, the variance of the C (cyan) channel, and the 
-    absolute sum of the A (alpha) channel.
+    absolute sum of the A (alpha) channel if present.
 
     Parameters:
     images (torch.Tensor): A batch of images. Expected shape is (N, C, H, W) where
-                           N is the batch size, C is the number of channels (assumed 4 for RGBA), 
+                           N is the batch size, C is the number of channels (3 for RGB or 4 for RGBA), 
                            H is the height, and W is the width.
     weights (tuple): A tuple of four floats representing the weights for each component of the loss
                      (default is (1.0, 1.0, 1.0, 1.0)).
@@ -220,15 +222,15 @@ def ymca_loss(images, weights=(1.0, 1.0, 1.0, 1.0)):
     Returns:
     torch.Tensor: The YMCA loss, combining the specified components.
     """
-    # Ensure the input tensor has the correct shape
-    if images.shape[1] != 4:
-        raise ValueError("Expected images with 4 channels (RGBA), but got shape {}".format(images.shape))
+    num_channels = images.shape[1]
+    
+    if num_channels not in [3, 4]:
+        raise ValueError("Expected images with 3 (RGB) or 4 (RGBA) channels, but got shape {}".format(images.shape))
     
-    # Extract the RGBA channels
+    # Extract the RGB channels
     R = images[:, 0, :, :]
     G = images[:, 1, :, :]
     B = images[:, 2, :, :]
-    A = images[:, 3, :, :]
     
     # Convert RGB to Y (luminance) channel
     Y = 0.299 * R + 0.587 * G + 0.114 * B
@@ -248,11 +250,15 @@ def ymca_loss(images, weights=(1.0, 1.0, 1.0, 1.0)):
     # Compute the variance of the C channel
     variance_C = torch.var(C)
     
-    # Compute the absolute sum of the A channel
-    abs_sum_A = torch.sum(torch.abs(A))
+    loss = weights[0] * mean_Y + weights[1] * mean_M - weights[2] * variance_C
     
-    # Combine the components with the given weights
-    loss = (weights[0] * mean_Y) + (weights[1] * mean_M) - (weights[2] * variance_C) + (weights[3] * abs_sum_A)
+    if num_channels == 4:
+        # Extract the alpha channel
+        A = images[:, 3, :, :]
+        # Compute the absolute sum of the A channel
+        abs_sum_A = torch.sum(torch.abs(A))
+        # Include the alpha component in the loss
+        loss += weights[3] * abs_sum_A
     
     return loss