import gradio as gr
import torch
from transformers import pipeline
from PIL import Image
import numpy as np
import cv2

def process_image(image, effect_type="Gaussian Blur", blur_intensity=15):
    """
    Process the image with selected effect
    """
    # Resize image to 512x512
    image = Image.fromarray(image).resize((512, 512))
    
    if effect_type == "Gaussian Blur":
        # Generate segmentation mask
        segmenter = pipeline("image-segmentation", 
                           model="openmmlab/upernet-swin-base",
                           device=0 if torch.cuda.is_available() else -1)
        
        results = segmenter(image)
        mask = np.zeros((512, 512), dtype=np.uint8)
        
        for segment in results:
            if segment['label'].lower() == 'person':
                segment_mask = np.array(segment['mask'])
                mask[segment_mask > 0] = 255
                
        # Apply gaussian blur
        img_np = np.array(image)
        blurred = cv2.GaussianBlur(img_np, (0, 0), blur_intensity)
        
        mask_np = mask / 255.0
        mask_np = np.stack([mask_np] * 3, axis=-1)
        
        result = img_np * mask_np + blurred * (1 - mask_np)
        return result.astype(np.uint8)
        
    else:  # Depth-based blur
        # Generate depth map
        depth_estimator = pipeline("depth-estimation", 
                                 model="Intel/dpt-large",
                                 device=0 if torch.cuda.is_available() else -1)
        
        depth_result = depth_estimator(image)
        depth_map = depth_result['predicted_depth']
        if torch.is_tensor(depth_map):
            depth_map = depth_map.cpu().numpy()
            
        # Apply depth-based blur
        img_np = np.array(image)
        depth_norm = blur_intensity * (1 - (depth_map - depth_map.min()) / 
                                     (depth_map.max() - depth_map.min()))
        
        result = np.zeros_like(img_np)
        for sigma in range(int(blur_intensity) + 1):
            if sigma == 0:
                continue
                
            kernel_size = 2 * int(4 * sigma + 0.5) + 1
            mask = (depth_norm >= sigma - 0.5) & (depth_norm < sigma + 0.5)
            
            if not mask.any():
                continue
            
            blurred = cv2.GaussianBlur(img_np, (kernel_size, kernel_size), sigma)
            result[mask] = blurred[mask]
        
        min_depth_mask = depth_norm > blur_intensity-0.5
        result[min_depth_mask] = img_np[min_depth_mask]
        
        return result

# Create Gradio interface
demo = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(label="Upload Image", type="numpy"),
        gr.Radio(["Gaussian Blur", "Depth-based Blur"], label="Effect Type", value="Gaussian Blur"),
        gr.Slider(minimum=1, maximum=30, value=15, label="Blur Intensity")
    ],
    outputs=gr.Image(label="Result"),
    title="Image Background Effects",
    description="""Upload an image to apply background effects:
    1. Gaussian Blur: Blurs the background while keeping the person sharp
    2. Depth-based Blur: Applies varying blur based on depth (bokeh effect)""",
    examples=[],  # You can add example images later
    cache_examples=False
)

if __name__ == "__main__":
    demo.launch()