Update app.py

app.py

@@ -235,14 +235,6 @@ def show_images(images_list):
         axs[c].imshow(images_list[c])
     plt.show()
 
-
-def invert_loss(gen_image):
-    inverter = T.RandomInvert(p=1.0)
-    inverted_img = inverter(gen_image)
-    #loss = torch.abs(gen_image - inverted_img).sum()
-    loss = torch.nn.functional.mse_loss(gen_image[:,0], gen_image[:,2]) + torch.nn.functional.mse_loss(gen_image[:,2], gen_image[:,1]) + torch.nn.functional.mse_loss(gen_image[:,0], gen_image[:,1])
-    return loss
-
 def brilliance_loss(image, target_brilliance=10):
     # Calculate the standard deviation of color channels
     std_dev = torch.std(image, dim=(2, 3))
@@ -252,6 +244,42 @@ def brilliance_loss(image, target_brilliance=10):
     loss = torch.abs(mean_std_dev - target_brilliance)
     return loss
 
+import numpy as np
+from PIL import Image
+
+import torch
+from scipy.stats import wasserstein_distance
+
+def exposure_loss(image, target_exposure = 3):
+    # Calculate the brightness (exposure) of the image.
+    image_brightness = torch.mean(image)
+    
+    # Calculate the loss as the absolute difference from the target exposure.
+    loss = torch.abs(image_brightness - target_exposure)
+    return loss
+
+def color_diversity_loss(images):
+    # Calculate color diversity by measuring the variance of color channels (R, G, B).
+    color_variance = torch.var(images, dim=(2, 3), keepdim=True)
+    # Sum the color variances for each channel to get the total color diversity.
+    total_color_diversity = torch.sum(color_variance, dim=1)
+    return total_color_diversity
+
+def sharpness_loss(images):
+    # Apply the Laplacian filter to the images to measure sharpness.
+    laplacian_filter = torch.Tensor([[-1, -1, -1],
+                                    [-1, 8, -1],
+                                    [-1, -1, -1]]).view(1, 1, 3, 3).to(images.device)
+    
+    # Expand the filter to match the number of channels in the input image.
+    laplacian_filter = laplacian_filter.expand(-1, images.shape[1], -1, -1)
+
+    # Apply the convolution operation.
+    laplacian = torch.abs(F.conv2d(images, laplacian_filter))
+
+    # Calculate sharpness as the negative of the Laplacian variance.
+    sharpness = torch.var(laplacian)
+    return sharpness
 
 def display_images_in_rows(images_with_titles, titles):
     num_images = len(images_with_titles)
@@ -280,41 +308,46 @@ def display_images_in_rows(images_with_titles, titles):
     # plt.show()
 
 
-def image_generator(prompt = "dog", loss_function=None):
-  images_without_loss = []
-  images_with_loss = []
+def image_generator(prompt="cat", loss_function=None):
+    images_without_loss = []
+    images_with_loss = []
 
-  for i in range(num_styles):
-      generated_img = generate_image_custom_style(prompt,style_num = i,random_seed = seed_values[i],custom_loss_fn = None)
-      images_without_loss.append(generated_img)
-      if loss_function:
-        generated_img = generate_image_custom_style(prompt,style_num = i,random_seed = seed_values[i],custom_loss_fn = loss_function)
-        images_with_loss.append(generated_img)
+    for i in range(num_styles):
+        generated_img = generate_image_custom_style(prompt, style_num=i, random_seed=seed_values[i], custom_loss_fn=None)
+        images_without_loss.append(generated_img)
 
-  generated_sd_images = []
-  titles = ["animal toy","fft style","mid journey","oil style","Space style"]
+        if loss_function:
+            if loss_function == "exposure_loss":
+                generated_img = generate_image_custom_style(prompt, style_num=i, random_seed=seed_values[i], custom_loss_fn=exposure_loss)
+            elif loss_function == "color_diversity_loss":
+                generated_img = generate_image_custom_style(prompt, style_num=i, random_seed=seed_values[i], custom_loss_fn=color_diversity_loss)
+            elif loss_function == "sharpness_loss":
+                generated_img = generate_image_custom_style(prompt, style_num=i, random_seed=seed_values[i], custom_loss_fn=sharpness_loss)
+            elif loss_function == "brilliance_loss":
+                generated_img = generate_image_custom_style(prompt, style_num=i, random_seed=seed_values[i], custom_loss_fn=brilliance_loss)
+            images_with_loss.append(generated_img)
 
-  for i in range(len(titles)):
-    generated_sd_images.append((images_without_loss[i], titles[i])) 
-    if images_with_loss != []:
-      generated_sd_images.append((images_with_loss[i], titles[i])) 
+    generated_sd_images = []
+    titles = ["animal toy", "fft style", "mid journey", "oil style", "Space style"]
 
-  return display_images_in_rows(generated_sd_images, titles)
+    for i in range(len(titles)):
+        generated_sd_images.append((images_without_loss[i], titles[i]))
+        if images_with_loss:
+            generated_sd_images.append((images_with_loss[i], titles[i]))
 
-# Create a wrapper function for show_misclassified_images()
-def image_generator_wrapper(prompt = "dog", loss_function=None):
-  if loss_function == "Yes":
-    loss_function = brilliance_loss
-  else:
-    loss_function = None
+    return generated_sd_images
 
-  return image_generator(prompt, loss_function)
+# Create a wrapper function for image_generator()
+def image_generator_wrapper(prompt="dog", selected_loss="None"):
+    return image_generator(prompt, selected_loss)
 
 description = 'Stable Diffusion is a generative artificial intelligence (generative AI) model that produces unique photorealistic images from text and image prompts.'
 title = 'Image Generation using Stable Diffusion'
 
 demo = gr.Interface(image_generator_wrapper,
                     inputs=[gr.Textbox(label="Enter prompt for generation", type="text", value="A ballerina cat dancing in space"),
-                            gr.Radio(["Yes", "No"], value="No"  , label="Apply Brilliance Loss")],
-                    outputs=gr.Plot(label="Generated Images"), title = "Stable Diffusion", description=description)
+                            gr.Radio(["None", "exposure_loss", "color_diversity_loss", "sharpness_loss", "brilliance_loss"], value="None", label="Select Loss")],
+                    outputs=gr.Plot(label="Generated Images"),
+                    title=title,
+                    description=description)
 demo.launch()