app.py

import gradio as gr
from transformers import AutoImageProcessor, AutoModelForDepthEstimation, AutoModelForImageSegmentation
import torch
import numpy as np
from PIL import Image, ImageFilter
import cv2
from torchvision.transforms.functional import normalize
import torch.nn.functional as F

# Load models
print("Loading models...")
depth_processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Base-hf")
depth_model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Base-hf")
seg_model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4", trust_remote_code=True)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
depth_model.to(device)
seg_model.to(device)
print("Models loaded!")

def depth_based_blur(image):
    image = image.resize((512, 512))
    inputs = depth_processor(images=image, 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_normalized = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min()) * 15
    
    img_array = np.array(image)
    blurred_image = np.zeros_like(img_array)
    
    for i in range(512):
        for j in range(512):
            depth_value = depth_normalized[i, j]
            kernel_size = int(depth_value * 2) | 1
            kernel_size = max(1, kernel_size)
            
            half_k = kernel_size // 2
            i_start, i_end = max(0, i-half_k), min(512, i+half_k+1)
            j_start, j_end = max(0, j-half_k), min(512, j+half_k+1)
            
            patch = img_array[i_start:i_end, j_start:j_end]
            blurred_patch = cv2.GaussianBlur(patch, (kernel_size, kernel_size), 0)
            blurred_image[i, j] = blurred_patch[i-i_start, j-j_start]
    
    return Image.fromarray(blurred_image)

def preprocess_image_seg(im: np.ndarray, model_input_size: list) -> torch.Tensor:
    if len(im.shape) < 3:
        im = im[:, :, np.newaxis]
    im_tensor = torch.tensor(im, dtype=torch.float32).permute(2, 0, 1)
    im_tensor = F.interpolate(torch.unsqueeze(im_tensor, 0), size=model_input_size, mode='bilinear')
    image = torch.divide(im_tensor, 255.0)
    image = normalize(image, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
    return image

def postprocess_image_seg(result: torch.Tensor, im_size: list) -> np.ndarray:
    result = torch.squeeze(F.interpolate(result, size=im_size, mode='bilinear'), 0)
    ma = torch.max(result)
    mi = torch.min(result)
    result = (result - mi) / (ma - mi)
    im_array = (result * 255).permute(1, 2, 0).cpu().data.numpy().astype(np.uint8)
    return np.squeeze(im_array)

def gaussian_blur_background(image, blur_radius):
    image = image.resize((512, 512))
    
    # Generate segmentation mask
    orig_im = np.array(image)
    orig_im_size = orig_im.shape[0:2]
    model_input = preprocess_image_seg(orig_im, [1024, 1024]).to(device)
    
    with torch.no_grad():
        result = seg_model(model_input)
    
    result_image = postprocess_image_seg(result[0][0], orig_im_size)
    binary_mask = (result_image > 128).astype(np.uint8) * 255
    binary_mask_pil = Image.fromarray(binary_mask).convert('L')
    
    # Apply blur: where mask is white (human) use original, where black (background) use blurred
    blurred_image = image.filter(ImageFilter.GaussianBlur(radius=blur_radius))
    output = Image.composite(image, blurred_image, binary_mask_pil)
    
    return output

# Gradio interface
with gr.Blocks(title="Image Blur Effects") as demo:
    gr.Markdown("# Image Blur Effects")
    gr.Markdown("Apply depth-based lens blur or Gaussian background blur to your images")
    
    with gr.Tab("Depth-Based Lens Blur"):
        gr.Markdown("### Simulates camera depth of field - farther objects blur more")
        with gr.Row():
            input_depth = gr.Image(type="pil", label="Input Image")
            output_depth = gr.Image(type="pil", label="Depth-Based Blur Output")
        btn_depth = gr.Button("Apply Depth Blur")
        btn_depth.click(fn=depth_based_blur, inputs=input_depth, outputs=output_depth)
    
    with gr.Tab("Gaussian Background Blur"):
        gr.Markdown("### Keeps human sharp, blurs background uniformly")
        with gr.Row():
            with gr.Column():
                input_gaussian = gr.Image(type="pil", label="Input Image")
                blur_slider = gr.Slider(minimum=1, maximum=30, value=15, step=1, label="Blur Radius (σ)")
            output_gaussian = gr.Image(type="pil", label="Blurred Background Output")
        btn_gaussian = gr.Button("Apply Background Blur")
        btn_gaussian.click(fn=gaussian_blur_background, inputs=[input_gaussian, blur_slider], outputs=output_gaussian)

demo.launch()