import gradio as gr
from transformers import AutoImageProcessor, SegformerForSemanticSegmentation, pipeline
from PIL import Image, ImageOps, ImageFilter
import numpy as np
import torch

# ----- Load models once -----
seg_model_id = "nvidia/segformer-b0-finetuned-ade-512-512"
depth_model_id = "depth-anything/Depth-Anything-V2-Base-hf"

seg_processor = AutoImageProcessor.from_pretrained(seg_model_id)
seg_model = SegformerForSemanticSegmentation.from_pretrained(seg_model_id)

depth_pipe = pipeline(
    task="depth-estimation",
    model=depth_model_id,
    device=0 if torch.cuda.is_available() else -1,
)

# ----- Gaussian background blur using segmentation -----
def gaussian_background_blur(img: Image.Image) -> Image.Image:
    img = ImageOps.fit(img.convert("RGB"), (512, 512), method=Image.BICUBIC)

    inputs = seg_processor(images=img, return_tensors="pt")
    with torch.no_grad():
        outputs = seg_model(**inputs)

    logits = outputs.logits
    upsampled = torch.nn.functional.interpolate(
        logits, size=(512, 512), mode="bilinear", align_corners=False
    )
    seg = upsampled.argmax(dim=1)[0].cpu().numpy()

    id2label = seg_model.config.id2label
    person_ids = [i for i, label in id2label.items() if "person" in label.lower()]
    mask = np.isin(seg, person_ids).astype(np.uint8)

    mask_pil = Image.fromarray(mask * 255, mode="L")
    blurred_bg = img.filter(ImageFilter.GaussianBlur(radius=15))
    out = Image.composite(img, blurred_bg, mask_pil)
    return out

# ----- Depth-based lens blur -----
def depth_lens_blur(img: Image.Image) -> Image.Image:
    img = ImageOps.fit(img.convert("RGB"), (512, 512), method=Image.BICUBIC)

    depth_output = depth_pipe(img)
    depth_tensor = depth_output["predicted_depth"]
    depth_np = depth_tensor.squeeze().cpu().numpy()

    # normalize, then invert so far = more blur, near = sharp
    d_min, d_max = depth_np.min(), depth_np.max()
    depth_norm = (depth_np - d_min) / (d_max - d_min + 1e-8)  # [0,1]
    blur_norm = 1.0 - depth_norm    # near≈1 -> 0 blur, far≈0 -> 1 blur

    max_radius = 15.0
    num_levels = 6
    radii = np.linspace(0, max_radius, num_levels)
    blurred_versions = [
        img.filter(ImageFilter.GaussianBlur(radius=float(r))) for r in radii
    ]
    blurred_np = [np.array(b) for b in blurred_versions]

    level_size = 1.0 / (num_levels - 1)
    blur_levels = np.floor(blur_norm / level_size).astype(np.int32)
    blur_levels = np.clip(blur_levels, 0, num_levels - 1)

    H, W = blur_levels.shape
    out_np = np.zeros((H, W, 3), dtype=np.uint8)

    for lvl in range(num_levels):
        mask = blur_levels == lvl
        if not np.any(mask):
            continue
        mask_3c = np.repeat(mask[:, :, None], 3, axis=2)
        out_np[mask_3c] = blurred_np[lvl][mask_3c]

    return Image.fromarray(out_np)

# ----- Gradio UI -----
def apply_effect(img, mode):
    if img is None:
        return None
    if mode == "Gaussian background blur":
        return gaussian_background_blur(img)
    elif mode == "Depth-based lens blur":
        return depth_lens_blur(img)
    else:
        return img

demo = gr.Interface(
    fn=apply_effect,
    inputs=[
        gr.Image(type="pil", label="Upload an image"),
        gr.Radio(
            ["Gaussian background blur", "Depth-based lens blur"],
            value="Gaussian background blur",
            label="Effect",
        ),
    ],
    outputs=gr.Image(label="Output"),
    title="Gaussian & Depth-based Lens Blur Demo",
    description="Upload a selfie or scene and choose Gaussian background blur or depth-based lens blur.",
)

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