import gradio as gr
from transformers import pipeline
from PIL import Image, ImageFilter
import numpy as np

# Load models from Hugging Face
segmentation_model = pipeline("image-segmentation", model="nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
depth_estimator = pipeline("depth-estimation", model="Intel/zoedepth-nyu-kitti")

def process_image(image, blur_type, sigma):
    # Step 1: Perform segmentation
    segmentation_results = segmentation_model(image)
    foreground_mask = segmentation_results[-1]["mask"]

    # Step 2: Apply Gaussian blur to background
    blurred_background = image.filter(ImageFilter.GaussianBlur(sigma))
    segmented_output = Image.composite(image, blurred_background, foreground_mask)

    # Step 3: Perform depth estimation
    depth_results = depth_estimator(image)
    depth_map = depth_results["depth"]

    # Step 4: Normalize depth map values
    depth_array = np.array(depth_map)
    normalized_depth = (depth_array - np.min(depth_array)) / (np.max(depth_array) - np.min(depth_array)) * 255
    normalized_depth_image = Image.fromarray(normalized_depth.astype('uint8'))

    # Step 5: Apply variable Gaussian blur based on depth map (Lens Blur)
    if blur_type == "Lens Blur":
        variable_blur_image = image.copy()
        for x in range(variable_blur_image.width):
            for y in range(variable_blur_image.height):
                blur_intensity = normalized_depth[y, x] / 255 * sigma  # Scale blur intensity by depth value
                pixel_value = image.getpixel((x, y))
                variable_blur_image.putpixel((x, y), tuple(int(p * blur_intensity) for p in pixel_value))
        output_image = variable_blur_image
    else:
        output_image = segmented_output

    return segmented_output, normalized_depth_image, output_image

# Create Gradio interface
app = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Radio(["Gaussian Blur", "Lens Blur"], label="Blur Type", value="Gaussian Blur"),
        gr.Slider(0, 50, step=1, label="Blur Intensity (Sigma)", value=15)
    ],
    outputs=[
        gr.Image(type="pil", label="Segmented Output with Background Blur"),
        gr.Image(type="pil", label="Depth Map Visualization"),
        gr.Image(type="pil", label="Final Output with Selected Blur")
    ],
    title="Vision Transformer Segmentation & Depth-Based Blur Effects",
    description="Upload an image and select the type of blur effect (Gaussian or Lens). Adjust the blur intensity using the slider."
)
app.launch()