Spaces:
Running
Running
| import gradio as gr | |
| from PIL import Image, ImageFilter | |
| import numpy as np | |
| import cv2 | |
| import torch | |
| from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation, DPTFeatureExtractor, DPTForDepthEstimation | |
| # Load models | |
| segformer_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b1-finetuned-ade-512-512") | |
| segformer_model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b1-finetuned-ade-512-512") | |
| dpt_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large") | |
| dpt_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large") | |
| # Gaussian Blur Background Function | |
| def gaussian_blur_background(image): | |
| # Preprocess image for segmentation | |
| inputs = segformer_extractor(images=image, return_tensors="pt") | |
| outputs = segformer_model(**inputs) | |
| logits = outputs.logits | |
| segmentation = torch.argmax(logits, dim=1)[0].numpy() | |
| # Create a binary mask for 'person' class (class index 12) | |
| human_mask = (segmentation == 12).astype(np.uint8) * 255 | |
| human_mask_image = Image.fromarray(human_mask).resize(image.size) | |
| # Apply Gaussian blur to the entire image | |
| blurred_background = image.filter(ImageFilter.GaussianBlur(15)) | |
| # Composite the original image with blurred background using the mask | |
| composite_image = Image.composite(image, blurred_background, human_mask_image) | |
| return composite_image | |
| # Depth-Based Lens Blur Function | |
| def lens_blur(image): | |
| # Preprocess image for depth estimation | |
| inputs = dpt_extractor(images=image, return_tensors="pt") | |
| with torch.no_grad(): | |
| outputs = dpt_model(**inputs) | |
| depth_map = outputs.predicted_depth.squeeze().cpu().numpy() | |
| # Normalize depth map to range [0, 15] and invert for blur intensity | |
| depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min()) * 15 | |
| depth_map = 15 - depth_map | |
| depth_map_resized = cv2.resize(depth_map, (image.width, image.height)) | |
| # Convert image to OpenCV format | |
| image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) | |
| blurred_image = np.zeros_like(image_cv, dtype=np.float32) | |
| # Apply variable blur based on depth | |
| for blur_radius in range(1, 16): | |
| blurred_layer = cv2.GaussianBlur(image_cv, (0, 0), sigmaX=blur_radius) | |
| mask = ((depth_map_resized >= (blur_radius - 1)) & (depth_map_resized < blur_radius)).astype(np.float32) | |
| mask = cv2.merge([mask] * 3) | |
| blurred_image += blurred_layer * mask | |
| blurred_image = np.clip(blurred_image, 0, 255).astype(np.uint8) | |
| blurred_image_pil = Image.fromarray(cv2.cvtColor(blurred_image, cv2.COLOR_BGR2RGB)) | |
| return blurred_image_pil | |
| # Gradio Interface | |
| def process_image(image, effect): | |
| if effect == "Gaussian Blur Background": | |
| return gaussian_blur_background(image) | |
| elif effect == "Lens Blur": | |
| return lens_blur(image) | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# BokehBot: Gaussian and Lens Blur Effects") | |
| with gr.Row(): | |
| with gr.Column(): | |
| uploaded_image = gr.Image(type="pil", label="Upload an Image") | |
| effect = gr.Radio(["Gaussian Blur Background", "Lens Blur"], label="Choose Effect") | |
| process_button = gr.Button("Apply Effect") | |
| with gr.Column(): | |
| output_image = gr.Image(type="pil", label="Processed Image") | |
| process_button.click(process_image, inputs=[uploaded_image, effect], outputs=output_image) | |
| demo.launch() | |