import gradio as gr
import torch
import numpy as np
from PIL import Image
from scipy.ndimage import gaussian_filter
from transformers import pipeline

def preprocess_image(image):
    """Resize and convert image to PIL format if needed."""
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    
    # Resize to 512x512 while maintaining aspect ratio
    image = image.resize((512, 512))
    return image

def segment_image(image, model_name="yolov8n-seg"):
    """
    Perform instance segmentation on the input image using YOLO segmentation model.
    """
    from ultralytics import YOLO
    
    # Load the YOLO segmentation model
    model = YOLO(model_name)
    
    # Run inference
    results = model(image)
    
    # Create a blank mask (1 for foreground, 0 for background)
    mask = np.zeros((image.size[1], image.size[0]), dtype=np.uint8)
    
    # Process each detected object
    for result in results:
        if result.masks is not None:
            for single_mask in result.masks:
                # Convert mask to numpy and resize if needed
                mask_array = single_mask.data.cpu().numpy().squeeze()
                mask_array = (mask_array > 0.5).astype(np.uint8)
                
                # Resize if needed
                if mask_array.shape != mask.shape:
                    mask_array = np.array(
                        Image.fromarray(mask_array).resize(
                            (image.size[0], image.size[1]), 
                            Image.NEAREST
                        )
                    )
                
                # Add this mask to the overall mask (OR operation)
                mask = np.maximum(mask, mask_array)
    
    return mask

def apply_gaussian_blur(image, sigma=15):
    """Apply Gaussian blur to the background."""
    # Convert image to numpy array
    image_array = np.array(image)
    
    # Get segmentation mask (1 for foreground, 0 for background)
    foreground_mask = segment_image(image)
    
    # Prepare blurred version
    blurred = np.zeros_like(image_array)
    for channel in range(3):
        blurred[:, :, channel] = gaussian_filter(image_array[:, :, channel], sigma=sigma)
    
    # Combine original and blurred images based on mask
    mask_3d = np.stack([foreground_mask] * 3, axis=2)
    result = image_array * mask_3d + blurred * (1 - mask_3d)
    
    return Image.fromarray(result.astype(np.uint8))

def estimate_depth(image, model_name="depth-anything/Depth-Anything-V2-Small-hf"):
    """Estimate depth of the image."""
    depth_estimator = pipeline(
        task="depth-estimation", 
        model=model_name,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
    )
    
    depth_output = depth_estimator(image)
    depth_map = np.array(depth_output["depth"])
    
    # Normalize depth map (0-1 where 1 is farthest)
    depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
    
    return depth_map

def apply_depth_aware_blur(image, max_sigma=15, min_sigma=0):
    """Apply depth-aware blur with farther objects more blurred."""
    # Estimate depth (1 = nearest, 0 = farthest)
    depth_map = estimate_depth(image)  # Returns 1 for near, 0 for far
    
    # INVERT the depth map (now 1 = farthest, 0 = nearest)
    depth_map = 1 - depth_map
    
    image_array = np.array(image)
    
    # Create single blurred version at max sigma
    max_blurred = np.zeros_like(image_array, dtype=np.float32)
    for channel in range(3):
        max_blurred[:, :, channel] = gaussian_filter(
            image_array[:, :, channel].astype(np.float32),
            sigma=max_sigma
        )
    
    # Create 3-channel depth map for blending
    depth_3d = np.stack([depth_map] * 3, axis=2)
    
    # Blend: More depth (farther) = more blur
    result = image_array * (1 - depth_3d) + max_blurred * depth_3d
    
    return Image.fromarray(result.astype(np.uint8))

def process_image(image, blur_type, sigma=15):
    """Process image based on blur type."""
    # Preprocess image
    pil_image = preprocess_image(image)
    
    # Apply appropriate blur
    if blur_type == "Gaussian Background Blur":
        result = apply_gaussian_blur(pil_image, sigma)
    elif blur_type == "Depth-Aware Lens Blur":
        result = apply_depth_aware_blur(pil_image, max_sigma=sigma)
    else:
        result = pil_image
    
    return result

# Gradio Interface
def create_blur_app():
    with gr.Blocks() as demo:
        gr.Markdown("# Image Blur Effects")
        
        with gr.Row():
            input_image = gr.Image(label="Input Image", type="pil")
            output_image = gr.Image(label="Processed Image")
        
        with gr.Row():
            blur_type = gr.Dropdown(
                choices=[
                    "Gaussian Background Blur", 
                    "Depth-Aware Lens Blur"
                ], 
                label="Blur Type"
            )
            sigma = gr.Slider(
                minimum=0, 
                maximum=30, 
                value=15, 
                label="Blur Intensity"
            )
        
        process_btn = gr.Button("Apply Blur Effect")
        
        process_btn.click(
            fn=process_image, 
            inputs=[input_image, blur_type, sigma], 
            outputs=output_image
        )
    
    return demo

# Launch the app
if __name__ == "__main__":
    demo = create_blur_app()
    demo.launch()