Update app.py

app.py

@@ -3,23 +3,54 @@ from PIL import Image, ImageFilter
 import numpy as np
 import torch
 import cv2
-from transformers import AutoImageProcessor, AutoModelForDepthEstimation
+from transformers import AutoImageProcessor, AutoModelForDepthEstimation, OneFormerProcessor, OneFormerForUniversalSegmentation
 
 # Load depth estimation model
 image_processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
-model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
+depth_model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
 
-def apply_gaussian_blur(image, mask):
-    """Applies Gaussian blur to the background based on a user-drawn mask."""
+# Load OneFormer processor and model
+processor = OneFormerProcessor.from_pretrained("shi-labs/oneformer_coco_swin_large")
+segmentation_model = OneFormerForUniversalSegmentation.from_pretrained("shi-labs/oneformer_coco_swin_large")
+
+def apply_gaussian_blur(image, foreground_label='person'):
+    """Applies Gaussian blur to the background based on a segmentation mask for the foreground."""
     
-    if mask is None:
-        return image  # If no mask is provided, return the original image
+    # Prepare input for semantic segmentation
+    inputs = processor(images=image, task_inputs=["semantic"], return_tensors="pt")
+
+    # Semantic segmentation
+    with torch.no_grad():
+        outputs = segmentation_model(**inputs)
+
+    # Processing semantic segmentation output
+    predicted_semantic_map = processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
+    segmentation_mask = predicted_semantic_map.cpu().numpy()
+
+    # Get the mapping of class IDs to labels from the processor
+    id2label = segmentation_model.config.id2label
+    foreground_class_id = None
+    for id, label in id2label.items():
+        if label == foreground_label:
+            foreground_class_id = id
+            break
+
+    if foreground_class_id is None:
+        print(f"Error: Could not find the label '{foreground_label}' in the model's class mapping.")
+        return image  # Return original image if foreground label is not found
+
+    # Create a black background mask and set the pixels corresponding to the foreground object to white
+    output_mask_array = np.zeros(segmentation_mask.shape, dtype=np.uint8)
+    output_mask_array[segmentation_mask == foreground_class_id] = 255
+
+    # Convert the output mask to a PIL Image (Grayscale)
+    mask_pil = Image.fromarray(output_mask_array, mode='L')
     
-    # Ensure mask is grayscale and resized to match image dimensions
-    mask_pil = Image.fromarray(mask).convert("L").resize(image.size)
-    mask_array = np.array(mask_pil)
+    # Resize the mask to match the image size
+    mask_pil = mask_pil.resize(image.size)
+    output_mask_array = np.array(mask_pil)
 
-    # Create a blurred background
+    # Create a blurred version of the input image
     blurred_background = image.filter(ImageFilter.GaussianBlur(radius=15))
 
     # Convert images to NumPy arrays
@@ -27,7 +58,7 @@ def apply_gaussian_blur(image, mask):
     blurred_array = np.array(blurred_background)
 
     # Create a boolean mask (foreground = True, background = False)
-    foreground_mask = mask_array > 0
+    foreground_mask = output_mask_array > 0
     foreground_mask_3d = np.stack([foreground_mask] * 3, axis=-1)
 
     # Blend the original image with the blurred background
@@ -47,25 +78,26 @@ def apply_lens_blur(image):
     inputs = image_processor(images=resized_image, return_tensors="pt")
 
     with torch.no_grad():
-        outputs = model(**inputs)
+        outputs = depth_model(**inputs)
         predicted_depth = outputs.predicted_depth
 
-    # Interpolate depth map to match the image size
+    # Interpolate to the original size
     prediction = torch.nn.functional.interpolate(
         predicted_depth.unsqueeze(1),
         size=resized_image.size[::-1],
         mode="bicubic",
         align_corners=False,
-    ).squeeze() 
+    ).squeeze()
 
     # Convert prediction to a NumPy array
     depth_map = prediction.cpu().numpy()
 
-    # Normalize the depth map
+    # Normalize the depth map to the range 0-1
     depth_norm = (depth_map - np.min(depth_map)) / (np.max(depth_map) - np.min(depth_map))
 
-    num_blur_levels = 5  
+    num_blur_levels = 5
     blurred_layers = []
+    
     for i in range(num_blur_levels):
         sigma = i * 0.5
         if sigma == 0:
@@ -77,6 +109,7 @@ def apply_lens_blur(image):
     depth_indices = ((1 - depth_norm) * (num_blur_levels - 1)).astype(np.uint8)
 
     final_blurred_image = np.zeros_like(image_np)
+    
     for y in range(image_np.shape[0]):
         for x in range(image_np.shape[1]):
             depth_index = depth_indices[y, x]
@@ -87,27 +120,24 @@ def apply_lens_blur(image):
 
     return final_blurred_pil_image
 
-def process_image(image, mask, blur_type):
+def process_image(image, blur_type, foreground_label='person'):
     """Processes the image based on the selected blur type."""
     if blur_type == "Gaussian Blur":
-        return apply_gaussian_blur(image, mask)
-    elif blur_type == "Lens Blur":
-        return apply_lens_blur(image)
+        return apply_gaussian_blur(image, foreground_label=foreground_label)
     else:
-        return image
+        return apply_lens_blur(image)
 
 interface = gr.Interface(
     fn=process_image,
     inputs=[
         gr.Image(type="pil", label="Upload an Image"),
-        gr.Sketchpad(height=256, width=256, label="Draw Mask (Only for Gaussian Blur)"),
-        gr.Radio(["Gaussian Blur", "Lens Blur"], label="Choose Blur Effect")
+        gr.Radio(["Gaussian Blur", "Lens Blur"], label="Choose Blur Effect"),
+        gr.Textbox(label="Foreground Label (for Gaussian Blur)", default="person")
     ],
-    outputs=gr.Image(type="pil"),
+    outputs=[gr.Image(type="pil"), gr.Image(type="pil")],
     title="Gaussian & Lens Blur Effects",
-    description="Upload an image and select either Gaussian blur (with mask) or depth-based lens blur."
+    description="Upload an image and select either Gaussian blur (with foreground segmentation) or depth-based lens blur."
 )
 
-
 if __name__ == "__main__":
     interface.launch()