Spaces:
Sleeping
Sleeping
| 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) |