app.py

from unicodedata import normalize

import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageFilter
from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
from scipy.ndimage import gaussian_filter
import torch
import requests
from io import BytesIO
import cv2
import warnings
warnings.filterwarnings('ignore')
from transformers import DPTImageProcessor, DPTForDepthEstimation,  AutoImageProcessor

model_cache = {
    "seg_name": None, "seg_proc": None, "seg_model": None,
    "depth_name": None, "depth_proc": None, "depth_model": None
}

MODEL_CONFIG = {
    "segmentation": {
        "Segformer (B0)": "nvidia/segformer-b0-finetuned-ade-512-512",
        "Segformer (B5)": "nvidia/segformer-b5-finetuned-ade-640-640",
    },
    "depth": {
        "DPT-Large": "Intel/dpt-large",
        "Facebook-DPT-Dinov2": "facebook/dpt-dinov2-small-nyu",
    }
}

def get_seg_model(model_name):
    global model_cache
    repo_id = MODEL_CONFIG["segmentation"][model_name]
    if model_cache["seg_name"] != model_name:
        print(f"Switching segmentation model to {model_name}...")
        model_cache["seg_proc"] = SegformerImageProcessor.from_pretrained(repo_id)
        model_cache["seg_model"] = SegformerForSemanticSegmentation.from_pretrained(repo_id)
        model_cache["seg_name"] = model_name
    return model_cache["seg_proc"], model_cache["seg_model"]

def get_depth_model(model_name):
    global model_cache
    repo_id = MODEL_CONFIG["depth"][model_name]
    if model_cache["depth_name"] != model_name:
        print(f"Switching depth model to {model_name}...")
        model_cache["depth_proc"] = DPTImageProcessor.from_pretrained(repo_id)
        model_cache["depth_model"] = DPTForDepthEstimation.from_pretrained(repo_id)
        model_cache["depth_name"] = model_name
    return model_cache["depth_proc"], model_cache["depth_model"]

def preprocess_image(image, target_size=512):
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    if image.mode != 'RGB':
        image = image.convert('RGB')
    return image.resize((target_size, target_size), Image.Resampling.LANCZOS)

def segment_human(image, processor, model):
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    upsampled = torch.nn.functional.interpolate(
        outputs.logits, size=(512, 512), mode="bilinear", align_corners=False
    )
    pred_seg = upsampled.argmax(dim=1)[0].cpu().numpy()
    # Note: Label 12 is 'person' in ADE20k dataset
    return (pred_seg == 12).astype(np.uint8) * 255

def apply_background_blur(image, mask, sigma=15):
    img_array = np.array(image).astype(np.float32)
    mask_normalized = mask.astype(np.float32) / 255.0
    mask_smooth = gaussian_filter(mask_normalized, sigma=2)
    mask_smooth = np.clip(mask_smooth, 0, 1)
    
    blurred_array = np.zeros_like(img_array)
    for i in range(3):
        blurred_array[:, :, i] = gaussian_filter(img_array[:, :, i], sigma=sigma)
    
    mask_3d = np.stack([mask_smooth] * 3, axis=2)
    result = (img_array * mask_3d + blurred_array * (1 - mask_3d)).astype(np.uint8)
    return Image.fromarray(result)

def estimate_depth(image, processor, model, invert):
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    prediction = torch.nn.functional.interpolate(
        outputs.predicted_depth.unsqueeze(1), size=(512, 512), mode="bicubic", align_corners=False,
    )
    depth_map = prediction.squeeze().cpu().numpy()
    depth_min, depth_max = depth_map.min(), depth_map.max()
    normalized = (depth_map - depth_min) / (depth_max - depth_min)
    if invert == True:
        normalized = 1.0 - normalized
    return normalized * 15.0

def apply_lens_blur(image, depth_map, max_sigma=15):
    img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR).astype(np.float32)
    
    # Create blur pyramid
    num_levels = 10
    blur_pyramid = []
    
    for i in range(num_levels):
        sigma = (i / (num_levels - 1)) * max_sigma
        if sigma < 0.5:
            blur_pyramid.append(img_cv.copy())
        else:
            ksize = int(2 * np.ceil(3 * sigma) + 1)
            if ksize % 2 == 0:
                ksize += 1
            blurred = cv2.GaussianBlur(img_cv, (ksize, ksize), sigma)
            blur_pyramid.append(blurred)
    
    # Apply variable blur based on depth
    depth_norm = depth_map / 15.0
    output = np.zeros_like(img_cv)
    
    depth_scaled = depth_norm * (num_levels - 1)
    level_low = np.floor(depth_scaled).astype(np.int32)
    level_high = np.ceil(depth_scaled).astype(np.int32)
    level_low = np.clip(level_low, 0, num_levels - 1)
    level_high = np.clip(level_high, 0, num_levels - 1)
    
    weight = depth_scaled - level_low
    weight = np.expand_dims(weight, axis=2)
    
    for y in range(img_cv.shape[0]):
        for x in range(img_cv.shape[1]):
            ll = level_low[y, x]
            lh = level_high[y, x]
            w = weight[y, x, 0]
            
            if ll == lh:
                output[y, x] = blur_pyramid[ll][y, x]
            else:
                output[y, x] = (1 - w) * blur_pyramid[ll][y, x] + w * blur_pyramid[lh][y, x]
    
    output = np.clip(output, 0, 255).astype(np.uint8)
    output_rgb = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
    return Image.fromarray(output_rgb)

def process_gaussian_blur(image, sigma, model_choice):
    if image is None: return None, "Upload an image!"
    try:
        proc, model = get_seg_model(model_choice)
        img = preprocess_image(image)
        mask = segment_human(img, proc, model)
        result = apply_background_blur(img, mask, sigma)
        return result, f"Applied {model_choice} with σ={sigma}"
    except Exception as e:
        return None, f"Error: {str(e)}"

def process_lens_blur(image, max_sigma, model_choice):
    if image is None: return None, None, "Upload an image!"
    try:
        proc, model = get_depth_model(model_choice)
        if model_choice == "Facebook-DPT-Dinov2":
            invert = False
        else:            
            invert = True
        img = preprocess_image(image)
        depth = estimate_depth(img, proc, model, invert)
        result = apply_lens_blur(img, depth, max_sigma)
        
        depth_vis = cv2.applyColorMap(((depth / 15.0) * 255).astype(np.uint8), cv2.COLORMAP_MAGMA)
        return result, Image.fromarray(cv2.cvtColor(depth_vis, cv2.COLOR_BGR2RGB)), f"Applied {model_choice}"
    except Exception as e:
        return None, None, f"Error: {str(e)}"

with gr.Blocks(title="AI Blur Studio", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# AI Blur Studio\nSelect your AI models and adjust blur intensity.")
    
    with gr.Tabs():
        with gr.Tab("📹 Gaussian Background Blur"):
            with gr.Row():
                with gr.Column():
                    gaussian_input = gr.Image(label="Input Image")
                    seg_model_dropdown = gr.Dropdown(
                        choices=list(MODEL_CONFIG["segmentation"].keys()),
                        value=list(MODEL_CONFIG["segmentation"].keys())[0],
                        label="Segmentation Model"
                    )
                    gaussian_sigma = gr.Slider(0, 30, 15, label="Blur σ")
                    gaussian_btn = gr.Button("Process", variant="primary")
                with gr.Column():
                    gaussian_output = gr.Image(label="Result")
                    gaussian_status = gr.Textbox(label="Status")

        with gr.Tab("📸 Depth-Based Lens Blur"):
            with gr.Row():
                with gr.Column():
                    lens_input = gr.Image(label="Input Image")
                    depth_model_dropdown = gr.Dropdown(
                        choices=list(MODEL_CONFIG["depth"].keys()),
                        value=list(MODEL_CONFIG["depth"].keys())[0],
                        label="Depth Estimation Model"
                    )
                    lens_sigma = gr.Slider(0, 25, 15, label="Max σ")
                    lens_btn = gr.Button("Process", variant="primary")
                with gr.Column():
                    lens_output = gr.Image(label="Blurred Result")
                    lens_depth = gr.Image(label="Depth Map")
                    lens_status = gr.Textbox(label="Status")

    gaussian_btn.click(process_gaussian_blur, [gaussian_input, gaussian_sigma, seg_model_dropdown], [gaussian_output, gaussian_status])
    lens_btn.click(process_lens_blur, [lens_input, lens_sigma, depth_model_dropdown], [lens_output, lens_depth, lens_status])

if __name__ == "__main__":
    demo.launch(share=True)