app.py

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

# Dictionary of available segmentation models
SEGMENTATION_MODELS = {
    "NVIDIA SegFormer (Cityscapes)": "nvidia/segformer-b1-finetuned-cityscapes-1024-1024",
    "NVIDIA SegFormer (ADE20K)": "nvidia/segformer-b0-finetuned-ade-512-512",
    "Facebook MaskFormer (COCO)": "facebook/maskformer-swin-base-ade",
    "OneFormer (COCO)": "shi-labs/oneformer_coco_swin_large",
    "NVIDIA SegFormer (B5)": "nvidia/segformer-b5-finetuned-cityscapes-1024-1024"
}

# Dictionary of available depth estimation models
DEPTH_MODELS = {
    "Intel ZoeDepth (NYU-KITTI)": "Intel/zoedepth-nyu-kitti",
    "DPT (Large)": "Intel/dpt-large",
    "DPT (Hybrid)": "Intel/dpt-hybrid-midas",
    "GLPDepth": "vinvino02/glpn-nyu"
}

# Initialize model placeholders
segmentation_model = None
depth_estimator = None

def load_segmentation_model(model_name):
    """Load the selected segmentation model"""
    global segmentation_model
    model_path = SEGMENTATION_MODELS[model_name]
    print(f"Loading segmentation model: {model_path}...")
    segmentation_model = pipeline("image-segmentation", model=model_path)
    return f"Loaded segmentation model: {model_name}"

def load_depth_model(model_name):
    """Load the selected depth estimation model"""
    global depth_estimator
    model_path = DEPTH_MODELS[model_name]
    print(f"Loading depth estimation model: {model_path}...")
    depth_estimator = pipeline("depth-estimation", model=model_path)
    return f"Loaded depth model: {model_name}"

def lens_blur(image, radius):
    """
    Apply a more realistic lens blur (bokeh effect) using OpenCV.
    """
    if radius < 1:
        return image
    
    # Convert PIL image to OpenCV format
    img_np = np.array(image)
    
    # Create a circular kernel for the bokeh effect
    kernel_size = 2 * radius + 1
    kernel = np.zeros((kernel_size, kernel_size), dtype=np.float32)
    center = radius
    for i in range(kernel_size):
        for j in range(kernel_size):
            # Create circular kernel
            if np.sqrt((i - center) ** 2 + (j - center) ** 2) <= radius:
                kernel[i, j] = 1.0
    
    # Normalize the kernel
    if kernel.sum() != 0:
        kernel = kernel / kernel.sum()
    
    # Apply the filter to each channel separately
    channels = cv2.split(img_np)
    blurred_channels = []
    
    for channel in channels:
        blurred_channel = cv2.filter2D(channel, -1, kernel)
        blurred_channels.append(blurred_channel)
    
    # Merge the channels back
    blurred_img = cv2.merge(blurred_channels)
    
    # Convert back to PIL image
    return Image.fromarray(blurred_img)

def process_image(input_image, method, blur_intensity, blur_type):
    """
    Process the input image using one of two methods:
    
    1. Segmented Background Blur:
       - Uses segmentation to extract a foreground mask.
       - Applies the selected blur (Gaussian or Lens) to the background.
       - Composites the final image.
       
    2. Depth-based Variable Blur:
       - Uses depth estimation to generate a depth map.
       - Normalizes the depth map to be used as a blending mask.
       - Blends a fully blurred version (using the selected blur) with the original image.
       
    Returns:
       - output_image: final composited image.
       - mask_image: the mask used (binary for segmentation, normalized depth for depth-based).
    """
    # Check if models are loaded
    if segmentation_model is None or depth_estimator is None:
        return input_image, input_image.convert("L")
    
    # Ensure image is in RGB mode
    input_image = input_image.convert("RGB")
    
    # Select blur function based on blur_type
    if blur_type == "Gaussian Blur":
        blur_fn = lambda img, rad: img.filter(ImageFilter.GaussianBlur(radius=rad))
    elif blur_type == "Lens Blur":
        blur_fn = lens_blur
    else:
        blur_fn = lambda img, rad: img.filter(ImageFilter.GaussianBlur(radius=rad))
    
    if method == "Segmented Background Blur":
        # Use segmentation to obtain a foreground mask.
        results = segmentation_model(input_image)
        # Assume the last result is the main foreground object.
        foreground_mask = results[-1]["mask"]
        # Ensure the mask is grayscale.
        foreground_mask = foreground_mask.convert("L")
        # Threshold to create a binary mask.
        binary_mask = foreground_mask.point(lambda p: 255 if p > 128 else 0)
        
        # Blur the background using the selected blur function.
        blurred_background = blur_fn(input_image, blur_intensity)
        
        # Composite the final image: keep foreground and use blurred background elsewhere.
        output_image = Image.composite(input_image, blurred_background, binary_mask)
        mask_image = binary_mask
        
    elif method == "Depth-based Variable Blur":
        # Generate depth map.
        depth_results = depth_estimator(input_image)
        depth_map = depth_results["depth"]
        
        # Convert depth map to numpy array and normalize to [0, 255]
        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)
        
        # Create fully blurred version using the selected blur function.
        blurred_image = blur_fn(input_image, blur_intensity)
        
        # Convert images to arrays for blending.
        orig_np = np.array(input_image).astype(np.float32)
        blur_np = np.array(blurred_image).astype(np.float32)
        # Reshape mask for broadcasting.
        alpha = normalized_depth[..., np.newaxis] / 255.0
        
        # Blend pixels: 0 = original; 1 = fully blurred.
        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

# Build a Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Image Processing App: Segmentation & Depth-based Blur")
    
    with gr.Tab("Model Selection"):
        with gr.Row():
            with gr.Column():
                seg_model_dropdown = gr.Dropdown(
                    label="Segmentation Model",
                    choices=list(SEGMENTATION_MODELS.keys()),
                    value=list(SEGMENTATION_MODELS.keys())[0]
                )
                seg_model_load_btn = gr.Button("Load Segmentation Model")
                seg_model_status = gr.Textbox(label="Status", value="No model loaded")
            
            with gr.Column():
                depth_model_dropdown = gr.Dropdown(
                    label="Depth Estimation Model",
                    choices=list(DEPTH_MODELS.keys()),
                    value=list(DEPTH_MODELS.keys())[0]
                )
                depth_model_load_btn = gr.Button("Load Depth Model")
                depth_model_status = gr.Textbox(label="Status", value="No model loaded")
    
    with gr.Tab("Image Processing"):
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Input Image", type="pil")
                method = gr.Radio(label="Processing Method", 
                                choices=["Segmented Background Blur", "Depth-based Variable Blur"],
                                value="Segmented Background Blur")
                blur_intensity = gr.Slider(label="Blur Intensity (Maximum Blur Radius)", 
                                        minimum=1, maximum=30, step=1, value=15)
                blur_type = gr.Dropdown(label="Blur Type", 
                                        choices=["Gaussian Blur", "Lens Blur"], 
                                        value="Gaussian Blur")
                run_button = gr.Button("Process Image")
            with gr.Column():
                output_image = gr.Image(label="Output Image")
                mask_output = gr.Image(label="Mask")
    
    # Set up event handlers
    seg_model_load_btn.click(
        fn=load_segmentation_model,
        inputs=[seg_model_dropdown],
        outputs=[seg_model_status]
    )
    
    depth_model_load_btn.click(
        fn=load_depth_model,
        inputs=[depth_model_dropdown],
        outputs=[depth_model_status]
    )
    
    run_button.click(
        fn=process_image, 
        inputs=[input_image, method, blur_intensity, blur_type], 
        outputs=[output_image, mask_output]
    )

    # Load default models on startup
    demo.load(
        fn=lambda: (
            load_segmentation_model(list(SEGMENTATION_MODELS.keys())[0]),
            load_depth_model(list(DEPTH_MODELS.keys())[0])
        ),
        inputs=None,
        outputs=[seg_model_status, depth_model_status]
    )

# Launch the app
demo.launch()