app.py

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

# Load models
def load_models():
    # Load segmentation model
    segmenter = pipeline("image-segmentation", model="facebook/maskformer-swin-base-ade")
    
    # Load depth estimation model
    depth_estimator = pipeline("depth-estimation", model="intel/dpt-large")
    
    return segmenter, depth_estimator

# Create binary mask
def create_binary_mask(segmentation_results, image_np, target_class="person"):
    # Initialize empty mask with black background
    mask = np.zeros((image_np.shape[0], image_np.shape[1]), dtype=np.uint8)
    
    # Look for segments with target class
    found = False
    for segment in segmentation_results:
        if target_class.lower() in segment['label'].lower():
            # Convert segment mask to numpy array
            segment_mask = np.array(segment['mask'])
            # Convert grayscale to binary (255 for white)
            binary_mask = np.where(segment_mask > 0.5, 255, 0).astype(np.uint8)
            # Add to overall mask
            mask = cv2.bitwise_or(mask, binary_mask)
            found = True
    
    # If target class not found, use the largest segment
    if not found:
        largest_area = 0
        largest_mask = None
        
        for segment in segmentation_results:
            segment_mask = np.array(segment['mask'])
            binary_mask = np.where(segment_mask > 0.5, 255, 0).astype(np.uint8)
            area = np.sum(binary_mask > 0)
            
            if area > largest_area:
                largest_area = area
                largest_mask = binary_mask
        
        if largest_mask is not None:
            mask = largest_mask
    
    return mask

# Apply Gaussian blur to background
def apply_gaussian_blur_to_background(image_np, mask, sigma=15):
    # Create a blurred version of the entire image
    blurred_image = cv2.GaussianBlur(image_np, (0, 0), sigma)
    
    # Ensure mask is in correct format
    if len(mask.shape) == 3 and mask.shape[2] == 3:
        mask_gray = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
    else:
        mask_gray = mask.copy()
    
    # Normalize mask to range 0-1
    if mask_gray.max() > 1:
        mask_gray = mask_gray / 255.0
    
    # Expand mask dimensions for elementwise multiplication
    mask_3channel = np.stack([mask_gray] * 3, axis=2)
    
    # Combine original foreground with blurred background
    result = image_np * mask_3channel + blurred_image * (1 - mask_3channel)
    result = result.astype(np.uint8)
    
    return result

# Normalize depth map
def normalize_depth_map(depth_map):
    """Normalize depth map to range [0, 1]"""
    depth_min = depth_map.min()
    depth_max = depth_map.max()
    if depth_min == depth_max:
        return np.zeros_like(depth_map)
    normalized_depth = (depth_map - depth_min) / (depth_max - depth_min)
    # Explicitly invert the depth map - this is the critical fix
    return 1.0 - normalized_depth

# Apply depth-based blur
# def apply_depth_based_blur(image, depth_map, max_blur=25):
#     """Apply variable Gaussian blur based on depth with enhanced effect"""
#     # Create output image
#     result = np.zeros_like(image)
    
#     # Normalize depth map
#     normalized_depth = normalize_depth_map(depth_map)
    
#     # Enhance depth contrast to make the effect more noticeable
#     # Apply gamma correction to increase contrast between foreground and background
#     gamma = 0.5  # Values less than 1 will enhance contrast
#     normalized_depth = np.power(normalized_depth, gamma)
    
#     # Apply blur with intensity proportional to depth
#     for blur_size in range(1, max_blur + 1, 2):  # Odd numbers for kernel size
#         # Create a mask for pixels that should receive this blur level
#         if blur_size == 1:
#             mask = (normalized_depth <= blur_size / max_blur).astype(np.float32)
#         else:
#             lower_bound = (blur_size - 2) / max_blur
#             upper_bound = blur_size / max_blur
#             mask = ((normalized_depth > lower_bound) & (normalized_depth <= upper_bound)).astype(np.float32)
        
#         # Skip if no pixels in this range
#         if not np.any(mask):
#             continue
        
#         # Apply Gaussian blur with current kernel size
#         if blur_size > 1:  # No need to blur with kernel size 1
#             try:
#                 blurred = cv2.GaussianBlur(image, (blur_size, blur_size), 0)
#                 # Add blurred result to output image
#                 mask_3d = np.stack([mask] * 3, axis=2)
#                 result += (blurred * mask_3d).astype(np.uint8)
#             except Exception as e:
#                 print(f"Error applying blur with size {blur_size}: {e}")
#                 continue
#         else:
#             # For blur_size=1, just copy the original pixels
#             mask_3d = np.stack([mask] * 3, axis=2)
#             result += (image * mask_3d).astype(np.uint8)
    
#     return result

def apply_depth_based_blur(image, depth_map, max_blur=25):
    """Apply variable Gaussian blur based on depth with foreground in focus"""
    # Start with a copy of the original image
    result = image.copy().astype(float)
    
    # Normalize depth map
    normalized_depth = normalize_depth_map(depth_map)
    
    # The depth map from intel/dpt-large is already set up so closer objects 
    # have lower values. We need to define a threshold to separate foreground from background.
    foreground_threshold = 0.3  # Adjust this value based on your depth map
    
    # Create increasingly blurred versions of the image
    blurred_images = []
    blur_strengths = []
    
    # Generate progressively blurred versions
    for blur_size in range(3, max_blur + 1, 4):  # Use larger step for efficiency
        blur_strengths.append(blur_size)
        blurred = cv2.GaussianBlur(image, (blur_size, blur_size), 0)
        blurred_images.append(blurred)
    
    # Apply the appropriate blur level based on depth
    for y in range(image.shape[0]):
        for x in range(image.shape[1]):
            depth_val = normalized_depth[y, x]
            
            # Keep foreground sharp
            if depth_val <= foreground_threshold:
                continue
                
            # Apply blur based on depth - background gets more blur
            relative_depth = (depth_val - foreground_threshold) / (1 - foreground_threshold)
            blur_index = min(int(relative_depth * len(blurred_images)), len(blurred_images) - 1)
            
            # Apply the appropriate blur level
            result[y, x] = blurred_images[blur_index][y, x]
    
    return result.astype(np.uint8)


# Process function for Gradio
def process_image(input_image, blur_effect_type, blur_strength, target_class, show_depth_map=False):
    try:
        # Load models if not already loaded
        if not hasattr(process_image, "models_loaded"):
            process_image.segmenter, process_image.depth_estimator = load_models()
            process_image.models_loaded = True
        
        # Convert to numpy array
        image_np = np.array(input_image)
        
        # Process based on selected effect
        if blur_effect_type == "Gaussian Background Blur":
            # Segment the image
            segmentation_results = process_image.segmenter(input_image)
            
            # Create binary mask
            binary_mask = create_binary_mask(segmentation_results, image_np, target_class)
            
            # Apply Gaussian blur to background
            result = apply_gaussian_blur_to_background(image_np, binary_mask, sigma=blur_strength)
            
            return result
            
        elif blur_effect_type == "Depth-Based Lens Blur":
            # Resize for depth estimation
            depth_input = cv2.resize(image_np, (512, 512))
            
            # Convert to PIL image for the depth estimator
            depth_input_pil = Image.fromarray(depth_input)
            
            # Get depth map
            depth_result = process_image.depth_estimator(depth_input_pil)
            depth_map = np.array(depth_result["depth"])
            
            # If show_depth_map is True, return a visualization of the depth map
            if show_depth_map:
                # Normalize depth map for visualization
                depth_vis = normalize_depth_map(depth_map)
                # Convert to colormap for better visualization
                depth_colormap = cv2.applyColorMap((depth_vis * 255).astype(np.uint8), cv2.COLORMAP_PLASMA)
                return depth_colormap
            
            # Apply depth-based blur
            result = apply_depth_based_blur(depth_input, depth_map, max_blur=blur_strength)
            
            # Resize back to original dimensions if needed
            if image_np.shape[:2] != (512, 512):
                result = cv2.resize(result, (image_np.shape[1], image_np.shape[0]))
            
            return result
        
        else:
            return image_np  # Return original if no effect selected
    
    except Exception as e:
        print(f"Error in process_image: {e}")
        # Return the original image if there's an error
        return input_image


# Create Gradio interface
demo = gr.Blocks(title="Image Blur Effects")

with demo:
    gr.Markdown("# Image Blur Effects using Segmentation and Depth Estimation")
    gr.Markdown("Upload an image to apply different blur effects. For best results, use an image with a clear foreground subject.")
    
    with gr.Row():
        input_image = gr.Image(label="Input Image", type="pil")
        output_image = gr.Image(label="Output Image")
    
    with gr.Row():
        blur_effect_type = gr.Radio(
            ["Gaussian Background Blur", "Depth-Based Lens Blur"],
            label="Blur Effect Type",
            value="Gaussian Background Blur"
        )
        
        blur_strength = gr.Slider(
            minimum=5,
            maximum=45,
            step=2,
            value=15,
            label="Blur Strength"
        )
        
        target_class = gr.Textbox(
            label="Target Class (for segmentation)",
            value="person",
            placeholder="e.g., person, cat, dog"
        )
    
    process_btn = gr.Button("Apply Effect")
    process_btn.click(
        fn=process_image,
        inputs=[input_image, blur_effect_type, blur_strength, target_class],
        outputs=output_image
    )
    
    error_output = gr.Textbox(label="Error Information", visible=False)
    
    gr.Markdown("""
    ## How to use:
    1. Upload an image with a clear foreground subject
    2. Choose a blur effect type:
       - **Gaussian Background Blur**: Blurs the background while keeping the foreground sharp
       - **Depth-Based Lens Blur**: Creates a realistic lens blur effect based on depth estimation
    3. Adjust the blur strength
    4. For Gaussian Background Blur, specify the target class to identify the foreground (e.g., person, cat, dog)
    5. Click "Apply Effect"
    """)

# Initialize models
segmenter, depth_estimator = load_models()

# Launch the app
demo.launch(show_error = True)