import gradio as gr
import torch
import numpy as np
from transformers import AutoImageProcessor, AutoModelForDepthEstimation
from PIL import Image, ImageFilter
import matplotlib.pyplot as plt
import matplotlib.cm as cm
def compute_depth_map(image: Image.Image, scale_factor: float) -> np.ndarray:
    image_processor = AutoImageProcessor.from_pretrained("LiheYoung/depth-anything-large-hf")
    model = AutoModelForDepthEstimation.from_pretrained("LiheYoung/depth-anything-large-hf")
    inputs = image_processor(images=image, return_tensors="pt")
    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],  
        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()
    depth_map_inverted = 1.0 - depth_map
    depth_map_inverted *= scale_factor
    return depth_map_inverted

def layered_blur(image: Image.Image, depth_map: np.ndarray, num_layers: int, max_blur: float) -> Image.Image:
    blur_radii = np.linspace(0, max_blur, num_layers)
    blur_versions = [image.filter(ImageFilter.GaussianBlur(radius)) for radius in blur_radii]
    thresholds = np.linspace(0, 1, 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):
    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
def create_heatmap(depth_map: np.ndarray, intensity: float) -> Image.Image:
    normalized = np.clip(depth_map * intensity, 0, 1)
    colormap = cm.get_cmap("inferno")
    colored = colormap(normalized) 
    heatmap = (colored[:, :, :3] * 255).astype(np.uint8) 
    return Image.fromarray(heatmap)

def process_depth_heatmap(uploaded_image, intensity):
    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=1.0)
    heatmap_img = create_heatmap(depth_map, intensity)
    return heatmap_img
def load_segmentation_model():
    global seg_model, seg_device
    if "seg_model" not in globals():
        from ben2 import BEN_Base  # Import BEN2
        seg_model = BEN_Base.from_pretrained("PramaLLC/BEN2")
        seg_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        seg_model.to(seg_device).eval()
    return seg_model, seg_device

def process_segmentation_blur(uploaded_image, seg_blur_radius: float):
    if not isinstance(uploaded_image, Image.Image):
        uploaded_image = Image.open(uploaded_image)
    image = uploaded_image.convert("RGB").resize((512, 512))
    seg_model, seg_device = load_segmentation_model()
    blurred_image = image.filter(ImageFilter.GaussianBlur(seg_blur_radius))
    foreground = seg_model.inference(image, refine_foreground=False)
    foreground_rgba = foreground.convert("RGBA")
    _, _, _, alpha = foreground_rgba.split()
    binary_mask = alpha.point(lambda x: 255 if x > 128 else 0, mode="L")
    final_image = Image.composite(image, blurred_image, binary_mask)
    return final_image


with gr.Blocks() as demo:
    gr.Markdown("Gaussian Blur & Lens Blur Effect")
    with gr.Tabs():
        with gr.Tab("Gaussian Blur"):
            seg_img = gr.Image(type="pil", label="Upload Image")
            seg_blur = gr.Slider(5, 30, value=15, step=1, label="Gaussian Blur Radius")
            seg_out = gr.Image(label="Gaussian Blurred Image")
            seg_button = gr.Button("Process Gaussian Blur")
            seg_button.click(process_segmentation_blur, inputs=[seg_img, seg_blur], outputs=seg_out)
        with gr.Tab("Depth Heatmap"):
            img_input2 = gr.Image(type="pil", label="Upload Image")
            intensity_slider = gr.Slider(0.5, 5.0, value=1.0, label="Heatmap Intensity")
            heatmap_output = gr.Image(label="Depth Heatmap")
            heatmap_button = gr.Button("Generate Depth Heatmap")
            heatmap_button.click(
                process_depth_heatmap, 
                inputs=[img_input2, intensity_slider], 
                outputs=heatmap_output
            )
        with gr.Tab("Lens Blur"):
            img_input = gr.Image(type="pil", label="Upload Image")
            blur_slider = gr.Slider(1, 50, value=6, label="Maximum Blur Radius")
            scale_slider = gr.Slider(0.1, 2.0, value=0.72, label="Depth Scale Factor")
            layers_slider = gr.Slider(2, 10, value=2.91, label="Number of Layers")
            blur_output = gr.Image(label="Lens Blur Result")
            blur_button = gr.Button("Process Blur")
            blur_button.click(
                process_depth_blur, 
                inputs=[img_input, blur_slider, scale_slider, layers_slider], 
                outputs=blur_output
            )
       

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