app.py

"""ViTMatte demo using TrimapEditor for trimap input."""

from __future__ import annotations

import base64
import json
import os
from io import BytesIO
from pathlib import Path

import gradio as gr
import numpy as np
import PIL.Image
import spaces
import torch
from transformers import VitMatteForImageMatting, VitMatteImageProcessor
from trimap_editor import TrimapEditor

DESCRIPTION = """\
# [ViTMatte](https://github.com/hustvl/ViTMatte)

Image matting with Vision Transformers — accurately extract the foreground
from an image, even tricky areas like hair and fur!

1. **Upload** an image (or click an example below).
2. **Draw** foreground (green) and unknown (blue) regions. Press **?** for shortcuts.
3. Click **Run** to generate the alpha matte.
"""

_ASSETS_DIR = Path(__file__).parent / "assets"

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1500"))
MODEL_ID = os.getenv("MODEL_ID", "hustvl/vitmatte-small-distinctions-646")

processor = VitMatteImageProcessor.from_pretrained(MODEL_ID)
model = VitMatteForImageMatting.from_pretrained(MODEL_ID).to(device)


def _resize_on_upload(value: str | None) -> PIL.Image.Image:
    """Downscale the image on upload if it exceeds MAX_IMAGE_SIZE.

    Only fires for fresh uploads (no trimapBase64 yet).
    Returns gr.skip() when no resize is needed so the editor is untouched.
    """
    if not value:
        return gr.skip()
    d = json.loads(value)
    if "trimapBase64" in d:
        return gr.skip()
    w, h = d.get("width", 0), d.get("height", 0)
    if max(w, h) <= MAX_IMAGE_SIZE:
        return gr.skip()
    image_url = d.get("image", "")
    if not image_url:
        return gr.skip()
    # .change() receives raw postprocess() values which still have the
    # Gradio file-serving prefix (JS strips it only in commitValue()).
    image_path = image_url.removeprefix("/gradio_api/file=")
    image = PIL.Image.open(image_path).convert("RGB")
    scale = MAX_IMAGE_SIZE / max(w, h)
    new_w, new_h = int(w * scale), int(h * scale)
    gr.Info(f"Image resized from {w}x{h} to {new_w}x{new_h} (max {MAX_IMAGE_SIZE}px).")
    return image.resize((new_w, new_h))


def _parse_editor(value: str | None) -> tuple[PIL.Image.Image, PIL.Image.Image]:
    """Extract image and trimap from TrimapEditor JSON value."""
    if not value:
        raise gr.Error("Upload an image and draw a trimap first.")
    d = json.loads(value)

    # Image
    image_url = d.get("image", "")
    if not image_url:
        raise gr.Error("No image loaded.")
    image = PIL.Image.open(image_url).convert("RGB")

    # Trimap: prefer trimapBase64 (user-drawn), fall back to trimap URL (from example)
    if "trimapBase64" in d:
        b64 = d["trimapBase64"]
        if "," in b64:
            b64 = b64.split(",", 1)[1]
        trimap = PIL.Image.open(BytesIO(base64.b64decode(b64))).convert("L")
    elif "trimap" in d:
        trimap = PIL.Image.open(d["trimap"]).convert("L")
    else:
        raise gr.Error("Draw a trimap first (mark foreground and unknown regions).")

    return image, trimap


def _adjust_background(bg: PIL.Image.Image, target_size: tuple[int, int]) -> PIL.Image.Image:
    """Crop-resize background to match target dimensions."""
    tw, th = target_size
    bw, bh = bg.size
    scale = max(tw / bw, th / bh)
    bg = bg.resize((int(bw * scale), int(bh * scale)))
    left = (bg.width - tw) // 2
    top = (bg.height - th) // 2
    return bg.crop((left, top, left + tw, top + th))


def _replace_background(
    image: PIL.Image.Image, alpha: np.ndarray, bg: PIL.Image.Image | None
) -> PIL.Image.Image | None:
    if bg is None:
        return None
    bg = _adjust_background(bg.convert("RGB"), image.size)
    fg = np.array(image).astype(float) / 255
    bg_arr = np.array(bg).astype(float) / 255
    result = fg * alpha[:, :, None] + bg_arr * (1 - alpha[:, :, None])
    return PIL.Image.fromarray((result * 255).astype(np.uint8))


@spaces.GPU
@torch.inference_mode()
def run(
    editor_value: str | None,
    apply_bg: bool,
    background_image: PIL.Image.Image | None,
) -> tuple:
    image, trimap = _parse_editor(editor_value)

    pixel_values = processor(images=image, trimaps=trimap, return_tensors="pt").to(device).pixel_values
    out = model(pixel_values=pixel_values)
    alpha = out.alphas[0, 0].to("cpu").numpy()

    w, h = image.size
    alpha = alpha[:h, :w]

    foreground = np.array(image).astype(float) / 255 * alpha[:, :, None] + (1 - alpha[:, :, None])
    foreground = PIL.Image.fromarray((foreground * 255).astype(np.uint8))

    res_bg = _replace_background(image, alpha, background_image) if apply_bg else None

    return (
        (image, alpha),
        (image, foreground),
        (image, res_bg) if res_bg is not None else None,
    )


with gr.Blocks(css_paths="style.css") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column():
            editor = TrimapEditor(label="Image & Trimap")
            with gr.Group():
                apply_bg = gr.Checkbox(label="Replace background", value=False)
                bg_image = gr.Image(label="Background image", type="pil", visible=False)
            run_btn = gr.Button("Run", variant="primary")
        with gr.Column():
            out_alpha = gr.ImageSlider(label="Alpha")
            out_foreground = gr.ImageSlider(label="Foreground")
            out_bg = gr.ImageSlider(label="Background replacement", visible=False)

    inputs = [editor, apply_bg, bg_image]
    outputs = [out_alpha, out_foreground, out_bg]

    gr.Examples(
        examples=[
            [
                [str(_ASSETS_DIR / "retriever_rgb.png"), str(_ASSETS_DIR / "retriever_trimap.png")],
                False,
                None,
            ],
            [
                [str(_ASSETS_DIR / "bulb_rgb.png"), str(_ASSETS_DIR / "bulb_trimap.png")],
                True,
                str(_ASSETS_DIR / "new_bg.jpg"),
            ],
        ],
        inputs=inputs,
        outputs=outputs,
        fn=run,
        cache_examples=False,
    )

    editor.input(
        fn=_resize_on_upload,
        inputs=editor,
        outputs=editor,
        api_name=False,
    )
    apply_bg.change(
        fn=lambda checked: (gr.Image(visible=checked), gr.ImageSlider(visible=checked)),
        inputs=apply_bg,
        outputs=[bg_image, out_bg],
        api_name=False,
    )

    run_btn.click(fn=run, inputs=inputs, outputs=outputs)

if __name__ == "__main__":
    demo.launch()