import gradio as gr
import cv2
import numpy as np
import torch
from PIL import Image
from transformers import (
    SegformerImageProcessor, 
    SegformerForSemanticSegmentation,
    AutoImageProcessor, 
    AutoModelForDepthEstimation
)

# Load Segformer model for Gaussian blur
segformer_processor = SegformerImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
segformer_model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")

# Load Depth-Anything model for lens blur
depth_processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
depth_model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")

def apply_blur(image, blur_type, blur_strength, depth_threshold):
    # Convert image to RGB
    img = image

    if blur_type == "Gaussian":
        # Use Segformer for Gaussian blur
        pil_image = Image.fromarray(img)
        inputs = segformer_processor(images=pil_image, return_tensors="pt")
        outputs = segformer_model(**inputs)
        logits = outputs.logits

        mask = logits[0, 12, :, :].detach().cpu().numpy() > depth_threshold
        mask = cv2.resize(mask.astype(np.uint8), (img.shape[1], img.shape[0]))

    elif blur_type == "Lens":
        # Use Depth-Anything for lens blur
        pil_image = Image.fromarray(img)
        inputs = depth_processor(images=pil_image, return_tensors="pt")

        with torch.no_grad():
            outputs = depth_model(**inputs)
            predicted_depth = outputs.predicted_depth

        prediction = torch.nn.functional.interpolate(
            predicted_depth.unsqueeze(1),
            size=img.shape[:2],
            mode="bicubic",
            align_corners=False,
        )

        mask = prediction[0, 0, :, :].detach().cpu().numpy() > depth_threshold
        mask = mask.astype(np.uint8)

    # Invert mask using cv2
    mask = cv2.bitwise_not(mask)
    mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2)

    # Apply blur based on selected type
    if blur_type == "Gaussian":
        blurred_image = cv2.GaussianBlur(img, (0, 0), sigmaX=blur_strength)
    elif blur_type == "Lens":
        # Simulate lens blur using a larger kernel
        kernel_size = int(blur_strength * 2) * 2 + 1
        blurred_image = cv2.GaussianBlur(img, (kernel_size, kernel_size), 0)

    # Combine blurred and original images using the mask
    output = np.where(mask == 255, blurred_image, img)
    
    return output


# Define Gradio interface
iface = gr.Interface(
    fn=apply_blur,
    inputs=[
        gr.Image(label="Input Image"),
        gr.Radio(["Gaussian", "Lens"], label="Blur Type", value="Gaussian"),
        gr.Slider(1, 30, value=15, step=1, label="Blur Strength"),
        gr.Slider(-20, 20, value=-4, step=0.1, label="Depth Threshold")
    ],
    outputs=gr.Image(label="Output Image"),
    title="Image Segmentation and Blurring",
    description="Upload an image and apply Gaussian or Lens blur to the background using different segmentation models."
)

# Launch the app
iface.launch(share=True)