app.py

"""
Gradio Space: remove background with rembg, crop to subject, normalize to 256×256.
"""

from __future__ import annotations

import os
import traceback
from typing import Optional, Tuple

import gradio as gr
import numpy as np
from PIL import Image, ImageEnhance, ImageOps
from rembg import remove


OUTPUT_SIZE = 256
# Subtle enhancement — tweak without blowing out edges on small subjects
CONTRAST_FACTOR = 1.08
SHARPNESS_FACTOR = 1.12
ALPHA_THRESHOLD = 8  # ignore nearly-transparent noise


def _alpha_bounding_box_rgba(img: Image.Image) -> Optional[Tuple[int, int, int, int]]:
    """Return (left, upper, right, lower) of non-transparent pixels, or None if empty."""
    if img.mode != "RGBA":
        img = img.convert("RGBA")
    alpha = np.array(img.split()[3])
    mask = alpha > ALPHA_THRESHOLD
    rows = np.any(mask, axis=1)
    cols = np.any(mask, axis=0)
    if not (np.any(rows) and np.any(cols)):
        return None
    y_indices = np.where(rows)[0]
    x_indices = np.where(cols)[0]
    return (
        int(x_indices[0]),
        int(y_indices[0]),
        int(x_indices[-1] + 1),
        int(y_indices[-1] + 1),
    )


def _resize_pad_square_rgba(img: Image.Image, size: int) -> Image.Image:
    """Scale uniformly to fit inside size×size, center on transparent canvas."""
    img = img.convert("RGBA")
    w, h = img.size
    if w == 0 or h == 0:
        return Image.new("RGBA", (size, size), (0, 0, 0, 0))
    scale = min(size / w, size / h)
    new_w = max(1, int(round(w * scale)))
    new_h = max(1, int(round(h * scale)))
    resized = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
    canvas = Image.new("RGBA", (size, size), (0, 0, 0, 0))
    ox = (size - new_w) // 2
    oy = (size - new_h) // 2
    canvas.paste(resized, (ox, oy), resized)
    return canvas


def _enhance_rgba(img: Image.Image, contrast: float, sharpness: float) -> Image.Image:
    """Apply contrast/sharpness to RGB channels only; preserve alpha."""
    img = img.convert("RGBA")
    r, g, b, a = img.split()
    rgb = Image.merge("RGB", (r, g, b))
    rgb = ImageEnhance.Contrast(rgb).enhance(contrast)
    rgb = ImageEnhance.Sharpness(rgb).enhance(sharpness)
    r2, g2, b2 = rgb.split()
    return Image.merge("RGBA", (r2, g2, b2, a))


def process_image(input_image: Optional[Image.Image]) -> Tuple[Optional[Image.Image], str]:
    """
    Pipeline: rembg → alpha bbox crop → 256×256 pad → mild enhance → PNG-ready RGBA.
    Returns (result_image, status_message).
    """
    if input_image is None:
        return None, "Please upload an image."

    try:
        # Normalize input (handles EXIF orientation, mode)
        pil = ImageOps.exif_transpose(input_image)
        if pil.mode not in ("RGB", "RGBA"):
            pil = pil.convert("RGBA") if pil.mode in ("P", "LA") else pil.convert("RGB")

        # rembg expects RGB or path/bytes; PIL RGB is fine
        if pil.mode == "RGBA":
            rgb = Image.new("RGB", pil.size, (255, 255, 255))
            rgb.paste(pil, mask=pil.split()[3])
            pil_rgb = rgb
        else:
            pil_rgb = pil.convert("RGB")

        cutout = remove(pil_rgb)
        if cutout.mode != "RGBA":
            cutout = cutout.convert("RGBA")

        bbox = _alpha_bounding_box_rgba(cutout)
        if bbox is None:
            return None, (
                "Could not detect a subject (empty alpha mask after background removal). "
                "Try another photo with a clearer foreground."
            )

        cropped = cutout.crop(bbox)
        out = _resize_pad_square_rgba(cropped, OUTPUT_SIZE)
        out = _enhance_rgba(out, CONTRAST_FACTOR, SHARPNESS_FACTOR)

        return out, "Done. Download the PNG below."

    except Exception as e:
        err = f"{type(e).__name__}: {e}"
        tb = traceback.format_exc()
        # Full traceback in console for Space logs; short message in UI
        print(tb)
        return None, f"Processing failed: {err}. Check the image format and try again."


def build_demo() -> gr.Blocks:
    theme = gr.themes.Soft(
        primary_hue="teal",
        secondary_hue="slate",
        font=[gr.themes.GoogleFont("Source Sans 3"), "ui-sans-serif", "system-ui", "sans-serif"],
    )

    with gr.Blocks(theme=theme, title="Background remover — 256×256 cutout") as demo:
        gr.Markdown(
            """
            # Subject cutout (256×256)
            Upload an image. The app removes the background with **rembg**, crops to the **alpha bounding box**,
            fits and **pads** to **256×256**, then applies a **light** contrast and sharpness boost.
            Output is a **PNG** with transparency, ready to download.
            """
        )

        with gr.Row():
            with gr.Column(scale=1):
                inp = gr.Image(
                    label="Upload image",
                    type="pil",
                    image_mode="RGB",
                    sources=["upload", "clipboard"],
                    height=400,
                )
                run_btn = gr.Button("Process", variant="primary")

            with gr.Column(scale=1):
                out_img = gr.Image(
                    label="Result (256×256 PNG)",
                    type="pil",
                    image_mode="RGBA",
                    format="png",
                    height=400,
                    buttons=["download", "fullscreen", "share"],
                )
                status = gr.Textbox(label="Status", interactive=False, lines=3)

        gr.Markdown(
            "*Large images may take a few seconds on first run while models load.*"
        )

        run_btn.click(fn=process_image, inputs=[inp], outputs=[out_img, status])

    return demo


if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    demo = build_demo()
    demo.queue()
    demo.launch(server_name="0.0.0.0", server_port=port)