import gradio as gr

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

# Load models
seg_processor = AutoImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
seg_model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")

depth_processor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas")
depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas")

# Segmentation + Gaussian Blur
def segmentation_blur(image):
    inputs = seg_processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = seg_model(**inputs)
    pred = outputs.logits.argmax(dim=1)[0].cpu().numpy()

    mask = (pred == 12).astype(np.uint8) * 255  # Person class
    mask_resized = cv2.resize(mask, image.size, interpolation=cv2.INTER_NEAREST)

    image_np = np.array(image)
    mask_3ch = np.stack([mask_resized] * 3, axis=-1) // 255
    blurred = cv2.GaussianBlur(image_np, (0, 0), sigmaX=15)
    combined = np.where(mask_3ch == 1, image_np, blurred)

    return Image.fromarray(combined)

# Depth-based Variable Blur
def apply_depth_based_blur(image_np, depth_map_norm, max_blur=25):
    h, w = depth_map_norm.shape
    blurred_image = np.zeros_like(image_np)
    image_np = image_np.astype(np.float32)

    for i in range(1, max_blur+1, 2):
        mask = ((depth_map_norm >= (i / max_blur)) & (depth_map_norm < ((i + 2) / max_blur))).astype(np.uint8)
        if np.sum(mask) == 0:
            continue
        blurred = cv2.GaussianBlur(image_np, (i, i), 0)
        mask_3ch = np.stack([mask]*3, axis=-1)
        blurred_image = np.where(mask_3ch == 1, blurred, blurred_image)

    sharpest_mask = (depth_map_norm < 0.1).astype(np.uint8)
    sharpest_mask_3ch = np.stack([sharpest_mask]*3, axis=-1)
    final_image = np.where(sharpest_mask_3ch == 1, image_np, blurred_image)

    return final_image.astype(np.uint8)

def depth_blur(image):
    inputs = depth_processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = depth_model(**inputs)

    depth_map = outputs.predicted_depth[0].cpu().numpy()
    depth_map = cv2.resize(depth_map, image.size)
    norm_depth = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())

    image_np = np.array(image)
    blurred_image = apply_depth_based_blur(image_np, norm_depth)
    return Image.fromarray(blurred_image)

# Gradio interface
def process(image, mode):
    if mode == "Segmentation Blur":
        return segmentation_blur(image)
    else:
        return depth_blur(image)

demo = gr.Interface(
    fn=process,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Radio(["Segmentation Blur", "Depth Blur"], value="Segmentation Blur", label="Blur Mode")
    ],
    outputs=gr.Image(label="Output"),
    title="Image Blur Effects (Segmentation & Depth)",
    description="Upload an image and apply either Zoom-style background blur or DSLR-style depth blur."
)

demo.launch()