Update app.py

app.py

@@ -1,300 +1,48 @@
-"""
-Background Remover Tool
-========================
-A fully open-source background removal app powered by rembg + Gradio.
-Runs on CPU. Ready for HuggingFace Spaces deployment.
-
-Supported models:
-  - u2net          → General purpose (best overall quality)
-  - u2netp         → Lightweight/fast (lower RAM usage)
-  - u2net_human_seg → Optimised for people / portraits
-  - isnet-general-use → High-resolution dichotomous segmentation
-  - birefnet-general  → State-of-the-art bilateral-reference net
-"""
-
-import os
-import io
-import time
-import numpy as np
 import gradio as gr
-from PIL import Image, ImageColor
-from rembg import remove, new_session
-
-# ─────────────────────────────────────────────
-# Model registry
-# ─────────────────────────────────────────────
-MODELS = {
-    "u2net  (General – best quality)":        "u2net",
-    "u2netp  (Lightweight – fastest)":        "u2netp",
-    "u2net_human_seg  (People / Portraits)":  "u2net_human_seg",
-    "isnet-general-use  (High-res objects)":  "isnet-general-use",
-    "birefnet-general  (State-of-the-art)":   "birefnet-general",
-}
-
-# Cache sessions so we don't reload the model on every call
-_session_cache: dict = {}
-
-def get_session(model_key: str):
-    model_name = MODELS[model_key]
-    if model_name not in _session_cache:
-        _session_cache[model_name] = new_session(model_name)
-    return _session_cache[model_name]
-
-
-# ─────────────────────────────────────────────
-# Core processing
-# ─────────────────────────────────────────────
-def remove_background(
-    input_image: Image.Image,
-    model_key: str,
-    bg_mode: str,
-    bg_color_hex: str,
-    bg_image: Image.Image | None,
-    alpha_matting: bool,
-    alpha_fg: int,
-    alpha_bg: int,
-    alpha_erode: int,
-    post_process: bool,
-) -> tuple[Image.Image, Image.Image, str]:
-    """Remove the background and optionally replace it."""
-
+from rembg import remove
+from PIL import Image
+
+def remove_background(input_image):
+    """
+    Takes a PIL Image, removes the background using rembg, 
+    and returns the processed image with a transparent background.
+    """
     if input_image is None:
-        raise gr.Error("Please upload an image first.")
-
-    # Ensure RGBA input
-    src = input_image.convert("RGBA")
-    start = time.time()
-
-    session = get_session(model_key)
-
-    # Build rembg kwargs
-    kwargs: dict = dict(
-        session=session,
-        alpha_matting=alpha_matting,
-        post_process_mask=post_process,
+        return None
+    
+    # The remove() function from rembg handles the heavy lifting
+    output_image = remove(input_image)
+    return output_image
+
+# Build the Gradio UI
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # ✂️ AI Background Remover
+        Upload any image to instantly remove its background. This tool uses the open-source `rembg` library and runs efficiently on CPU.
+        """
     )
-    if alpha_matting:
-        kwargs["alpha_matting_foreground_threshold"] = alpha_fg
-        kwargs["alpha_matting_background_threshold"] = alpha_bg
-        kwargs["alpha_matting_erode_size"]           = alpha_erode
-
-    result: Image.Image = remove(src, **kwargs)   # RGBA, transparent BG
-
-    elapsed = time.time() - start
-
-    # ── Compose background ────────────────────
-    if bg_mode == "Transparent (PNG)":
-        final = result
-    elif bg_mode == "Solid colour":
-        try:
-            colour = ImageColor.getrgb(bg_color_hex)
-        except ValueError:
-            colour = (255, 255, 255)
-        bg = Image.new("RGBA", result.size, colour + (255,))
-        bg.paste(result, mask=result.split()[3])
-        final = bg.convert("RGB")
-    elif bg_mode == "Custom image" and bg_image is not None:
-        bg = bg_image.convert("RGBA").resize(result.size, Image.LANCZOS)
-        bg.paste(result, mask=result.split()[3])
-        final = bg.convert("RGB")
-    else:
-        # fallback: white
-        bg = Image.new("RGBA", result.size, (255, 255, 255, 255))
-        bg.paste(result, mask=result.split()[3])
-        final = bg.convert("RGB")
-
-    # ── Alpha mask preview ────────────────────
-    mask_preview = result.split()[3].convert("RGB")
-
-    status = (
-        f"✅  Done in {elapsed:.2f}s  |  "
-        f"Model: {MODELS[model_key]}  |  "
-        f"Size: {input_image.width}×{input_image.height}px"
-    )
-    return result, final, mask_preview, status
-
-
-# ─────────────────────────────────────────────
-# UI helpers
-# ─────────────────────────────────────────────
-DESCRIPTION = """
-# 🖼️ AI Background Remover
-**Remove backgrounds instantly — 100 % open-source, runs on CPU.**
-
-Choose a model, upload your photo, and hit **Remove Background**.  
-You can replace the background with a solid colour, a custom image, or keep it transparent.
-"""
-
-TIPS = """
-### 💡 Tips
-- **u2net** gives the best overall results for mixed content.  
-- **u2net_human_seg** is optimised for portraits & selfies.  
-- **birefnet-general** is the most accurate but slowest on CPU.  
-- Enable **Alpha matting** for fine hair / fur details.  
-- Download the transparent PNG to use in design tools.
-"""
-
-
-# ─────────────────────────────────────────────
-# Gradio Blocks UI
-# ─────────────────────────────────────────────
-with gr.Blocks(
-    title="AI Background Remover",
-    theme=gr.themes.Soft(
-        primary_hue=gr.themes.colors.violet,
-        secondary_hue=gr.themes.colors.blue,
-        neutral_hue=gr.themes.colors.slate,
-        font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui"],
-    ),
-    css="""
-        #header { text-align: center; padding: 8px 0 4px 0; }
-        #status_box { font-size: 0.9rem; color: #555; }
-        .gr-image { border-radius: 12px; }
-        footer { display: none !important; }
-    """,
-) as demo:
-
-    gr.Markdown(DESCRIPTION, elem_id="header")
-
+    
     with gr.Row():
-        # ── LEFT: inputs ─────────────────────
-        with gr.Column(scale=1):
-            gr.Markdown("### 📤 Input")
-            input_img = gr.Image(
-                label="Upload Image",
-                type="pil",
-                height=340,
-                sources=["upload", "clipboard"],
-            )
-
-            model_dd = gr.Dropdown(
-                label="🤖 Model",
-                choices=list(MODELS.keys()),
-                value=list(MODELS.keys())[0],
-                info="Different models suit different subjects.",
-            )
-
-            bg_mode_radio = gr.Radio(
-                label="🎨 Background Mode",
-                choices=["Transparent (PNG)", "Solid colour", "Custom image"],
-                value="Transparent (PNG)",
-            )
-
-            with gr.Group(visible=False) as colour_group:
-                bg_color = gr.ColorPicker(label="Colour", value="#ffffff")
-
-            with gr.Group(visible=False) as image_group:
-                bg_img = gr.Image(
-                    label="Background Image",
-                    type="pil",
-                    height=180,
-                )
-
-            with gr.Accordion("⚙️ Advanced Options", open=False):
-                alpha_chk = gr.Checkbox(
-                    label="Alpha Matting (better edges for hair/fur)",
-                    value=False,
-                )
-                with gr.Group(visible=False) as alpha_group:
-                    alpha_fg_sl = gr.Slider(1, 255, value=240, step=1,
-                                            label="Foreground Threshold")
-                    alpha_bg_sl = gr.Slider(1, 255, value=10,  step=1,
-                                            label="Background Threshold")
-                    alpha_er_sl = gr.Slider(0, 40,  value=10,  step=1,
-                                            label="Erode Size")
-                post_chk = gr.Checkbox(
-                    label="Post-process mask (smoother edges)",
-                    value=False,
-                )
-
-            run_btn = gr.Button("✂️  Remove Background", variant="primary", size="lg")
-
-        # ── RIGHT: outputs ────────────────────
-        with gr.Column(scale=1):
-            gr.Markdown("### 📥 Output")
-            with gr.Tabs():
-                with gr.Tab("Result (transparent)"):
-                    out_transparent = gr.Image(
-                        label="Transparent PNG",
-                        type="pil",
-                        height=280,
-                        format="png",
-                        show_download_button=True,
-                    )
-                with gr.Tab("Result (with background)"):
-                    out_composed = gr.Image(
-                        label="Composed Result",
-                        type="pil",
-                        height=280,
-                        show_download_button=True,
-                    )
-                with gr.Tab("Mask preview"):
-                    out_mask = gr.Image(
-                        label="Alpha Mask",
-                        type="pil",
-                        height=280,
-                    )
-
-            status_box = gr.Textbox(
-                label="",
-                interactive=False,
-                elem_id="status_box",
-                placeholder="Status will appear here after processing…",
-            )
-
-    gr.Markdown(TIPS)
-
-    # ── Example images ─────────────────────
-    gr.Examples(
-        examples=[
-            ["examples/portrait.jpg",   list(MODELS.keys())[2], "Transparent (PNG)", "#ffffff", None, False, 240, 10, 10, False],
-            ["examples/product.jpg",    list(MODELS.keys())[0], "Solid colour",      "#f0f0f0", None, False, 240, 10, 10, False],
-            ["examples/landscape.jpg",  list(MODELS.keys())[3], "Transparent (PNG)", "#ffffff", None, True,  240, 10, 10, True],
-        ],
-        inputs=[
-            input_img, model_dd, bg_mode_radio, bg_color, bg_img,
-            alpha_chk, alpha_fg_sl, alpha_bg_sl, alpha_er_sl, post_chk,
-        ],
-        label="📷 Example Images (click to try)",
-        cache_examples=False,
-    )
-
-    # ── Wiring ─────────────────────────────
-    def toggle_bg_mode(mode):
-        return (
-            gr.update(visible=mode == "Solid colour"),
-            gr.update(visible=mode == "Custom image"),
-        )
-
-    bg_mode_radio.change(
-        toggle_bg_mode,
-        inputs=bg_mode_radio,
-        outputs=[colour_group, image_group],
-    )
-
-    alpha_chk.change(
-        lambda v: gr.update(visible=v),
-        inputs=alpha_chk,
-        outputs=alpha_group,
-    )
-
-    run_btn.click(
+        with gr.Column():
+            input_img = gr.Image(type="pil", label="Original Image")
+            submit_btn = gr.Button("Remove Background", variant="primary")
+            clear_btn = gr.ClearButton(components=[input_img])
+            
+        with gr.Column():
+            # Output format is explicitly set to png to preserve transparency
+            output_img = gr.Image(type="pil", label="Transparent Result", format="png")
+            
+    # Connect the button to the function
+    submit_btn.click(
         fn=remove_background,
-        inputs=[
-            input_img, model_dd, bg_mode_radio, bg_color, bg_img,
-            alpha_chk, alpha_fg_sl, alpha_bg_sl, alpha_er_sl, post_chk,
-        ],
-        outputs=[out_transparent, out_composed, out_mask, status_box],
-        api_name="remove_background",
+        inputs=input_img,
+        outputs=output_img
     )
+    
+    # Connect the clear button to the output as well
+    clear_btn.add(components=[output_img])
 
-
-# ─────────────────────────────────────────────
-# Launch
-# ─────────────────────────────────────────────
+# Launch the app (Hugging Face Spaces automatically handles the port and server)
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-    )
+    demo.launch()