import gradio as gr
from transformers import pipeline
from PIL import Image
import numpy as np
from scipy.ndimage import gaussian_filter
import matplotlib.pyplot as plt

# Load the depth estimation model
depth_estimator = pipeline("depth-estimation", model="Intel/zoedepth-nyu-kitti")

def process_image(input_image):
    # Convert Gradio input (numpy array) to PIL Image
    input_image = Image.fromarray(input_image.astype('uint8'), 'RGB')

    # Perform depth estimation
    depth_results = depth_estimator(input_image)
    depth_map = depth_results["depth"]

    # Convert depth map to numpy array and normalize to [0, 1]
    depth_array = np.array(depth_map)
    normalized_depth = (depth_array - np.min(depth_array)) / (np.max(depth_array) - np.min(depth_array))

    # Convert input image to numpy array
    img_array = np.array(input_image)

    # Create variable blur effect
    max_blur = 5.0  # Maximum blur radius
    min_blur = 0.5  # Minimum blur radius to avoid completely sharp areas
    n_steps = 10  # Number of blur levels

    # Create output array
    blurred_array = np.zeros_like(img_array, dtype=np.float32)

    # Apply variable blur by processing the image with multiple blur levels
    for i in range(n_steps):
        sigma = min_blur + (max_blur - min_blur) * i / (n_steps - 1)
        # Apply Gaussian blur with current sigma to the whole image
        blurred_r = gaussian_filter(img_array[:,:,0], sigma=sigma)
        blurred_g = gaussian_filter(img_array[:,:,1], sigma=sigma)
        blurred_b = gaussian_filter(img_array[:,:,2], sigma=sigma)
        blurred_temp = np.stack([blurred_r, blurred_g, blurred_b], axis=2)
        
        # Create a mask for this blur level
        lower_bound = i / n_steps
        upper_bound = (i + 1) / n_steps
        mask = (normalized_depth >= lower_bound) & (normalized_depth < upper_bound)
        mask = mask[..., np.newaxis]  # Add channel dimension
        
        # Apply this blur level to the appropriate regions
        blurred_array = np.where(mask, blurred_temp, blurred_array)

    # Convert back to uint8
    blurred_image = Image.fromarray(blurred_array.astype('uint8'))

    # Create side-by-side visualization
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    axes[0].imshow(input_image)
    axes[0].set_title("Input Image")
    axes[0].axis("off")

    axes[1].imshow(depth_map, cmap="gray")
    axes[1].set_title("Depth Map")
    axes[1].axis("off")

    axes[2].imshow(blurred_image)
    axes[2].set_title("Variable Blur Output")
    axes[2].axis("off")

    plt.tight_layout()

    # Save the figure to a temporary file or buffer to display in Gradio
    output_path = "output.png"
    plt.savefig(output_path, bbox_inches='tight')
    plt.close()

    return output_path

# Define Gradio interface
interface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(label="Upload an Image"),
    outputs=gr.Image(label="Processed Output"),
    title="Depth-Based Variable Blur App",
    description="Upload an image to apply a variable blur effect based on depth estimation."
)

# Launch the app
interface.launch()