app.py

import gradio as gr
from transformers import OneFormerProcessor, OneFormerForUniversalSegmentation, AutoImageProcessor, AutoModelForDepthEstimation
from PIL import Image, ImageFilter
import numpy as np
import torch
from scipy.ndimage import gaussian_filter
import cv2

# Load the OneFormer processor and model globally
oneformer_processor = None
oneformer_model = None
try:
    oneformer_processor = OneFormerProcessor.from_pretrained("shi-labs/oneformer_coco_swin_large")
    oneformer_model = OneFormerForUniversalSegmentation.from_pretrained("shi-labs/oneformer_coco_swin_large")
except Exception as e:
    print(f"Error loading OneFormer model: {e}")

# Load the Depth Estimation processor and model globally
depth_processor = None
depth_model = None
try:
    depth_processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
    depth_model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
except Exception as e:
    print(f"Error loading Depth Anything model: {e}")

def apply_gaussian_blur_background(image, mask, radius):
    """Applies Gaussian blur to the background of the image."""
    img_array = np.array(image)
    
    # Apply Gaussian blur to the entire image
    blurred_background = image.filter(ImageFilter.GaussianBlur(radius=radius))
    blurred_array = np.array(blurred_background)

    # Ensure the mask is in the same size and format as the image
    foreground_mask = np.array(mask) > 0
    foreground_mask_3d = np.stack([foreground_mask] * 3, axis=-1)  # Ensure it's 3 channels (RGB)
    final_image_array = np.where(foreground_mask_3d, img_array, blurred_array)

    return Image.fromarray(final_image_array.astype(np.uint8))

def apply_depth_based_blur_background(image, mask, strength):
    """Applies lens blur to the background of the image based on depth estimation."""
    resized_image = image.resize((512, 512))
    image_np = np.array(resized_image)

    if depth_processor is None or depth_model is None:
        return "Error: Depth Anything model not loaded."

    # Prepare image for the depth estimation model
    inputs = depth_processor(images=resized_image, return_tensors="pt")

    with torch.no_grad():
        outputs = depth_model(**inputs)
        predicted_depth = outputs.predicted_depth

    # Interpolate depth map to the resized image size
    prediction = torch.nn.functional.interpolate(
        predicted_depth.unsqueeze(1),
        size=resized_image.size[::-1],
        mode="bicubic",
        align_corners=False,
    ).squeeze().cpu().numpy()

    # Normalize the depth map to the range 0-1
    depth_norm = (prediction - np.min(prediction)) / (np.max(prediction) - np.min(prediction))

    num_blur_levels = 5
    blurred_layers = []
    for i in range(num_blur_levels):
        sigma = i * (strength / 5)  # Adjust sigma based on strength
        if sigma == 0:
            blurred = image_np
        else:
            blurred = cv2.GaussianBlur(image_np, (15, 15), sigmaX=sigma, sigmaY=sigma, borderType=cv2.BORDER_REPLICATE)
        blurred_layers.append(blurred)

    depth_indices = ((1 - depth_norm) * (num_blur_levels - 1)).astype(np.uint8)

    final_blurred_image_resized = np.zeros_like(image_np)
    for y in range(image_np.shape[0]):
        for x in range(image_np.shape[1]):
            depth_index = depth_indices[y, x]
            final_blurred_image_resized[y, x] = blurred_layers[depth_index][y, x]

    final_blurred_pil_resized = Image.fromarray(final_blurred_image_resized.astype(np.uint8))
    final_blurred_pil = final_blurred_pil_resized.resize(image.size)
    final_blurred_array = np.array(final_blurred_pil)
    original_array = np.array(image)
    mask_resized = mask.resize(image.size)
    mask_array = np.array(mask_resized) > 0
    mask_array_3d = np.stack([mask_array] * 3, axis=-1)

    # Apply the mask to combine the original foreground with the blurred background
    final_output_array = np.where(mask_array_3d, original_array, final_blurred_array)
    return Image.fromarray(final_output_array.astype(np.uint8))


def segment_and_blur(input_image, blur_type, gaussian_radius=15, lens_strength=5):
    try:
        if oneformer_processor is None or oneformer_model is None:
            return "Error: OneFormer model not loaded."  # <-- Fix indentation here

        image = input_image.convert("RGB")
        image = image.rotate(-90, expand=True)

        inputs = oneformer_processor(images=image, task_inputs=["semantic"], return_tensors="pt")
        with torch.no_grad():
            outputs = oneformer_model(**inputs)

        predicted_semantic_map = oneformer_processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
        segmentation_mask = predicted_semantic_map.cpu().numpy()

        id2label = oneformer_model.config.id2label
        print(id2label)  # <-- Add this to debug

        foreground_label = 'person'
        foreground_class_id = None
        for id, label in id2label.items():
            if label.lower() == foreground_label.lower():
                foreground_class_id = id
                break

        if foreground_class_id is None:
            return f"Error: Could not find the label '{foreground_label}' in the model's class mapping."

        output_mask_array = np.zeros(segmentation_mask.shape, dtype=np.uint8)
        output_mask_array[segmentation_mask == foreground_class_id] = 255
        mask_pil = Image.fromarray(output_mask_array, mode='L')

        if blur_type == "Gaussian":
            blurred_image = apply_gaussian_blur_background(image, mask_pil, gaussian_radius)
        elif blur_type == "Lens":
            blurred_image = apply_depth_based_blur_background(image, mask_pil, lens_strength)
        else:
            return "Error: Invalid blur type selected."

        return blurred_image
    except Exception as e:
        return f"Error during processing: {str(e)}"



iface = gr.Interface(
    fn=segment_and_blur,
    inputs=[
        gr.Image(type="pil", label="Input Image"),  # <-- ADD type="pil" here
        gr.Radio(["Gaussian", "Lens"], label="Blur Type", value="Gaussian"),
        gr.Slider(0, 30, step=1, value=15, label="Gaussian Blur Radius"),
        gr.Slider(0, 10, step=1, value=5, label="Lens Blur Strength"),
    ],
    outputs=gr.Image(label="Output Image"),
    title="Image Background Blur App",
    description="Upload an image, select a blur type (Gaussian or Lens), and adjust the blur parameters to blur the background while keeping the person in focus."
)


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