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

# Load models
segformer_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b1-finetuned-ade-512-512")
segformer_model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b1-finetuned-ade-512-512")
dpt_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
dpt_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")

# Gaussian Blur Background Function
def gaussian_blur_background(image):
    # Preprocess image for segmentation
    inputs = segformer_extractor(images=image, return_tensors="pt")
    outputs = segformer_model(**inputs)
    logits = outputs.logits
    segmentation = torch.argmax(logits, dim=1)[0].numpy()

    # Create a binary mask for 'person' class (class index 12)
    human_mask = (segmentation == 12).astype(np.uint8) * 255
    human_mask_image = Image.fromarray(human_mask).resize(image.size)

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

    # Composite the original image with blurred background using the mask
    composite_image = Image.composite(image, blurred_background, human_mask_image)
    return composite_image

# Depth-Based Lens Blur Function
def lens_blur(image):
    # Preprocess image for depth estimation
    inputs = dpt_extractor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = dpt_model(**inputs)
        depth_map = outputs.predicted_depth.squeeze().cpu().numpy()

    # Normalize depth map to range [0, 15] and invert for blur intensity
    depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min()) * 15
    depth_map = 15 - depth_map
    depth_map_resized = cv2.resize(depth_map, (image.width, image.height))

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

    # Apply variable blur based on depth
    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)
    blurred_image_pil = Image.fromarray(cv2.cvtColor(blurred_image, cv2.COLOR_BGR2RGB))
    return blurred_image_pil

# Gradio Interface
def process_image(image, effect):
    if effect == "Gaussian Blur Background":
        return gaussian_blur_background(image)
    elif effect == "Lens Blur":
        return lens_blur(image)

with gr.Blocks() as demo:
    gr.Markdown("# BokehBot: Gaussian and Lens Blur Effects")
    with gr.Row():
        with gr.Column():
            uploaded_image = gr.Image(type="pil", label="Upload an Image")
            effect = gr.Radio(["Gaussian Blur Background", "Lens Blur"], label="Choose Effect")
            process_button = gr.Button("Apply Effect")
        with gr.Column():
            output_image = gr.Image(type="pil", label="Processed Image")
    
    process_button.click(process_image, inputs=[uploaded_image, effect], outputs=output_image)

demo.launch()