import gradio as gr
from PIL import Image, ImageFilter
import numpy as np
import cv2
import torch
from transformers import DPTFeatureExtractor, DPTForDepthEstimation

# Load model and feature extractor
feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")

# Gaussian Blur function
def apply_gaussian_blur(image, blur_radius):
    return image.filter(ImageFilter.GaussianBlur(blur_radius))

# Lens Blur function
def apply_lens_blur(image):
    # Get depth map
    inputs = feature_extractor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        depth_map = outputs.predicted_depth.squeeze().cpu().numpy()
    depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min()) * 15
    depth_map_resized = cv2.resize(depth_map, (image.width, image.height))
    depth_map_resized = 15 - depth_map_resized

    # Convert to OpenCV format
    image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    blurred_image = np.zeros_like(image_cv, dtype=np.float32)

    for blur_radius in range(1, 16):
        blurred_layer = cv2.GaussianBlur(image_cv, (0, 0), sigmaX=blur_radius)
        mask = ((depth_map_resized >= (blur_radius - 1)) & (depth_map_resized < blur_radius)).astype(np.float32)
        mask = cv2.merge([mask] * 3)
        blurred_image += blurred_layer * mask

    blurred_image = np.clip(blurred_image, 0, 255).astype(np.uint8)
    return Image.fromarray(cv2.cvtColor(blurred_image, cv2.COLOR_BGR2RGB))

# Gradio app interface
def process_image(image, effect, blur_radius):
    if effect == "Gaussian Blur":
        return apply_gaussian_blur(image, blur_radius)
    elif effect == "Lens Blur":
        return apply_lens_blur(image)
    else:
        return image

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Gaussian and Lens Blur Effects")
    with gr.Row():
        with gr.Column():
            uploaded_image = gr.Image(type="pil")
            effect = gr.Radio(["Gaussian Blur", "Lens Blur"], value="Gaussian Blur", label="Effect")
            blur_radius = gr.Slider(1, 15, value=5, step=1, label="Blur Radius (for Gaussian Blur)")
            submit_button = gr.Button("Apply Effect")
        with gr.Column():
            output_image = gr.Image(type="pil", label="Processed Image")

    submit_button.click(process_image, inputs=[uploaded_image, effect, blur_radius], outputs=output_image)

# Launch the app
demo.launch()