Claude remodel

README.md

@@ -13,6 +13,13 @@ preload_from_hub:
 
 # BEN2 Background Remover
 
-Upload an image and download a transparent PNG. Edge refinement is optional.
+Upload, paste, or capture an image and get a clean cut-out. Compare the result
+against the original with a before/after slider, optionally refine hair and fine
+edges, and either keep the background transparent or drop the subject onto a
+solid color. Download the result as a PNG.
 
 Select **ZeroGPU** in the Space hardware settings after creating the Space.
+
+To add one-click examples, create an `examples/` folder next to `app.py` and
+drop image files (`.png`, `.jpg`, `.jpeg`, `.webp`, `.bmp`) into it — they'll
+appear automatically. No folder means no examples, and nothing breaks.

app.py

@@ -3,12 +3,20 @@ from __future__ import annotations
 # ZeroGPU must patch PyTorch before anything else imports it.
 import spaces
 
+import os
+import tempfile
+
 import gradio as gr
+import numpy as np
 import torch
 from ben2 import BEN_Base
 from PIL import Image, ImageOps
 
 MODEL_ID = "PramaLLC/BEN2"
+MODEL_URL = "https://huggingface.co/PramaLLC/BEN2"
+
+# Guard rails: the model happily eats memory on huge inputs, and edge
+# refinement is far more expensive, so it gets a tighter ceiling.
 MAX_IMAGE_PIXELS = 40_000_000
 MAX_REFINEMENT_PIXELS = 12_000_000
 
@@ -17,7 +25,81 @@ model.to("cuda").eval()
 model.requires_grad_(False)
 
 
-def gpu_duration(_image: Image.Image | None, refine_foreground: bool) -> int:
+# --------------------------------------------------------------------------- #
+# Background options
+# --------------------------------------------------------------------------- #
+BG_TRANSPARENT = "Transparent"
+BG_WHITE = "White"
+BG_BLACK = "Black"
+BG_CUSTOM = "Custom color"
+BG_CHOICES = [BG_TRANSPARENT, BG_WHITE, BG_BLACK, BG_CUSTOM]
+
+_SOLID_RGB = {BG_WHITE: (255, 255, 255), BG_BLACK: (0, 0, 0)}
+_DEFAULT_RGB = (255, 255, 255)
+
+
+def _parse_color(value: str | None) -> tuple[int, int, int]:
+    """Turn a ColorPicker value ('#rrggbb', '#rgb', 'rgb(...)', 'rgba(...)') into RGB."""
+    if not value:
+        return _DEFAULT_RGB
+    v = value.strip()
+    if v.startswith("#"):
+        v = v[1:]
+        if len(v) == 3:
+            v = "".join(c * 2 for c in v)
+        try:
+            return int(v[0:2], 16), int(v[2:4], 16), int(v[4:6], 16)
+        except (ValueError, IndexError):
+            return _DEFAULT_RGB
+    if v.lower().startswith(("rgb(", "rgba(")):
+        try:
+            parts = v[v.index("(") + 1 : v.index(")")].split(",")
+            rgb = (int(round(float(p))) for p in parts[:3])
+            r, g, b = (max(0, min(255, n)) for n in rgb)
+            return r, g, b
+        except (ValueError, IndexError):
+            return _DEFAULT_RGB
+    return _DEFAULT_RGB
+
+
+def _checkerboard(size: tuple[int, int], square: int = 16) -> Image.Image:
+    """A transparency checkerboard sized to the image, so cut-outs read clearly."""
+    w, h = size
+    rows = (np.arange(h) // square)[:, None]
+    cols = (np.arange(w) // square)[None, :]
+    dark = ((rows + cols) & 1).astype(bool)
+    board = np.empty((h, w, 3), dtype=np.uint8)
+    board[~dark] = (245, 245, 245)
+    board[dark] = (208, 208, 208)
+    return Image.fromarray(board, mode="RGB")
+
+
+def _flatten(cutout: Image.Image, background: Image.Image) -> Image.Image:
+    """Alpha-composite an RGBA cut-out over an opaque RGB background."""
+    flat = background.convert("RGB").copy()
+    flat.paste(cutout, mask=cutout.split()[-1])
+    return flat
+
+
+def _compose(cutout: Image.Image, bg_choice: str, custom_color: str | None):
+    """
+    From the RGBA cut-out, build both:
+      - download: exactly what the user asked for (transparent RGBA, or flattened RGB)
+      - preview:  an always-visible RGB image for the before/after slider
+    """
+    if bg_choice == BG_TRANSPARENT:
+        preview = _flatten(cutout, _checkerboard(cutout.size))
+        return cutout, preview  # RGBA download, checkerboard preview
+
+    color = _SOLID_RGB.get(bg_choice) or _parse_color(custom_color)
+    flat = _flatten(cutout, Image.new("RGB", cutout.size, color))
+    return flat, flat  # RGB download == preview
+
+
+# --------------------------------------------------------------------------- #
+# Inference
+# --------------------------------------------------------------------------- #
+def gpu_duration(image, refine_foreground, bg_choice, custom_color) -> int:
     return 45 if refine_foreground else 30
 
 
@@ -25,7 +107,9 @@ def gpu_duration(_image: Image.Image | None, refine_foreground: bool) -> int:
 def remove_background(
     image: Image.Image | None,
     refine_foreground: bool,
-) -> Image.Image:
+    bg_choice: str,
+    custom_color: str | None,
+):
     if image is None:
         raise gr.Error("Upload, paste, or capture an image first.")
 
@@ -45,46 +129,155 @@ def remove_background(
             image.copy(),
             refine_foreground=bool(refine_foreground),
         )
+    cutout = foreground.convert("RGBA")
+
+    download_img, preview = _compose(cutout, bg_choice, custom_color)
 
-    return foreground.convert("RGBA")
-
-
-demo = gr.Interface(
-    fn=remove_background,
-    inputs=[
-        gr.Image(
-            type="pil",
-            image_mode="RGB",
-            sources=["upload", "clipboard", "webcam"],
-            label="Input image",
-        ),
-        gr.Checkbox(
-            value=False,
-            label="Refine edges",
-            info="Slower, but sometimes improves hair and fine boundaries.",
-        ),
-    ],
-    outputs=gr.Image(
-        type="pil",
-        image_mode="RGBA",
-        format="png",
-        label="Transparent PNG",
-        buttons=["download", "fullscreen"],
-    ),
-    title="BEN2 Background Remover",
-    description="Upload an image, remove its background, and download the transparent PNG.",
-    submit_btn="Remove background",
-    clear_btn="Clear",
-    flagging_mode="never",
-    api_name="remove_background",
-    concurrency_limit=1,
-    time_limit=60,
-    fill_width=True,
+    out_path = os.path.join(tempfile.mkdtemp(), "background-removed.png")
+    download_img.save(out_path)
+
+    return (image, preview), gr.DownloadButton(value=out_path, visible=True)
+
+
+# --------------------------------------------------------------------------- #
+# UI
+# --------------------------------------------------------------------------- #
+CSS = """
+#app-wrap { max-width: 1140px; margin: 0 auto; }
+#app-header { text-align: center; padding: 10px 0 2px; }
+#app-header h1 {
+  margin: 0; font-size: 1.9rem; font-weight: 700; letter-spacing: -0.02em;
+}
+#app-header p {
+  margin: 6px auto 0; max-width: 640px;
+  color: var(--body-text-color-subdued); font-size: 0.98rem; line-height: 1.5;
+}
+#app-header a { color: var(--link-text-color); text-decoration: none; }
+#app-header a:hover { text-decoration: underline; }
+/* soft transparency checkerboard behind the result, visible in any letterboxing */
+.checker-bg .image-frame, .checker-bg .empty {
+  background-color: #fbfbfb;
+  background-image:
+    linear-gradient(45deg, #e6e6e6 25%, transparent 25%),
+    linear-gradient(-45deg, #e6e6e6 25%, transparent 25%),
+    linear-gradient(45deg, transparent 75%, #e6e6e6 75%),
+    linear-gradient(-45deg, transparent 75%, #e6e6e6 75%);
+  background-size: 22px 22px;
+  background-position: 0 0, 0 11px, 11px -11px, -11px 0;
+}
+"""
+
+HEADER = f"""
+<div id="app-header">
+  <h1>BEN2 Background Remover</h1>
+  <p>Cut the background out of any photo and download a clean transparent PNG —
+  or drop your subject onto a solid color. Powered by
+  <a href="{MODEL_URL}" target="_blank" rel="noopener">BEN2</a>.</p>
+</div>
+"""
+
+# Optional: drop image files into an `examples/` folder next to this file and
+# they'll appear as one-click examples. No folder -> no examples, no error.
+_EXT = (".png", ".jpg", ".jpeg", ".webp", ".bmp")
+_EX_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples")
+EXAMPLE_FILES = (
+    sorted(
+        os.path.join(_EX_DIR, f)
+        for f in os.listdir(_EX_DIR)
+        if f.lower().endswith(_EXT)
+    )
+    if os.path.isdir(_EX_DIR)
+    else []
 )
 
+
+with gr.Blocks(title="BEN2 Background Remover", fill_width=True) as demo:
+    with gr.Column(elem_id="app-wrap"):
+        gr.HTML(HEADER)
+
+        with gr.Row(equal_height=False):
+            with gr.Column(scale=1):
+                input_image = gr.Image(
+                    type="pil",
+                    image_mode="RGB",
+                    sources=["upload", "clipboard", "webcam"],
+                    label="Input",
+                    placeholder="Drop an image here, paste from the clipboard, "
+                    "or capture from your webcam",
+                    height=400,
+                )
+                with gr.Group():
+                    refine = gr.Checkbox(
+                        value=False,
+                        label="Refine edges",
+                        info="Slower, but recovers hair and fine boundaries.",
+                    )
+                    bg_choice = gr.Radio(
+                        BG_CHOICES,
+                        value=BG_TRANSPARENT,
+                        label="Background",
+                        info="Keep it transparent, or fill it with a color.",
+                    )
+                    custom_color = gr.ColorPicker(
+                        value="#00b140",  # chroma-key green
+                        label="Custom color",
+                        visible=False,
+                    )
+                with gr.Row():
+                    clear_btn = gr.ClearButton(value="Clear")
+                    run_btn = gr.Button("Remove background", variant="primary")
+
+            with gr.Column(scale=1):
+                output_slider = gr.ImageSlider(
+                    label="Before / after",
+                    type="pil",
+                    image_mode="RGB",
+                    slider_position=50,
+                    buttons=["fullscreen"],
+                    interactive=False,
+                    height=400,
+                    elem_classes=["checker-bg"],
+                )
+                download_btn = gr.DownloadButton(
+                    "Download PNG",
+                    variant="primary",
+                    visible=False,
+                )
+
+        if EXAMPLE_FILES:
+            gr.Examples(
+                examples=[[f] for f in EXAMPLE_FILES],
+                inputs=[input_image],
+                label="Try an example",
+            )
+
+    # Reveal the color picker only when a custom color is selected.
+    bg_choice.change(
+        lambda choice: gr.ColorPicker(visible=choice == BG_CUSTOM),
+        inputs=bg_choice,
+        outputs=custom_color,
+        show_progress="hidden",
+    )
+
+    clear_btn.add([input_image, output_slider])
+    clear_btn.click(
+        lambda: gr.DownloadButton(visible=False),
+        outputs=download_btn,
+        show_progress="hidden",
+    )
+
+    run_btn.click(
+        remove_background,
+        inputs=[input_image, refine, bg_choice, custom_color],
+        outputs=[output_slider, download_btn],
+        api_name="remove_background",
+        concurrency_limit=1,
+    )
+
+
 if __name__ == "__main__":
     demo.queue(
         max_size=20,
         api_open=False,
         default_concurrency_limit=1,
-    ).launch()
+    ).launch(theme=gr.themes.Ocean(), css=CSS)

requirements.txt

@@ -1,5 +1,7 @@
 torch==2.11.0
 torchvision==0.26.0
+numpy
+pillow
 opencv-python-headless
 safetensors
 git+https://github.com/PramaLLC/BEN2.git