Update MVPR.py

MVPR.py

@@ -0,0 +1,95 @@
+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()