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)