test examples

app.py

@@ -1,5 +1,6 @@
 from typing import Tuple
 
+import requests
 import random
 import numpy as np
 import gradio as gr
@@ -20,14 +21,42 @@ MAX_SEED = np.iinfo(np.int32).max
 IMAGE_SIZE = 1024
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+
+def remove_background(image: Image.Image, threshold: int = 50) -> Image.Image:
+    image = image.convert("RGBA")
+    data = image.getdata()
+    new_data = []
+    for item in data:
+        avg = sum(item[:3]) / 3
+        if avg < threshold:
+            new_data.append((0, 0, 0, 0))
+        else:
+            new_data.append(item)
+
+    image.putdata(new_data)
+    return image
+
+
 EXAMPLES = [
     [
         {
-            "background": "https://media.roboflow.com/spaces/doge-2-image.jpeg",
-            "layers": ["https://media.roboflow.com/spaces/doge-2-mask.png"],
-            "composite": None
+            "background": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-image.png", stream=True).raw),
+            "layers": [remove_background(Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-mask-2.png", stream=True).raw))],
+            "composite": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-composite-2.png", stream=True).raw),
+        },
+        "little lion",
+        42,
+        False,
+        0.9,
+        30
+    ],
+[
+        {
+            "background": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-image.png", stream=True).raw),
+            "layers": [remove_background(Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-mask-3.png", stream=True).raw))],
+            "composite": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-composite-3.png", stream=True).raw),
         },
-        "spider tattoo",
+        "tattoos",
         42,
         False,
         0.9,
@@ -76,18 +105,18 @@ def process(
 ):
     if not input_text:
         gr.Info("Please enter a text prompt.")
-        return None
+        return None, None
 
     image = input_image_editor['background']
     mask = input_image_editor['layers'][0]
 
     if not image:
         gr.Info("Please upload an image.")
-        return None
+        return None, None
 
     if not mask:
         gr.Info("Please draw a mask on the image.")
-        return None
+        return None, None
 
     width, height = resize_image_dimensions(original_resolution_wh=image.size)
     resized_image = image.resize((width, height), Image.LANCZOS)
@@ -168,10 +197,10 @@ with gr.Blocks() as demo:
                     )
         with gr.Column():
             output_image_component = gr.Image(
-                type='pil', image_mode='RGB', label='Generated image')
+                type='pil', image_mode='RGB', label='Generated image', format="png")
             with gr.Accordion("Debug", open=False):
                 output_mask_component = gr.Image(
-                    type='pil', image_mode='RGB', label='Input mask')
+                    type='pil', image_mode='RGB', label='Input mask', format="png")
     with gr.Row():
         gr.Examples(
             fn=process,