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

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(input_image, method, blur_intensity):
    """
    Process the input image using one of two methods:
    
    1. Segmentation Blur Model:
       - Uses segmentation to extract a foreground mask.
       - Applies Gaussian blur to the background.
       - Composites the final image.
       
    2. Monocular Depth Estimation Model:
       - Uses depth estimation to generate a depth map.
       - Normalizes the depth map to be used as a blending mask.
       - Blends a fully blurred version with the original image.
       
    Returns:
       - output_image: final composited image.
       - mask_image: the mask used (binary for segmentation, normalized depth for depth-based).
    """
    input_image = input_image.convert("RGB")
    
    if method == "Segmentation Blur Model":
        results = segmentation_model(input_image)
        foreground_mask = results[-1]["mask"]
        foreground_mask = foreground_mask.convert("L")
        binary_mask = foreground_mask.point(lambda p: 255 if p > 128 else 0)
        
        blurred_background = input_image.filter(ImageFilter.GaussianBlur(radius=blur_intensity))
        
        output_image = Image.composite(input_image, blurred_background, binary_mask)
        mask_image = binary_mask
        
    elif method == "Monocular Depth Estimation Model":
        depth_results = depth_estimator(input_image)
        depth_map = depth_results["depth"]
        
        depth_array = np.array(depth_map).astype(np.float32)
        norm = (depth_array - depth_array.min()) / (depth_array.max() - depth_array.min() + 1e-8)
        normalized_depth = (norm * 255).astype(np.uint8)
        mask_image = Image.fromarray(normalized_depth)
        
        blurred_image = input_image.filter(ImageFilter.GaussianBlur(radius=blur_intensity))
        
        orig_np = np.array(input_image).astype(np.float32)
        blur_np = np.array(blurred_image).astype(np.float32)
        alpha = normalized_depth[..., np.newaxis] / 255.0
        
        blended_np = (1 - alpha) * orig_np + alpha * blur_np
        blended_np = np.clip(blended_np, 0, 255).astype(np.uint8)
        output_image = Image.fromarray(blended_np)
    
    else:
        output_image = input_image
        mask_image = input_image.convert("L")
    
    return output_image, mask_image

with gr.Blocks() as demo:
    gr.Markdown("## FocusFusion: Segmentation & Depth Blur")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input Image", type="pil")
            method = gr.Radio(label="Processing Method", 
                            choices=["Segmentation Blur Model", "Monocular Depth Estimation Model"],
                            value="Segmentation Blur Model")
            blur_intensity = gr.Slider(label="Blur Intensity (sigma)", 
                                    minimum=1, maximum=30, step=1, value=15)
            run_button = gr.Button("Process Image")
        with gr.Column():
            output_image = gr.Image(label="Output Image")
            mask_output = gr.Image(label="Mask")
    
    run_button.click(
        fn=process_image, 
        inputs=[input_image, method, blur_intensity], 
        outputs=[output_image, mask_output]
    )

demo.launch()