added remove background

.gitignore

@@ -0,0 +1,4 @@
+.gitattributes
+examples/*
+.vscode/*
+__pycache__/*

app.py

@@ -7,8 +7,17 @@ import svgwrite
 import base64
 import requests
 from io import BytesIO
+from rembg import new_session, remove
 import argparse
 
+# -------------------------------------------------
+# Background removal session (CPU friendly)
+# -------------------------------------------------
+
+print("Loading rembg model...")
+rembg_session = new_session("u2net")
+print("rembg model loaded")
+
 
 # -------------------------------------------------
 # Example Images
@@ -110,16 +119,23 @@ def generate(
     background_mode,
     rounded,
     clipping,
+    remove_bg,
 ):
-
     if image is None:
         return None, None, None
+    
+    if remove_bg:
+        image = remove(image, session=rembg_session)
 
     border_color = rgb_to_hex(border_r, border_g, border_b)
+    
     padding_color = None
-
     if use_padding:
-        padding_color = rgb_to_hex(pad_r, pad_g, pad_b)
+        # If the image has transparency -> keep padding transparent
+        if image.mode == "RGBA":
+            padding_color = None
+        else:
+            padding_color = rgb_to_hex(pad_r, pad_g, pad_b)
 
     mosaic, _ = make_vit_mosaic(
         image,
@@ -211,6 +227,7 @@ with gr.Blocks() as demo:
 
                 rounded = gr.Checkbox(value=True, label="Enable rounded corners")
                 clipping = gr.Checkbox(value=True, label="True rounded clipping")
+                remove_bg = gr.Checkbox(value=False, label="Remove Background (rembg)")
 
                 supersample = gr.Slider(1, 4, value=2, step=1)
                 scale_mode = gr.Radio(["keep", "downscale"], value="keep")
@@ -224,7 +241,7 @@ with gr.Blocks() as demo:
 
             gr.Markdown("### 📥 Selected Image")
 
-            input_image = gr.Image(type="pil", height=250)
+            input_image = gr.Image(type="pil", image_mode="RGBA", height=250)
             credit_display = gr.Markdown("")
 
             gr.Markdown("### 🖼 Mosaic Preview")
@@ -270,6 +287,7 @@ with gr.Blocks() as demo:
             background_mode,
             rounded,
             clipping,
+            remove_bg,
         ],
         outputs=[output_image, download_png, download_svg],
     )

requirements.txt

@@ -4,4 +4,6 @@ pillow~=10.4.0
 svgwrite==1.4.3
 requests==2.31.0
 huggingface-hub==0.34.3
-pydantic==2.10.6
+pydantic==2.10.6
+rembg
+onnxruntime