# app.py
import gradio as gr
import torch
import numpy as np
from PIL import Image, ImageFilter
from transformers import AutoImageProcessor, AutoModelForSemanticSegmentation, DPTFeatureExtractor, DPTForDepthEstimation
import torchvision.transforms as T

# Load segmentation model
seg_model_name = "nvidia/segformer-b0-finetuned-ade-512-512"
seg_processor = AutoImageProcessor.from_pretrained(seg_model_name)
seg_model = AutoModelForSemanticSegmentation.from_pretrained(seg_model_name)

# Load depth model
depth_model_name = "Intel/dpt-hybrid-midas"
depth_processor = DPTFeatureExtractor.from_pretrained(depth_model_name)
depth_model = DPTForDepthEstimation.from_pretrained(depth_model_name)

def process(image):
    image = image.convert("RGB").resize((512, 512))
    image_np = np.array(image)

    # --- Segmentation ---
    inputs = seg_processor(images=image, return_tensors="pt")
    with torch.no_grad():
        logits = seg_model(**inputs).logits
    upsampled_logits = torch.nn.functional.interpolate(
        logits, size=image.size[::-1], mode="bilinear", align_corners=False
    )
    pred_mask = upsampled_logits.argmax(dim=1)[0]
    foreground_mask = (pred_mask == 12).byte().cpu().numpy() * 255

    # --- Gaussian Blur (Zoom Style) ---
    blurred_image = image.filter(ImageFilter.GaussianBlur(radius=15))
    mask_img = Image.fromarray(foreground_mask.astype(np.uint8)).convert("L")
    gaussian_blur_result = Image.composite(image, blurred_image, mask_img)

    # --- Depth Estimation ---
    inputs = depth_processor(images=image, return_tensors="pt")
    with torch.no_grad():
        depth = depth_model(**inputs).predicted_depth[0]

    depth_resized = torch.nn.functional.interpolate(
        depth.unsqueeze(0).unsqueeze(0), size=image.size[::-1], mode="bicubic", align_corners=False
    ).squeeze().cpu().numpy()
    depth_norm = (depth_resized - depth_resized.min()) / (depth_resized.max() - depth_resized.min())
    depth_norm = 1.0 - depth_norm  # Invert so farther = more blur

    # --- Depth-Based Variable Blur ---
    num_levels = 10
    max_radius = 20
    blurred_layers = []
    for i in range(num_levels):
        r = (i / (num_levels - 1)) * max_radius
        blurred = image.filter(ImageFilter.GaussianBlur(radius=r))
        blurred_layers.append(np.array(blurred, dtype=np.float32))

    depth_indices = depth_norm * (num_levels - 1)
    output = np.zeros_like(image_np, dtype=np.float32)
    mask_np = np.array(mask_img)

    for y in range(image_np.shape[0]):
        for x in range(image_np.shape[1]):
            if mask_np[y, x] > 128:
                output[y, x] = image_np[y, x]  # preserve foreground
            else:
                d = depth_indices[y, x]
                low = int(np.floor(d))
                high = min(low + 1, num_levels - 1)
                alpha = d - low
                pixel = (1 - alpha) * blurred_layers[low][y, x] + alpha * blurred_layers[high][y, x]
                output[y, x] = pixel

    depth_blur_result = Image.fromarray(np.uint8(np.clip(output, 0, 255)))

    return image, mask_img, gaussian_blur_result, depth_blur_result

# Gradio Interface
iface = gr.Interface(
    fn=process,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=[
        gr.Image(type="pil", label="Original Image"),
        gr.Image(type="pil", label="Foreground Mask"),
        gr.Image(type="pil", label="Gaussian Background Blur"),
        gr.Image(type="pil", label="Depth-Based Lens Blur")
    ],
    title="Image Blur Effects Demo",
    description="Upload an image to apply Gaussian background blur and depth-based lens blur using Hugging Face models."
)

iface.launch()