Update app.py

app.py

@@ -17,17 +17,8 @@ def preprocess_image(image):
 def segment_image(image, model_name="yolov8n-seg"):
     """
     Perform instance segmentation on the input image using YOLO segmentation model.
-    
-    Args:
-        image (PIL.Image): Input image
-        model_name (str): Name of the YOLO segmentation model
-    
-    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)
@@ -35,76 +26,38 @@ def segment_image(image, model_name="yolov8n-seg"):
     # Run inference
     results = model(image)
     
-    # Create a blank mask
-    mask = np.zeros(image.size[::-1], dtype=np.uint8)
+    # Create a blank mask (1 for foreground, 0 for background)
+    mask = np.zeros((image.size[1], image.size[0]), 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:
+        if result.masks is not None:
+            for single_mask in result.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
+                # Resize if needed
                 if mask_array.shape != mask.shape:
-                    from PIL import Image
                     mask_array = np.array(
                         Image.fromarray(mask_array).resize(
-                            image.size[::-1], 
+                            (image.size[0], image.size[1]), 
                             Image.NEAREST
                         )
                     )
                 
-                # Add this mask to the overall mask
+                # Add this mask to the overall mask (OR operation)
                 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."""
     # Convert image to numpy array
     image_array = np.array(image)
     
-    # Create segmentation mask (assuming we want to keep the foreground)
-    segmentation_mask = segment_image(image)
-    
-    # Choose a prominent object class (e.g., person with ID 24 in Cityscapes)
-    foreground_mask = (segmentation_mask == 24).astype(np.uint8)
+    # Get segmentation mask (1 for foreground, 0 for background)
+    foreground_mask = segment_image(image)
     
     # Prepare blurred version
     blurred = np.zeros_like(image_array)
@@ -128,48 +81,49 @@ def estimate_depth(image, model_name="depth-anything/Depth-Anything-V2-Small-hf"
     depth_output = depth_estimator(image)
     depth_map = np.array(depth_output["depth"])
     
-    # Normalize depth map
+    # Normalize depth map (0-1 where 1 is farthest)
     depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
     
     return depth_map
 
 def apply_depth_aware_blur(image, max_sigma=10, min_sigma=0):
-    """Apply depth-aware blur to the image (REVERSED version)."""
-    # Estimate depth
+    """Apply depth-aware blur with farther objects more blurred."""
+    # Estimate depth (1 = farthest)
     depth_map = estimate_depth(image)
     
     image_array = np.array(image)
-    blurred = np.zeros_like(image_array, dtype=np.float32)
-    
-    # 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 = {}
-    for sigma in np.unique(sigmas):
-        if sigma > 0:
-            blurred_layer = np.zeros_like(image_array, dtype=np.float32)
-            for channel in range(3):
-                blurred_layer[:, :, channel] = gaussian_filter(
-                    image_array[:, :, channel].astype(np.float32),
-                    sigma=sigma
-                )
-            blur_stack[sigma] = blurred_layer
-    
-    # Blend based on depth
-    for sigma in np.unique(sigmas):
-        if sigma > 0:
-            mask = (sigmas == sigma)
-            mask_3d = np.stack([mask] * 3, axis=2)
-            blurred += mask_3d * blur_stack[sigma]
-        else:
-            mask = (sigmas == 0)
-            mask_3d = np.stack([mask] * 3, axis=2)
-            blurred += mask_3d * image_array
-    
-    return Image.fromarray(blurred.astype(np.uint8))
-
+    
+    # Create single blurred version at max sigma for efficiency
+    max_blurred = np.zeros_like(image_array, dtype=np.float32)
+    for channel in range(3):
+        max_blurred[:, :, channel] = gaussian_filter(
+            image_array[:, :, channel].astype(np.float32),
+            sigma=max_sigma
+        )
+    
+    # Create 3-channel depth map for blending
+    depth_3d = np.stack([depth_map] * 3, axis=2)
+    
+    # Blend between original (near) and blurred (far) based on depth
+    # Higher depth values (farther) get more blur
+    result = image_array * (1 - depth_3d) + max_blurred * depth_3d
+    
+    return Image.fromarray(result.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():