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

# Load the pre-trained segmentation model
feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
segmentation_model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b1-finetuned-cityscapes-1024-1024")

def apply_blur_effect(image, blur_type):
    """
    Applies either Gaussian blur to the whole image, 
    depth-based blur, or background blur while keeping the foreground sharp.
    """
    image = image.resize((512, 512))  # Resize input image

    if blur_type == "Depth-Based Blur(Lens Blur)":
        # Use depth estimation model to get depth map
        depth_estimator = pipeline(task="depth-estimation", model="Intel/zoedepth-nyu-kitti")
        outputs = depth_estimator(image)
        depth_map = np.array(outputs["depth"])

        # Normalize depth map
        depth_map_normalized = (depth_map - np.min(depth_map)) / (np.max(depth_map) - np.min(depth_map))
        depth_array = np.clip(depth_map_normalized * 5, 0, 5).astype(int)  # Scale depth to select blur levels

        # Generate different levels of Gaussian blur
        blur_levels = [image.filter(ImageFilter.GaussianBlur(radius=r)) for r in range(6)]
        
        # Create depth-based blur image
        depth_based_blur_image = Image.new("RGB", image.size)
        for i in range(image.width):
            for j in range(image.height):
                blur_level = depth_array[j, i]
                depth_based_blur_image.putpixel((i, j), blur_levels[blur_level].getpixel((i, j)))
        
        return depth_based_blur_image

    elif blur_type == "Gaussian Background Blur":
        # Perform segmentation to get foreground and background masks
        inputs = feature_extractor(images=image, return_tensors="pt")
        outputs = segmentation_model(**inputs)
        logits = outputs.logits  # Shape: (batch, num_classes, height, width)
        predicted_mask = torch.argmax(logits, dim=1).squeeze().cpu().numpy()

        # Create a binary mask (1 = foreground, 0 = background)
        foreground_mask = (predicted_mask == 11).astype(np.uint8)

        # Convert the mask into a PIL image for processing
        mask_image = Image.fromarray((foreground_mask * 255).astype(np.uint8)).resize(image.size)

        # Apply Gaussian blur to the entire image
        blurred_image = image.filter(ImageFilter.GaussianBlur(radius=15))

        # Blend the sharp foreground and blurred background
        final_image = Image.composite(image, blurred_image, mask_image)

        return final_image

    return image

# Gradio UI
interface = gr.Interface(
    fn=apply_blur_effect,
    inputs=[
        gr.Image(type="pil"),  
        gr.Radio(["Depth-Based Blur(Lens Blur)", "Gaussian Background Blur"], label="Blur Type"),
    ],
    outputs="image",
    title="Image Blur Effects: Gaussian Background Blur, Depth-Based Blur(Lens Blur)",
    description="Upload a selfie (needs some background, not just the face!) and apply Gaussian Background Blur or Depth-Based Blur(Lens Blur)",
)

# Launch the Gradio app
if __name__ == "__main__":
    interface.launch()