#!/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()