import cv2
import gradio as gr
import numpy as np

from functools import lru_cache
from rembg import new_session, remove


GOOD_MODELS = [
    "u2net_human_seg",
    "silueta",
    "u2netp",
    "isnet-general-use",
    "bria-rmbg",
    "birefnet-general-lite",
    "birefnet-portrait",
]


MODEL_NOTES = {
    "u2net_human_seg": "Fast human-specific baseline. Good for full body/person masks.",
    "silueta": "Small and fast. Lower quality, useful for latency testing.",
    "u2netp": "Very light U2Net variant. Fastest baseline, lower quality.",
    "isnet-general-use": "General-purpose background removal fallback.",
    "bria-rmbg": "High-quality general background remover. Check license for your use.",
    "birefnet-general-lite": "Lighter BiRefNet. Better quality than tiny models, slower than U2Net variants.",
    "birefnet-portrait": "Best BiRefNet choice for portraits / people, usually slower.",
}


@lru_cache(maxsize=3)
def get_session(model_name: str):
    return new_session(model_name)


def ensure_uint8_rgb(image: np.ndarray) -> np.ndarray:
    if image is None:
        raise gr.Error("Please upload an image.")

    image = np.asarray(image)

    if image.ndim == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)

    if image.shape[-1] == 4:
        image = image[:, :, :3]

    if image.dtype != np.uint8:
        image = np.clip(image, 0, 255).astype(np.uint8)

    return image


def normalize_mask(mask, width: int, height: int) -> np.ndarray:
    mask = np.asarray(mask)

    if mask.ndim == 3:
        mask = mask[:, :, 0]

    if mask.shape[:2] != (height, width):
        mask = cv2.resize(
            mask,
            (width, height),
            interpolation=cv2.INTER_LANCZOS4,
        )

    if mask.dtype != np.uint8:
        mask = np.clip(mask, 0, 255).astype(np.uint8)

    return mask


def composite_on_white(
    image_rgb: np.ndarray,
    mask_u8: np.ndarray,
    edge_blur: float,
) -> np.ndarray:
    alpha = mask_u8.astype(np.float32) / 255.0

    if edge_blur > 0:
        alpha = cv2.GaussianBlur(alpha, (0, 0), sigmaX=edge_blur)
        alpha = np.clip(alpha, 0.0, 1.0)

    alpha = alpha[:, :, None]

    white = np.full_like(image_rgb, 255)

    result = (
        image_rgb.astype(np.float32) * alpha
        + white.astype(np.float32) * (1.0 - alpha)
    )

    return np.clip(result, 0, 255).astype(np.uint8)


def inference(
    image,
    model,
    post_process_mask,
    alpha_matting,
    edge_blur,
):
    image_rgb = ensure_uint8_rgb(image)
    height, width = image_rgb.shape[:2]

    session = get_session(model)

    mask = remove(
        image_rgb,
        session=session,
        only_mask=True,
        post_process_mask=post_process_mask,
        alpha_matting=alpha_matting,
    )

    mask_u8 = normalize_mask(mask, width=width, height=height)

    white_result = composite_on_white(
        image_rgb=image_rgb,
        mask_u8=mask_u8,
        edge_blur=edge_blur,
    )

    return white_result, mask_u8, MODEL_NOTES.get(model, "")


def compare_all(
    image,
    post_process_mask,
    alpha_matting,
    edge_blur,
):
    image_rgb = ensure_uint8_rgb(image)

    outputs = []

    for model in GOOD_MODELS:
        result, _, _ = inference(
            image=image_rgb,
            model=model,
            post_process_mask=post_process_mask,
            alpha_matting=alpha_matting,
            edge_blur=edge_blur,
        )

        outputs.append((result, model))

    return outputs


with gr.Blocks() as app:
    gr.Markdown("# Person Background Removal Benchmark")
    gr.Markdown(
        """
Remove background with [Rembg](https://github.com/danielgatis/rembg) models.
"""
    )

    with gr.Row():
        input_image = gr.Image(type="numpy", label="Input Image")

        with gr.Column():
            output_image = gr.Image(type="numpy", label="White Background Result")
            output_mask = gr.Image(type="numpy", label="Mask")
            model_note = gr.Textbox(label="Model note", interactive=False)

    with gr.Row():
        model_selector = gr.Dropdown(
            GOOD_MODELS,
            value="u2net_human_seg",
            label="Model",
        )

        edge_blur = gr.Slider(
            minimum=0.0,
            maximum=3.0,
            value=0.0,
            step=0.25,
            label="Edge blur",
        )

    with gr.Row():
        post_process_mask = gr.Checkbox(
            value=False,
            label="Post-process mask",
        )

        alpha_matting = gr.Checkbox(
            value=False,
            label="Alpha matting",
        )

    with gr.Row():
        process_button = gr.Button("Process selected model", variant="primary")
        compare_button = gr.Button("Compare all models")

    gallery = gr.Gallery(
        label="Compare all models",
        columns=3,
        height="auto",
        object_fit="contain",
    )

    process_button.click(
        inference,
        inputs=[
            input_image,
            model_selector,
            post_process_mask,
            alpha_matting,
            edge_blur,
        ],
        outputs=[
            output_image,
            output_mask,
            model_note,
        ],
    )

    compare_button.click(
        compare_all,
        inputs=[
            input_image,
            post_process_mask,
            alpha_matting,
            edge_blur,
        ],
        outputs=gallery,
    )


app.launch()