app.py

"""AssetsMaker UI — drag a PNG, get a Unity-ready atlas.

Run:
    .venv/bin/python app.py

Opens http://127.0.0.1:7860 in your browser.
"""

from __future__ import annotations

import shutil
import tempfile
from pathlib import Path
from typing import Tuple

import gradio as gr
from PIL import Image, ImageDraw

from atlas import detect_sprites, pack_logical_grid, write_atlas
from bg_remover import RemoveOptions, remove_background


MODEL_CHOICES = [
    "birefnet-general",
    "birefnet-general-lite",
    "birefnet-portrait",
    "birefnet-dis",
    "isnet-general-use",
    "u2net",
]


def _composite_on_magenta(rgba: Image.Image) -> Image.Image:
    bg = Image.new("RGBA", rgba.size, (255, 0, 255, 255))
    return Image.alpha_composite(bg, rgba.convert("RGBA")).convert("RGB")


def _draw_debug_overlay(rgba: Image.Image, sprites) -> Image.Image:
    composite = _composite_on_magenta(rgba).convert("RGBA")
    draw = ImageDraw.Draw(composite)
    for s in sprites:
        x0, y0 = s.source_x, s.source_y
        x1, y1 = x0 + s.width, y0 + s.height
        draw.rectangle([x0, y0, x1 - 1, y1 - 1], outline=(0, 255, 0, 255), width=2)
    return composite.convert("RGB")


def process(
    input_image: Image.Image,
    model: str,
    edge_sharpen: bool,
    sharpen_strength: float,
    decontaminate: bool,
    detect_holes: bool,
    bg_tolerance: float,
    detect_method: str,
    separation_distance: int,
    grid_size_str: str,
    alpha_threshold: int,
    shadow_padding: int,
    min_area: int,
    atlas_padding: int,
    atlas_max_width: int,
    atlas_pow2: bool,
    size_tolerance: float,
    progress=gr.Progress(),
) -> Tuple[Image.Image, Image.Image, Image.Image, Image.Image, str, list]:
    if input_image is None:
        raise gr.Error("Drop an image first.")

    progress(0.05, desc="Removing background…")
    cutout = remove_background(
        input_image,
        RemoveOptions(
            model=model,
            edge_sharpen=edge_sharpen,
            edge_sharpen_strength=sharpen_strength,
            decontaminate=decontaminate,
            detect_holes=detect_holes,
            bg_color_tolerance=bg_tolerance,
        ),
    )

    progress(0.55, desc="Detecting sprites…")
    grid_size = None
    if detect_method == "grid":
        try:
            gw, gh = (int(v) for v in grid_size_str.lower().split("x"))
            grid_size = (gw, gh)
        except Exception:
            raise gr.Error(f"Grid size must be 'WxH' (got {grid_size_str!r}).")

    sprites = detect_sprites(
        cutout,
        method=detect_method,
        core_alpha_threshold=alpha_threshold,
        shadow_padding=shadow_padding,
        min_area=min_area,
        separation_distance=separation_distance,
        grid_size=grid_size,
    )

    if not sprites:
        raise gr.Error(
            "No sprites detected. Try lowering Alpha threshold or Min area."
        )

    progress(0.80, desc="Packing atlas…")
    atlas = pack_logical_grid(
        sprites,
        padding=atlas_padding,
        max_width=int(atlas_max_width),
        size_tolerance=size_tolerance,
        pow2=atlas_pow2,
    )

    progress(0.92, desc="Writing files…")
    out_dir = Path(tempfile.mkdtemp(prefix="assetsmaker_"))
    cutout_path = out_dir / "cutout.png"
    cutout.save(cutout_path, format="PNG", optimize=True)
    paths = write_atlas(atlas, out_dir, name="atlas")

    cutout_preview = _composite_on_magenta(cutout)
    overlay_preview = _draw_debug_overlay(cutout, sprites)
    atlas_preview = atlas.image  # RGBA, Gradio shows checkerboard

    summary = (
        f"**{len(sprites)} sprites detected** · "
        f"atlas **{atlas.width} × {atlas.height}** · "
        f"output: `{out_dir}`"
    )

    downloads = [str(cutout_path), str(paths["png"]), str(paths["json"])]
    return cutout_preview, overlay_preview, atlas_preview, atlas.image, summary, downloads


with gr.Blocks(title="AssetsMaker") as demo:
    gr.Markdown(
        "# AssetsMaker\n"
        "Drop a sprite-sheet PNG → sharp background removal → Unity-ready atlas."
    )

    with gr.Row():
        with gr.Column(scale=1, min_width=320):
            input_img = gr.Image(
                label="Input image",
                type="pil",
                image_mode="RGBA",
                height=320,
            )
            run_btn = gr.Button("Process", variant="primary", size="lg")

            with gr.Accordion("Background removal", open=False):
                model = gr.Dropdown(
                    choices=MODEL_CHOICES, value="birefnet-general",
                    label="Model",
                    info="BiRefNet-general = SOTA quality. -lite is ~3× faster.",
                )
                edge_sharpen = gr.Checkbox(value=True, label="Edge sharpening")
                sharpen_strength = gr.Slider(
                    0.0, 1.5, value=0.6, step=0.1,
                    label="Sharpen strength",
                )
                decontaminate = gr.Checkbox(
                    value=True, label="Foreground decontamination",
                    info="Removes bg-color halo around partial-alpha edges.",
                )
                detect_holes = gr.Checkbox(
                    value=False, label="Chroma-key bg color (interior holes)",
                    info="Aggressive: harms objects of similar color.",
                )
                bg_tolerance = gr.Slider(
                    0, 80, value=18, step=1,
                    label="BG color tolerance",
                )

            with gr.Accordion("Sprite detection", open=True):
                detect_method = gr.Radio(
                    choices=["components", "watershed", "grid"],
                    value="components",
                    label="Method",
                    info=(
                        "components = isolated sprites only. "
                        "watershed = splits touching sprites. "
                        "grid = fixed-cell tilemap."
                    ),
                )
                separation_distance = gr.Slider(
                    4, 200, value=24, step=2,
                    label="Watershed separation distance (px)",
                    info="≈ expected sprite half-size. Smaller = more splits.",
                )
                grid_size_str = gr.Textbox(
                    value="64x64", label="Grid cell size (WxH)",
                    info="Used only when method = grid.",
                )
                alpha_threshold = gr.Slider(
                    1, 200, value=80, step=1,
                    label="Alpha threshold",
                    info="Higher = avoid bridging sprites via shadow alpha.",
                )
                shadow_padding = gr.Slider(
                    0, 40, value=6, step=1,
                    label="Shadow padding (px)",
                    info="Pixels added around each bbox to keep drop shadows.",
                )
                min_area = gr.Slider(
                    50, 5000, value=400, step=10,
                    label="Min sprite area (px)",
                )

            with gr.Accordion("Atlas packing", open=False):
                atlas_padding = gr.Slider(
                    0, 16, value=2, step=1,
                    label="Padding between sprites",
                )
                atlas_max_width = gr.Slider(
                    512, 8192, value=2048, step=256,
                    label="Max atlas width",
                )
                atlas_pow2 = gr.Checkbox(
                    value=True, label="Power-of-2 atlas size",
                    info="Unity-friendly for GPU compression.",
                )
                size_tolerance = gr.Slider(
                    0.0, 0.4, value=0.10, step=0.01,
                    label="Same-size grouping tolerance",
                )

        with gr.Column(scale=2):
            summary = gr.Markdown()
            with gr.Tabs():
                with gr.Tab("Cutout"):
                    cutout_view = gr.Image(
                        label="On magenta — to spot any halo",
                        type="pil", height=520,
                    )
                with gr.Tab("Detected sprites"):
                    overlay_view = gr.Image(
                        label="Green = bbox of each detected sprite",
                        type="pil", height=520,
                    )
                with gr.Tab("Atlas"):
                    atlas_view = gr.Image(
                        label="Final atlas (RGBA)",
                        type="pil", height=520,
                    )
                with gr.Tab("Atlas (RGBA)"):
                    atlas_view_alpha = gr.Image(
                        label="Direct alpha render",
                        type="pil", image_mode="RGBA", height=520,
                    )
            downloads = gr.File(
                label="Download — cutout.png · atlas.png · atlas.json",
                file_count="multiple",
            )

    run_btn.click(
        process,
        inputs=[
            input_img, model, edge_sharpen, sharpen_strength,
            decontaminate, detect_holes, bg_tolerance,
            detect_method, separation_distance, grid_size_str,
            alpha_threshold, shadow_padding, min_area,
            atlas_padding, atlas_max_width, atlas_pow2, size_tolerance,
        ],
        outputs=[cutout_view, overlay_view, atlas_view, atlas_view_alpha, summary, downloads],
    )


if __name__ == "__main__":
    demo.launch(
        server_name="127.0.0.1",
        server_port=7860,
        inbrowser=True,
        theme=gr.themes.Soft(primary_hue="violet"),
        css="footer {visibility: hidden}",
    )