Spaces:
Running
Running
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| import cv2 | |
| from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation | |
| import gradio as gr | |
| # Initialize the SegFormer model for segmentation | |
| segformer_processor = SegformerImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512") | |
| segformer_model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512") | |
| # Function to segment the person in the image | |
| def segment_person(image_input): | |
| # Convert input image (numpy array in RGB) to PIL Image | |
| image = Image.fromarray(image_input).convert("RGB") | |
| original_width, original_height = image.size | |
| # Resize image to 512x512 for the model | |
| model_input = image.resize((512, 512), Image.Resampling.LANCZOS) | |
| # Prepare the image for SegFormer | |
| inputs = segformer_processor(images=model_input, return_tensors="pt") | |
| # Perform inference | |
| with torch.no_grad(): | |
| outputs = segformer_model(**inputs) | |
| logits = outputs.logits | |
| # Upsample logits to 512x512 | |
| upsampled_logits = torch.nn.functional.interpolate( | |
| logits, size=(512, 512), mode="bilinear", align_corners=False | |
| ) | |
| # Get the predicted segmentation mask (person class = 12 in ADE20K dataset) | |
| person_class_id = 12 | |
| predicted_mask = upsampled_logits.argmax(dim=1)[0] # Shape: (512, 512) | |
| binary_mask = (predicted_mask == person_class_id).cpu().numpy() # Boolean mask | |
| # Post-process the mask | |
| mask_uint8 = (binary_mask * 255).astype(np.uint8) | |
| kernel = np.ones((5, 5), np.uint8) | |
| mask_cleaned = cv2.morphologyEx(mask_uint8, cv2.MORPH_CLOSE, kernel, iterations=2) | |
| mask_cleaned = cv2.morphologyEx(mask_cleaned, cv2.MORPH_OPEN, kernel, iterations=2) | |
| mask_smoothed = cv2.GaussianBlur(mask_cleaned, (7, 7), 0) | |
| _, mask_final = cv2.threshold(mask_smoothed, 127, 255, cv2.THRESH_BINARY) | |
| # Resize mask back to original dimensions | |
| mask_pil = Image.fromarray(mask_final) | |
| mask_resized = mask_pil.resize((original_width, original_height), Image.Resampling.LANCZOS) | |
| mask_array = np.array(mask_resized) > 0 # Boolean mask | |
| return mask_array | |
| # Function to apply background blur | |
| def blur_background(image_input, blur_strength): | |
| # Ensure image is in numpy array format (RGB) | |
| image_array = np.array(image_input) | |
| # Segment the person | |
| mask = segment_person(image_array) | |
| # Apply Gaussian blur to the entire image | |
| sigma = blur_strength | |
| blurred_image = cv2.GaussianBlur(image_array, (0, 0), sigmaX=sigma, sigmaY=sigma) | |
| # Composite the original foreground with the blurred background | |
| mask_3d = mask[:, :, np.newaxis] # Add channel dimension for broadcasting | |
| result = np.where(mask_3d, image_array, blurred_image).astype(np.uint8) | |
| return result | |
| # Gradio interface function | |
| def gradio_interface(image, blur_strength): | |
| if image is None: | |
| raise ValueError("Please upload an image.") | |
| # Process the image | |
| output_image = blur_background(image, blur_strength) | |
| return output_image | |
| # Create the Gradio app | |
| app = gr.Interface( | |
| fn=gradio_interface, | |
| inputs=[ | |
| gr.Image(type="numpy", label="Upload Image"), | |
| gr.Slider(minimum=1, maximum=25, value=10, step=1, label="Blur Strength (Sigma)") | |
| ], | |
| outputs=gr.Image(type="numpy", label="Output Image"), | |
| title="Person Segmentation and Background Blur", | |
| description="Upload an image to segment the person and blur the background. Adjust the blur strength using the slider." | |
| ) | |
| # Launch the app | |
| app.launch() |