import gradio as gr
import torch
import numpy as np
from transformers import AutoImageProcessor, AutoModelForDepthEstimation
from PIL import Image, ImageFilter

def load_depth_model():
    """
    Loads the depth estimation model and processor.
    Returns (processor, model, device).
    """
    global processor, model, device
    if "model" not in globals():
        processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2")
        model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2")
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model.to(device).eval()
    return processor, model, device

def compute_depth_map(image: Image.Image, scale_factor: float) -> np.ndarray:
    """
    Computes the depth map for a PIL image.
    Inverts the map (i.e. force invert_depth=True) and scales it.
    Returns a NumPy array in [0, 1]*scale_factor.
    """
    processor, model, device = load_depth_model()
    inputs = processor(images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
        predicted_depth = outputs.predicted_depth

    prediction = torch.nn.functional.interpolate(
        predicted_depth.unsqueeze(1),
        size=image.size[::-1],  # PIL image size: (width, height)
        mode="bicubic",
        align_corners=False,
    )
    depth_min = prediction.min()
    depth_max = prediction.max()
    depth_vis = (prediction - depth_min) / (depth_max - depth_min + 1e-8)
    depth_map = depth_vis.squeeze().cpu().numpy()
    # Always invert depth so that near=0 and far=1
    depth_map = 1.0 - depth_map
    depth_map *= scale_factor
    return depth_map

def layered_blur(image: Image.Image, depth_map: np.ndarray, num_layers: int, max_blur: float) -> Image.Image:
    """
    Creates multiple blurred versions of 'image' (radii from 0 to max_blur)
    and composites them based on the depth map split into num_layers bins.
    """
    blur_radii = np.linspace(0, max_blur, num_layers)
    blur_versions = [image.filter(ImageFilter.GaussianBlur(r)) for r in blur_radii]
    upper_bound = depth_map.max()
    thresholds = np.linspace(0, upper_bound, num_layers + 1)
    final_image = blur_versions[-1]
    for i in range(num_layers - 1, -1, -1):
        mask_array = np.logical_and(
            depth_map >= thresholds[i],
            depth_map < thresholds[i + 1]
        ).astype(np.uint8) * 255
        mask_image = Image.fromarray(mask_array, mode="L")
        final_image = Image.composite(blur_versions[i], final_image, mask_image)
    return final_image

def process_depth_blur(uploaded_image, max_blur_value, scale_factor, num_layers):
    """
    Processes the image with a depth-based blur.
    The image is resized to 512x512, its depth is computed (with invert_depth always True),
    and a layered blur is applied.
    """
    if not isinstance(uploaded_image, Image.Image):
        uploaded_image = Image.open(uploaded_image)
    image = uploaded_image.convert("RGB").resize((512, 512))
    depth_map = compute_depth_map(image, scale_factor)
    final_image = layered_blur(image, depth_map, int(num_layers), max_blur_value)
    return final_image

with gr.Blocks() as demo:
    gr.Markdown("# Depth-Based Lens Blur")
    depth_img = gr.Image(type="pil", label="Upload Image")
    depth_max_blur = gr.Slider(1.0, 5.0, value=3.0, step=0.1, label="Maximum Blur Radius")
    depth_scale = gr.Slider(0.1, 1.0, value=0.5, step=0.1, label="Depth Scale Factor")
    depth_layers = gr.Slider(2, 20, value=8, step=1, label="Number of Layers")
    depth_out = gr.Image(label="Depth-Based Blurred Image")
    depth_button = gr.Button("Process Depth Blur")
    depth_button.click(process_depth_blur, 
                       inputs=[depth_img, depth_max_blur, depth_scale, depth_layers],
                       outputs=depth_out)

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