test examples

app.py

@@ -22,29 +22,44 @@ 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": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-image.png", stream=True).raw),
-            "layers": [Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-mask-2.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.8,
+        0.85,
         30
     ],
-[
+    [
         {
             "background": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-image.png", stream=True).raw),
-            "layers": [Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-mask-3.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),
         },
         "tattoos",
         42,
         False,
-        0.8,
+        0.85,
         30
     ]
 ]