Update app.py

app.py

@@ -14,30 +14,86 @@ def preprocess_image(image):
     image = image.resize((512, 512))
     return image
 
-def segment_image(image, model_name="facebook/mask2former-swin-large-cityscapes-semantic"):
-    """Perform semantic segmentation on the input image."""
-    from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
+def segment_image(image, model_name="yolov8n-seg"):
+    """
+    Perform instance segmentation on the input image using YOLO segmentation model.
     
-    # Load processor and model
-    processor = AutoImageProcessor.from_pretrained(model_name)
-    model = Mask2FormerForUniversalSegmentation.from_pretrained(model_name)
+    Args:
+        image (PIL.Image): Input image
+        model_name (str): Name of the YOLO segmentation model
     
-    # Prepare inputs
-    inputs = processor(images=image, return_tensors="pt")
+    Returns:
+        numpy.ndarray: Segmentation mask with instance segmentation
+    """
+    from ultralytics import YOLO
+    import numpy as np
+    import torch
+    
+    # Load the YOLO segmentation model
+    model = YOLO(model_name)
     
     # Run inference
-    with torch.no_grad():
-        outputs = model(**inputs)
-    
-    # Post-process segmentation
-    semantic_map = processor.post_process_semantic_segmentation(
-        outputs, 
-        target_sizes=[image.size[::-1]]
-    )[0]
-    
-    # Convert to numpy and create binary mask
-    semantic_map = semantic_map.numpy()
-    return semantic_map
+    results = model(image)
+    
+    # Create a blank mask
+    mask = np.zeros(image.size[::-1], dtype=np.uint8)
+    
+    # Process each detected object
+    for result in results:
+        # Get masks for all detected objects
+        masks = result.masks
+        
+        if masks is not None:
+            # Convert masks to numpy and add to the overall mask
+            for single_mask in masks:
+                # Convert mask to numpy and resize if needed
+                mask_array = single_mask.data.cpu().numpy().squeeze()
+                mask_array = (mask_array > 0.5).astype(np.uint8)
+                
+                # If mask size doesn't match image, resize
+                if mask_array.shape != mask.shape:
+                    from PIL import Image
+                    mask_array = np.array(
+                        Image.fromarray(mask_array).resize(
+                            image.size[::-1], 
+                            Image.NEAREST
+                        )
+                    )
+                
+                # Add this mask to the overall mask
+                mask = np.maximum(mask, mask_array)
+    
+    return mask
+
+def process_image(image, blur_type, sigma=15):
+    """Process image based on blur type."""
+    # Preprocess image
+    pil_image = preprocess_image(image)
+    
+    # Apply appropriate blur
+    if blur_type == "Gaussian Background Blur":
+        # Get segmentation mask
+        segmentation_mask = segment_image(pil_image)
+        
+        # Convert to 3-channel mask
+        mask_3d = np.stack([segmentation_mask] * 3, axis=2)
+        
+        # Apply Gaussian blur
+        image_array = np.array(pil_image)
+        blurred = np.zeros_like(image_array)
+        for channel in range(3):
+            blurred[:, :, channel] = gaussian_filter(image_array[:, :, channel], sigma=sigma)
+        
+        # Combine original and blurred images
+        result = image_array * mask_3d + blurred * (1 - mask_3d)
+        result = Image.fromarray(result.astype(np.uint8))
+    
+    elif blur_type == "Depth-Aware Lens Blur":
+        result = apply_depth_aware_blur(pil_image, max_sigma=sigma)
+    else:
+        result = pil_image
+    
+    return result
 
 def apply_gaussian_blur(image, sigma=15):
     """Apply Gaussian blur to the background."""
@@ -78,19 +134,15 @@ def estimate_depth(image, model_name="depth-anything/Depth-Anything-V2-Small-hf"
     return depth_map
 
 def apply_depth_aware_blur(image, max_sigma=10, min_sigma=0):
-    """Apply depth-aware blur to the image."""
+    """Apply depth-aware blur to the image (REVERSED version)."""
     # Estimate depth
     depth_map = estimate_depth(image)
     
     image_array = np.array(image)
     blurred = np.zeros_like(image_array, dtype=np.float32)
-
-    inverted_depth_map = 1 - depth_map
     
-    # Interpolate sigmas based on depth
+    # REVERSED: Now we use depth_map directly (no inversion) so farther objects get more blur
     sigmas = np.interp(depth_map, [0, 1], [min_sigma, max_sigma])
-
-    
     
     # Precompute blurred layers
     blur_stack = {}
@@ -117,20 +169,7 @@ def apply_depth_aware_blur(image, max_sigma=10, min_sigma=0):
     
     return Image.fromarray(blurred.astype(np.uint8))
 
-def process_image(image, blur_type, sigma=15):
-    """Process image based on blur type."""
-    # Preprocess image
-    pil_image = preprocess_image(image)
-    
-    # Apply appropriate blur
-    if blur_type == "Gaussian Background Blur":
-        result = apply_gaussian_blur(pil_image, sigma)
-    elif blur_type == "Depth-Aware Lens Blur":
-        result = apply_depth_aware_blur(pil_image, max_sigma=sigma)
-    else:
-        result = pil_image
-    
-    return result
+
 
 # Gradio Interface
 def create_blur_app():