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."""
    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

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."""
    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()
    depth_map = (depth_map - depth_min) / (depth_max - depth_min)
    depth_map = 1 - depth_map  # Invert: higher values = farther
    return depth_map

def apply_depth_based_blur(image: Image.Image, depth_map: np.ndarray, max_radius: float = 30, foreground_percentile: float = 5.0):
    """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=30, foreground_percentile=5.0):
    """Processes the image based on the selected blur type and parameters."""
    if image is None:
        return None, "Please upload an image."

    image = Image.fromarray(image).convert("RGB")

    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})"

    return output_image, title

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Image Blur Effects with Gaussian and Lens Blur")
    gr.Markdown("Upload an image and choose a blur effect. Adjust the parameters to customize the blur intensity.")

    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 (for Gaussian Blur)")
            max_radius = gr.Slider(minimum=1, maximum=50, step=1, value=15, label="Max Blur Radius (for Lens Blur)")
            foreground_percentile = gr.Slider(minimum=1, maximum=50, step=1, value=30, label="Foreground Percentile (for Lens Blur)")

    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()