Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| import cv2 | |
| from transformers import pipeline | |
| import os | |
| # Load models | |
| def load_models(): | |
| # Load segmentation model | |
| segmenter = pipeline("image-segmentation", model="facebook/maskformer-swin-base-ade") | |
| # Load depth estimation model | |
| depth_estimator = pipeline("depth-estimation", model="intel/dpt-large") | |
| return segmenter, depth_estimator | |
| # Create binary mask | |
| def create_binary_mask(segmentation_results, image_np, target_class="person"): | |
| # Initialize empty mask with black background | |
| mask = np.zeros((image_np.shape[0], image_np.shape[1]), dtype=np.uint8) | |
| # Look for segments with target class | |
| found = False | |
| for segment in segmentation_results: | |
| if target_class.lower() in segment['label'].lower(): | |
| # Convert segment mask to numpy array | |
| segment_mask = np.array(segment['mask']) | |
| # Convert grayscale to binary (255 for white) | |
| binary_mask = np.where(segment_mask > 0.5, 255, 0).astype(np.uint8) | |
| # Add to overall mask | |
| mask = cv2.bitwise_or(mask, binary_mask) | |
| found = True | |
| # If target class not found, use the largest segment | |
| if not found: | |
| largest_area = 0 | |
| largest_mask = None | |
| for segment in segmentation_results: | |
| segment_mask = np.array(segment['mask']) | |
| binary_mask = np.where(segment_mask > 0.5, 255, 0).astype(np.uint8) | |
| area = np.sum(binary_mask > 0) | |
| if area > largest_area: | |
| largest_area = area | |
| largest_mask = binary_mask | |
| if largest_mask is not None: | |
| mask = largest_mask | |
| return mask | |
| # Apply Gaussian blur to background | |
| def apply_gaussian_blur_to_background(image_np, mask, sigma=15): | |
| # Create a blurred version of the entire image | |
| blurred_image = cv2.GaussianBlur(image_np, (0, 0), sigma) | |
| # Ensure mask is in correct format | |
| if len(mask.shape) == 3 and mask.shape[2] == 3: | |
| mask_gray = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY) | |
| else: | |
| mask_gray = mask.copy() | |
| # Normalize mask to range 0-1 | |
| if mask_gray.max() > 1: | |
| mask_gray = mask_gray / 255.0 | |
| # Expand mask dimensions for elementwise multiplication | |
| mask_3channel = np.stack([mask_gray] * 3, axis=2) | |
| # Combine original foreground with blurred background | |
| result = image_np * mask_3channel + blurred_image * (1 - mask_3channel) | |
| result = result.astype(np.uint8) | |
| return result | |
| # Normalize depth map | |
| def normalize_depth_map(depth_map): | |
| depth_min = depth_map.min() | |
| depth_max = depth_map.max() | |
| normalized_depth = (depth_map - depth_min) / (depth_max - depth_min) | |
| return normalized_depth | |
| # Apply depth-based blur | |
| def apply_depth_based_blur(image, depth_map, max_blur=25): | |
| # Create output image | |
| result = np.zeros_like(image) | |
| # Apply blur with intensity proportional to depth | |
| for blur_size in range(1, max_blur + 1, 2): # Odd numbers for kernel size | |
| # Create a mask for pixels that should receive this blur level | |
| if blur_size == 1: | |
| mask = (depth_map <= blur_size / max_blur).astype(np.float32) | |
| else: | |
| lower_bound = (blur_size - 2) / max_blur | |
| upper_bound = blur_size / max_blur | |
| mask = ((depth_map > lower_bound) & (depth_map <= upper_bound)).astype(np.float32) | |
| # Skip if no pixels in this range | |
| if not np.any(mask): | |
| continue | |
| # Apply Gaussian blur with current kernel size | |
| if blur_size > 1: | |
| blurred = cv2.GaussianBlur(image, (blur_size, blur_size), 0) | |
| mask_3d = np.stack([mask] * 3, axis=2) | |
| result += (blurred * mask_3d).astype(np.uint8) | |
| else: | |
| mask_3d = np.stack([mask] * 3, axis=2) | |
| result += (image * mask_3d).astype(np.uint8) | |
| return result | |
| # Process function for Gradio | |
| def process_image(input_image, blur_effect_type, blur_strength, target_class): | |
| # Load models if not already loaded | |
| if not hasattr(process_image, "models_loaded"): | |
| process_image.segmenter, process_image.depth_estimator = load_models() | |
| process_image.models_loaded = True | |
| # Convert to numpy array | |
| image_np = np.array(input_image) | |
| # Process based on selected effect | |
| if blur_effect_type == "Gaussian Background Blur": | |
| # Segment the image | |
| segmentation_results = process_image.segmenter(input_image) | |
| # Create binary mask | |
| binary_mask = create_binary_mask(segmentation_results, image_np, target_class) | |
| # Apply Gaussian blur to background | |
| result = apply_gaussian_blur_to_background(image_np, binary_mask, sigma=blur_strength) | |
| return result | |
| elif blur_effect_type == "Depth-Based Lens Blur": | |
| # Resize for depth estimation | |
| depth_input = cv2.resize(image_np, (512, 512)) | |
| # Get depth map | |
| depth_result = process_image.depth_estimator(depth_input) | |
| depth_map = np.array(depth_result["depth"]) | |
| # Normalize depth map | |
| normalized_depth = normalize_depth_map(depth_map) | |
| # Apply depth-based blur | |
| result = apply_depth_based_blur(depth_input, normalized_depth, max_blur=blur_strength) | |
| # Resize back to original dimensions if needed | |
| if image_np.shape[:2] != (512, 512): | |
| result = cv2.resize(result, (image_np.shape[1], image_np.shape[0])) | |
| return result | |
| else: | |
| return image_np # Return original if no effect selected | |
| # Create Gradio interface | |
| demo = gr.Blocks(title="Image Blur Effects") | |
| with demo: | |
| gr.Markdown("# Image Blur Effects using Segmentation and Depth Estimation") | |
| gr.Markdown("Upload an image to apply different blur effects. For best results, use an image with a clear foreground subject.") | |
| with gr.Row(): | |
| input_image = gr.Image(label="Input Image", type="pil") | |
| output_image = gr.Image(label="Output Image") | |
| with gr.Row(): | |
| blur_effect_type = gr.Radio( | |
| ["Gaussian Background Blur", "Depth-Based Lens Blur"], | |
| label="Blur Effect Type", | |
| value="Gaussian Background Blur" | |
| ) | |
| blur_strength = gr.Slider( | |
| minimum=5, | |
| maximum=45, | |
| step=2, | |
| value=15, | |
| label="Blur Strength" | |
| ) | |
| target_class = gr.Textbox( | |
| label="Target Class (for segmentation)", | |
| value="person", | |
| placeholder="e.g., person, cat, dog" | |
| ) | |
| process_btn = gr.Button("Apply Effect") | |
| process_btn.click( | |
| fn=process_image, | |
| inputs=[input_image, blur_effect_type, blur_strength, target_class], | |
| outputs=output_image | |
| ) | |
| gr.Markdown(""" | |
| ## How to use: | |
| 1. Upload an image with a clear foreground subject | |
| 2. Choose a blur effect type: | |
| - **Gaussian Background Blur**: Blurs the background while keeping the foreground sharp | |
| - **Depth-Based Lens Blur**: Creates a realistic lens blur effect based on depth estimation | |
| 3. Adjust the blur strength | |
| 4. For Gaussian Background Blur, specify the target class to identify the foreground (e.g., person, cat, dog) | |
| 5. Click "Apply Effect" | |
| """) | |
| # Initialize models | |
| segmenter, depth_estimator = load_models() | |
| # Launch the app | |
| demo.launch() | |