app.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Mar 27 13:56:42 2025

@author: perghect
"""
import gradio as gr
import requests
import io
import torch
import numpy as np
from PIL import Image, ImageFilter
from torchvision import transforms
from transformers import AutoModelForImageSegmentation, AutoImageProcessor, AutoModelForDepthEstimation

# Set device and precision
device = 'cuda' if torch.cuda.is_available() else 'cpu'
torch.set_float32_matmul_precision('high')

# Load models at startup
rmbg_model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-2.0", trust_remote_code=True).to(device).eval()
depth_processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
depth_model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf").to(device)

def load_image_from_link(url: str) -> Image.Image:
    """Downloads an image from a URL and returns a Pillow Image."""
    response = requests.get(url)
    response.raise_for_status()
    image = Image.open(io.BytesIO(response.content)).convert("RGB")
    return image

# Gaussian Blur Functions
def run_rmbg(image: Image.Image, threshold=0.5):
    """Runs the RMBG-2.0 model on the image and returns a binary mask."""
    try:
        image_size = (1024, 1024)
        transform_image = transforms.Compose([
            transforms.Resize(image_size),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

        input_images = transform_image(image).unsqueeze(0).to(device)

        with torch.no_grad():
            preds = rmbg_model(input_images)
            if isinstance(preds, list):
                mask_logits = preds[-1]
            else:
                raise ValueError(f"Unexpected output format: {type(preds)}")

        mask_prob = mask_logits.sigmoid().cpu()[0].squeeze()
        pred_pil = transforms.ToPILImage()(mask_prob)
        mask_pil = pred_pil.resize(image.size, resample=Image.BILINEAR)

        mask_np = np.array(mask_pil, dtype=np.uint8) / 255.0
        binary_mask = (mask_np > threshold).astype(np.uint8)
        return binary_mask
    except Exception as e:
        raise Exception(f"Error in background removal: {str(e)}")

def apply_background_blur(image: Image.Image, mask: np.ndarray, sigma: float = 15):
    """Applies a Gaussian blur to the background while keeping the foreground sharp."""
    image_np = np.array(image)
    mask_np = mask.astype(np.uint8)

    blurred_image = image.filter(ImageFilter.GaussianBlur(radius=sigma))
    blurred_np = np.array(blurred_image)

    output_np = np.where(mask_np[..., None] == 1, image_np, blurred_np)
    output_image = Image.fromarray(output_np.astype(np.uint8))
    return output_image

# Lens Blur Functions
def run_depth_estimation(image: Image.Image, target_size=(512, 512)):
    """Runs the Depth-Anything-V2-Small model and returns the depth map."""
    try:
        image_resized = image.resize(target_size, resample=Image.BILINEAR)
        inputs = depth_processor(images=image_resized, return_tensors="pt").to(device)

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

        prediction = torch.nn.functional.interpolate(
            predicted_depth.unsqueeze(1),
            size=image.size[::-1],
            mode="bicubic",
            align_corners=False,
        )

        depth_map = prediction.squeeze().cpu().numpy()
        depth_max = depth_map.max()
        depth_min = depth_map.min()
        if depth_max == depth_min:
            depth_max = depth_min + 1e-6  # Avoid division by zero
        depth_map = (depth_map - depth_min) / (depth_max - depth_min)
        depth_map = 1 - depth_map  # Invert: higher values = farther
        return depth_map
    except Exception as e:
        raise Exception(f"Error in depth estimation: {str(e)}")

def apply_depth_based_blur(image: Image.Image, depth_map: np.ndarray, max_radius: float = 15, foreground_percentile: float = 30):
    """Applies a variable Gaussian blur based on the depth map."""
    image_np = np.array(image)

    if depth_map.shape != image_np.shape[:2]:
        depth_map = np.array(Image.fromarray(depth_map).resize(image.size, resample=Image.BILINEAR))

    foreground_threshold = np.percentile(depth_map.flatten(), foreground_percentile)

    output_np = np.zeros_like(image_np)
    mask_foreground = (depth_map <= foreground_threshold)
    output_np[mask_foreground] = image_np[mask_foreground]

    depth_max = depth_map.max()
    depth_range = depth_max - foreground_threshold
    if depth_range == 0:
        depth_range = 1e-6
    normalized_depth = np.zeros_like(depth_map)
    mask_above_foreground = (depth_map > foreground_threshold)
    normalized_depth[mask_above_foreground] = (depth_map[mask_above_foreground] - foreground_threshold) / depth_range
    normalized_depth = np.clip(normalized_depth, 0, 1)

    depth_levels = np.linspace(0, 1, 20)
    for i in range(len(depth_levels) - 1):
        depth_min = depth_levels[i]
        depth_max = depth_levels[i + 1]
        mask = (normalized_depth >= depth_min) & (normalized_depth < depth_max) & (depth_map > foreground_threshold)
        if not np.any(mask):
            continue

        avg_depth = (depth_min + depth_max) / 2
        blur_radius = max_radius * avg_depth

        blurred_image = image.filter(ImageFilter.GaussianBlur(radius=blur_radius))
        blurred_np = np.array(blurred_image)
        output_np[mask] = blurred_np[mask]

    mask_farthest = (normalized_depth >= depth_levels[-1]) & (depth_map > foreground_threshold)
    if np.any(mask_farthest):
        blurred_max = image.filter(ImageFilter.GaussianBlur(radius=max_radius))
        output_np[mask_farthest] = np.array(blurred_max)[mask_farthest]

    output_image = Image.fromarray(output_np.astype(np.uint8))
    return output_image

# Main Processing Function for Gradio
def process_image(image, blur_type, sigma=15, max_radius=15, foreground_percentile=30):
    """Processes the image based on the selected blur type."""
    if image is None:
        return None, "Please upload an image."

    try:
        image = Image.fromarray(image).convert("RGB")
    except Exception as e:
        return None, f"Error processing image: {str(e)}"

    # Resize image if too large
    max_size = (1024, 1024)
    if image.size[0] > max_size[0] or image.size[1] > max_size[1]:
        image.thumbnail(max_size, Image.Resampling.LANCZOS)

    try:
        if blur_type == "Gaussian Blur":
            mask = run_rmbg(image, threshold=0.5)
            output_image = apply_background_blur(image, mask, sigma=sigma)
            title = f"Gaussian Blur (sigma={sigma})"
        else:  # Lens Blur
            depth_map = run_depth_estimation(image, target_size=(512, 512))
            output_image = apply_depth_based_blur(image, depth_map, max_radius=max_radius, foreground_percentile=foreground_percentile)
            title = f"Lens Blur (max_radius={max_radius}, foreground_percentile={foreground_percentile})"
    except Exception as e:
        return None, f"Error applying blur: {str(e)}"

    return output_image, title

# Gradio Interface with Conditional Parameter Display
with gr.Blocks() as demo:
    gr.Markdown("# Image Blur Effects with Gaussian and Lens Blur")
    gr.Markdown("""
    This app applies blur effects to your images. Follow these steps to use it:

    **Note**: This app is hosted on Hugging Face Spaces’ free tier and may go to "Sleeping" mode after 48 hours of inactivity. If it doesn’t load immediately, please wait a few seconds while it wakes up.

    1. **Upload an Image**: Click the "Upload Image" box to upload an image from your device.
    2. **Choose a Blur Type**:
       - **Gaussian Blur**: Applies a uniform blur to the background, keeping the foreground sharp. Adjust the sigma parameter to control blur intensity.
       - **Lens Blur**: Applies a depth-based blur, simulating a depth-of-field effect (closer objects are sharp, farther objects are blurred). Adjust the max radius and foreground percentile to fine-tune the effect.
    3. **Adjust Parameters**:
       - For Gaussian Blur, use the "Gaussian Blur Sigma" slider to control blur intensity (higher values = more blur).
       - For Lens Blur, use the "Max Blur Radius" slider to control the maximum blur intensity and the "Foreground Percentile" slider to adjust the depth threshold for the foreground.
    4. **Apply the Blur**: Click the "Apply Blur" button to process the image.
    5. **View the Result**: The processed image will appear in the "Output Image" box, along with a description of the effect applied.

    **Example**: Try uploading an image with a clear foreground and background (e.g., a person in front of a landscape) to see the effects in action.
    """)

    with gr.Row():
        image_input = gr.Image(label="Upload Image", type="numpy")
        with gr.Column():
            blur_type = gr.Radio(choices=["Gaussian Blur", "Lens Blur"], label="Blur Type", value="Gaussian Blur")
            sigma = gr.Slider(minimum=1, maximum=50, step=1, value=15, label="Gaussian Blur Sigma", visible=True)
            max_radius = gr.Slider(minimum=1, maximum=50, step=1, value=15, label="Max Lens Blur Radius", visible=False)
            foreground_percentile = gr.Slider(minimum=1, maximum=50, step=1, value=30, label="Foreground Percentile", visible=False)

    # Update visibility of parameters based on blur type
    def update_visibility(blur_type):
        if blur_type == "Gaussian Blur":
            return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
        else:  # Lens Blur
            return gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)

    blur_type.change(
        fn=update_visibility,
        inputs=blur_type,
        outputs=[sigma, max_radius, foreground_percentile]
    )

    process_button = gr.Button("Apply Blur")
    with gr.Row():
        output_image = gr.Image(label="Output Image")
        output_text = gr.Textbox(label="Effect Applied")

    process_button.click(
        fn=process_image,
        inputs=[image_input, blur_type, sigma, max_radius, foreground_percentile],
        outputs=[output_image, output_text]
    )

# Launch the app
demo.launch()