Create app.py

app.py

@@ -0,0 +1,469 @@
+import os
+import io
+import time
+import glob
+import math
+from dataclasses import dataclass
+from typing import Dict, Optional, Tuple, List
+
+import gradio as gr
+import spaces
+
+import numpy as np
+from PIL import Image
+
+import torch
+from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
+
+# Slider component (same one BRIA Space uses)
+from gradio_imageslider import ImageSlider
+
+# InSPyReNet wrapper (same approach as your sample Space)
+from transparent_background import Remover
+
+# rembg (U2Net + IS-Net via ONNX)
+from rembg import new_session, remove as rembg_remove
+
+
+# ----------------------------
+# Utilities
+# ----------------------------
+
+def pil_to_rgb(pil: Image.Image) -> Image.Image:
+    if pil.mode != "RGB":
+        return pil.convert("RGB")
+    return pil
+
+
+def ensure_rgba(pil: Image.Image) -> Image.Image:
+    if pil.mode != "RGBA":
+        return pil.convert("RGBA")
+    return pil
+
+
+def make_checkerboard(w: int, h: int, block: int = 16) -> Image.Image:
+    # Neutral checkerboard
+    cols = int(math.ceil(w / block))
+    rows = int(math.ceil(h / block))
+    board = np.zeros((rows * block, cols * block, 3), dtype=np.uint8)
+
+    c1 = np.array([235, 235, 235], dtype=np.uint8)
+    c2 = np.array([200, 200, 200], dtype=np.uint8)
+
+    for r in range(rows):
+        for c in range(cols):
+            color = c1 if (r + c) % 2 == 0 else c2
+            board[r * block:(r + 1) * block, c * block:(c + 1) * block] = color
+
+    board = board[:h, :w, :]
+    return Image.fromarray(board, mode="RGB")
+
+
+def rgba_on_checkerboard(rgba: Image.Image) -> Image.Image:
+    rgba = ensure_rgba(rgba)
+    w, h = rgba.size
+    bg = make_checkerboard(w, h)
+    comp = Image.alpha_composite(bg.convert("RGBA"), rgba)
+    return comp.convert("RGB")
+
+
+def save_temp_png(rgba: Image.Image, out_dir: str = "output_images") -> str:
+    os.makedirs(out_dir, exist_ok=True)
+    path = os.path.join(out_dir, "no_bg.png")
+    ensure_rgba(rgba).save(path, format="PNG")
+    return path
+
+
+def now_ms() -> float:
+    return time.perf_counter() * 1000.0
+
+
+@dataclass
+class Timing:
+    preprocess_ms: float
+    inference_ms: float
+    postprocess_ms: float
+    total_ms: float
+
+    def to_text(self) -> str:
+        return (
+            f"preprocess:  {self.preprocess_ms:.2f} ms\n"
+            f"inference:   {self.inference_ms:.2f} ms\n"
+            f"postprocess: {self.postprocess_ms:.2f} ms\n"
+            f"TOTAL:       {self.total_ms:.2f} ms"
+        )
+
+
+# ----------------------------
+# Model Manager
+# ----------------------------
+
+class ModelManager:
+    """
+    Loads and runs:
+      1) InSPyReNet via transparent_background.Remover()
+      2) BiRefNet via AutoModelForImageSegmentation("ZhengPeng7/BiRefNet", trust_remote_code=True)
+      3) U2Net via rembg (onnxruntime; can use CUDA provider if available)
+      4) BRIA RMBG 2.0 via AutoModelForImageSegmentation("briaai/RMBG-2.0", trust_remote_code=True)
+      5) IS-Net (isnet-general-use) via rembg
+    """
+    def __init__(self):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self._inspy: Optional[Remover] = None
+
+        self._torch_models: Dict[str, AutoModelForImageSegmentation] = {}
+        self._torch_model_on_gpu: Optional[str] = None
+
+        # rembg sessions
+        self._rembg_sessions: Dict[str, object] = {}
+
+        # Common transforms for BiRefNet / BRIA RMBG inference
+        self._tf_1024 = transforms.Compose([
+            transforms.Resize((1024, 1024)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ])
+
+        # Try to set matmul precision nicely
+        try:
+            torch.set_float32_matmul_precision("high")
+        except Exception:
+            pass
+
+    def _maybe_sync(self):
+        if self.device == "cuda":
+            torch.cuda.synchronize()
+
+    def _load_inspy(self) -> Remover:
+        if self._inspy is None:
+            # jit=False like your sample
+            self._inspy = Remover(jit=False)
+        return self._inspy
+
+    def _offload_torch_models_from_gpu(self, keep_name: str):
+        if self.device != "cuda":
+            return
+        if self._torch_model_on_gpu and self._torch_model_on_gpu != keep_name:
+            prev = self._torch_models.get(self._torch_model_on_gpu)
+            if prev is not None:
+                prev.to("cpu")
+            self._torch_model_on_gpu = None
+            torch.cuda.empty_cache()
+
+    def _load_torch_model(self, key: str) -> AutoModelForImageSegmentation:
+        """
+        key in {"birefnet", "bria_rmbg_2"}
+        """
+        if key in self._torch_models:
+            return self._torch_models[key]
+
+        if key == "birefnet":
+            model_id = "ZhengPeng7/BiRefNet"
+        elif key == "bria_rmbg_2":
+            model_id = "briaai/RMBG-2.0"
+        else:
+            raise ValueError(f"Unknown torch model key: {key}")
+
+        m = AutoModelForImageSegmentation.from_pretrained(model_id, trust_remote_code=True)
+        m.eval()
+        # Keep on CPU initially; move to GPU on-demand to avoid T4 OOM.
+        m.to("cpu")
+        self._torch_models[key] = m
+        return m
+
+    def _get_rembg_session(self, name: str):
+        """
+        name: "u2net" or "isnet-general-use"
+        """
+        if name in self._rembg_sessions:
+            return self._rembg_sessions[name]
+
+        # Prefer CUDA provider if onnxruntime-gpu is installed; otherwise CPU works.
+        # rembg will pass this into onnxruntime internally.
+        providers = None
+        try:
+            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        except Exception:
+            providers = None
+
+        sess = new_session(name, providers=providers) if providers else new_session(name)
+        self._rembg_sessions[name] = sess
+        return sess
+
+    def _run_torch_alpha_model(self, model_key: str, image_rgb: Image.Image) -> Image.Image:
+        """
+        Runs a torch segmentation model that returns a single-channel mask (alpha matte-ish).
+        Returns RGBA (with alpha).
+        """
+        m = self._load_torch_model(model_key)
+
+        # Put model on GPU for inference if possible
+        if self.device == "cuda":
+            self._offload_torch_models_from_gpu(keep_name=model_key)
+            if self._torch_model_on_gpu != model_key:
+                m.to("cuda")
+                self._torch_model_on_gpu = model_key
+
+        image_rgb = pil_to_rgb(image_rgb)
+        orig_size = image_rgb.size
+
+        x = self._tf_1024(image_rgb).unsqueeze(0)
+        x = x.to(self.device)
+
+        with torch.inference_mode():
+            if self.device == "cuda":
+                with torch.autocast(device_type="cuda", dtype=torch.float16):
+                    preds = m(x)[-1].sigmoid()
+            else:
+                preds = m(x)[-1].sigmoid()
+
+        # Convert prediction to PIL alpha channel
+        pred = preds[0].squeeze().detach().float().cpu()
+        alpha = transforms.ToPILImage()(pred).resize(orig_size)
+
+        out = image_rgb.convert("RGBA")
+        out.putalpha(alpha)
+        return out
+
+    def run(self, model_name: str, input_image: Image.Image) -> Tuple[Image.Image, Timing]:
+        """
+        Returns (output_rgba, timing).
+        """
+        if input_image is None:
+            raise ValueError("No input image")
+
+        t0 = now_ms()
+
+        # --- preprocess ---
+        pre0 = now_ms()
+        img_rgb = pil_to_rgb(input_image)
+        pre1 = now_ms()
+
+        # --- inference ---
+        inf0 = now_ms()
+        if model_name == "InSPyReNet":
+            remover = self._load_inspy()
+            # The library returns various modes; we want alpha mask and apply ourselves for consistent output
+            mask = remover.process(input_image, type="map")
+            if isinstance(mask, Image.Image):
+                mask = mask.convert("L")
+            else:
+                mask = Image.fromarray((mask * 255).astype(np.uint8), mode="L")
+
+            out = img_rgb.convert("RGBA")
+            out.putalpha(mask)
+
+        elif model_name == "BiRefNet":
+            out = self._run_torch_alpha_model("birefnet", img_rgb)
+
+        elif model_name == "U2Net":
+            sess = self._get_rembg_session("u2net")
+            # rembg returns bytes (PNG RGBA)
+            out_bytes = rembg_remove(img_rgb, session=sess)
+            out = Image.open(io.BytesIO(out_bytes)).convert("RGBA")
+
+        elif model_name == "BRIA RMBG 2.0":
+            out = self._run_torch_alpha_model("bria_rmbg_2", img_rgb)
+
+        elif model_name == "IS-Net":
+            sess = self._get_rembg_session("isnet-general-use")
+            out_bytes = rembg_remove(img_rgb, session=sess)
+            out = Image.open(io.BytesIO(out_bytes)).convert("RGBA")
+
+        else:
+            raise ValueError(f"Unknown model: {model_name}")
+
+        # Make sure GPU timing is accurate
+        self._maybe_sync()
+        inf1 = now_ms()
+
+        # --- postprocess ---
+        post0 = now_ms()
+        out = ensure_rgba(out)
+        post1 = now_ms()
+
+        t1 = now_ms()
+
+        timing = Timing(
+            preprocess_ms=pre1 - pre0,
+            inference_ms=inf1 - inf0,
+            postprocess_ms=post1 - post0,
+            total_ms=t1 - t0,
+        )
+        return out, timing
+
+
+MANAGER = ModelManager()
+
+MODEL_CHOICES = [
+    "InSPyReNet",
+    "BiRefNet",
+    "U2Net",
+    "BRIA RMBG 2.0",
+    "IS-Net",
+]
+
+
+# ----------------------------
+# Gradio handlers
+# ----------------------------
+
+@spaces.GPU
+def run_single(model_name: str, image: Image.Image):
+    if image is None:
+        return None, None, "Upload an image first.", None
+
+    # Warmup-ish for fairer timing (tiny; avoids huge overhead in UI)
+    # Note: real benchmark tab does proper warmups.
+    out_rgba, timing = MANAGER.run(model_name, image)
+
+    # Slider wants (processed, original) or (after, before) depending on component;
+    # we’ll show: left=original, right=on-checkerboard preview of transparent output.
+    preview = rgba_on_checkerboard(out_rgba)
+
+    out_path = save_temp_png(out_rgba)
+    return (image, preview), out_rgba, timing.to_text(), out_path
+
+
+def list_bench_images() -> List[str]:
+    # Put your 10–15 images under bench/
+    exts = ("*.jpg", "*.jpeg", "*.png", "*.webp")
+    files = []
+    for e in exts:
+        files += glob.glob(os.path.join("bench", e))
+    files = sorted(files)
+
+    # Fallback to repo-root examples like your sample Space
+    if not files:
+        fallback = []
+        for f in ["1.jpg", "2.jpg", "3.png", "4.webp"]:
+            if os.path.exists(f):
+                fallback.append(f)
+        files = fallback
+    return files
+
+
+@spaces.GPU
+def run_benchmark(model_name: str, repeats: int = 1):
+    files = list_bench_images()
+    if not files:
+        return gr.Dataframe(value=[]), "No benchmark images found. Add 10–15 images under bench/."
+
+    # Warmup: 2 runs on first image (not timed)
+    warm_img = Image.open(files[0]).convert("RGB")
+    for _ in range(2):
+        _ = MANAGER.run(model_name, warm_img)
+
+    rows = []
+    total_ms = 0.0
+    n_images = 0
+
+    for f in files:
+        img = Image.open(f).convert("RGB")
+        for r in range(repeats):
+            out, timing = MANAGER.run(model_name, img)
+            rows.append({
+                "file": os.path.basename(f),
+                "repeat": r + 1,
+                "total_ms": round(timing.total_ms, 2),
+                "inference_ms": round(timing.inference_ms, 2),
+            })
+            total_ms += timing.total_ms
+            n_images += 1
+
+    avg_ms = total_ms / max(1, n_images)
+    ips = 1000.0 / avg_ms if avg_ms > 0 else 0.0
+
+    summary = (
+        f"Model: {model_name}\n"
+        f"Images: {len(files)} (repeats={repeats}) => runs={n_images}\n"
+        f"Avg total: {avg_ms:.2f} ms\n"
+        f"Estimated throughput: {ips:.2f} images/sec\n"
+        f"Device: {'T4 GPU' if torch.cuda.is_available() else 'CPU'}"
+    )
+
+    df = gr.Dataframe(
+        headers=["file", "repeat", "total_ms", "inference_ms"],
+        value=[[r["file"], r["repeat"], r["total_ms"], r["inference_ms"]] for r in rows],
+        datatype=["str", "number", "number", "number"],
+        interactive=False
+    )
+    return df, summary
+
+
+# ----------------------------
+# UI
+# ----------------------------
+
+with gr.Blocks(title="Background Removal Benchmark (T4)") as demo:
+    gr.Markdown(
+        """
+# Background Removal Benchmark (T4)
+
+Benchmarked models:
+1) InSPyReNet  
+2) BiRefNet  
+3) U2Net  
+4) BRIA RMBG 2.0  
+5) IS-Net (isnet-general-use)
+
+**Notes**
+- Output download is a true transparent PNG (RGBA).
+- The slider preview composites the transparent result over a checkerboard for visibility.
+- For the benchmark tab, add **10–15 images** under `bench/` in your Space repo.
+"""
+    )
+
+    with gr.Tab("Try single image"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                inp = gr.Image(type="pil", label="Upload image", height=420)
+                model = gr.Dropdown(choices=MODEL_CHOICES, value="InSPyReNet", label="Model")
+                run_btn = gr.Button("Run", variant="primary")
+            with gr.Column(scale=2):
+                slider = ImageSlider(label="Before / After", type="pil")
+                out_img = gr.Image(type="pil", label="Output (RGBA)", height=420)
+                timing_box = gr.Textbox(label="Timing", lines=5)
+                out_file = gr.File(label="Download PNG (transparent)")
+
+        run_btn.click(
+            fn=run_single,
+            inputs=[model, inp],
+            outputs=[slider, out_img, timing_box, out_file]
+        )
+
+    with gr.Tab("Benchmark (throughput estimate)"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                bench_model = gr.Dropdown(choices=MODEL_CHOICES, value="InSPyReNet", label="Model")
+                repeats = gr.Slider(1, 5, value=1, step=1, label="Repeats per image (higher = more stable averages)")
+                bench_btn = gr.Button("Run benchmark", variant="primary")
+            with gr.Column(scale=2):
+                bench_table = gr.Dataframe(
+                    headers=["file", "repeat", "total_ms", "inference_ms"],
+                    datatype=["str", "number", "number", "number"],
+                    interactive=False
+                )
+                bench_summary = gr.Textbox(label="Summary", lines=6)
+
+        bench_btn.click(
+            fn=run_benchmark,
+            inputs=[bench_model, repeats],
+            outputs=[bench_table, bench_summary]
+        )
+
+    # Examples (optional) — if these files exist, they show up like your sample Space
+    example_files = []
+    for f in ["1.jpg", "2.jpg", "3.png", "4.webp"]:
+        if os.path.exists(f):
+            example_files.append([f, "InSPyReNet"])
+    if example_files:
+        gr.Examples(
+            examples=example_files,
+            inputs=[inp, model],
+            label="Examples"
+        )
+
+if __name__ == "__main__":
+    demo.launch(show_error=True)