Update app.py

app.py

@@ -93,6 +93,42 @@ def segment_image(image):
     
     return mask_array, mask_viz
 
+# Function to get depth map
+def get_depth_map(image):
+    """Get depth map from image using DPT model"""
+    print("Running depth estimation...")
+    
+    # Convert to PIL Image if needed
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+    
+    # Original dimensions
+    original_size = image.size
+    
+    # Higher resolution for depth estimation
+    model_size = (640, 640)
+    model_image = image.resize(model_size, Image.LANCZOS)
+    
+    # Process image for depth estimation
+    inputs = depth_processor(images=model_image, return_tensors="pt")
+    
+    # Run inference
+    with torch.no_grad():
+        outputs = depth_model(**inputs)
+        predicted_depth = outputs.predicted_depth
+    
+    # Process depth map
+    depth = predicted_depth.squeeze().cpu().numpy()
+    depth_map = (depth - depth.min()) / (depth.max() - depth.min())
+    
+    # Create colored depth map for visualization
+    depth_map_colored = plt.cm.viridis(depth_map)[:, :, :3]
+    depth_map_viz = Image.fromarray((depth_map_colored * 255).astype(np.uint8))
+    depth_map_viz_resized = depth_map_viz.resize(original_size, Image.LANCZOS)
+    
+    # Return both visualization and raw depth map
+    return np.array(depth_map_viz_resized), depth_map
+
 # Function to apply Gaussian blur to background
 def apply_background_blur(image, mask, sigma=15):
     """Apply Gaussian blur to background while keeping foreground sharp"""
@@ -121,10 +157,10 @@ def apply_background_blur(image, mask, sigma=15):
     
     return result
 
-# Function for depth estimation and depth-based blur
-def apply_depth_based_blur(image, mask=None, max_sigma=15):
-    """Apply depth-based blur using DPT model"""
-    print(f"Running depth estimation and applying depth-based blur with max_sigma={max_sigma}...")
+# Function for depth-based blur
+def apply_depth_based_blur(image, mask, depth_map, max_sigma=15):
+    """Apply depth-based blur using provided depth map"""
+    print(f"Applying depth-based blur with max_sigma={max_sigma}...")
     
     # Convert to PIL Image if needed
     if not isinstance(image, Image.Image):
@@ -133,38 +169,32 @@ def apply_depth_based_blur(image, mask=None, max_sigma=15):
     # Original dimensions
     original_size = image.size
     
-    # Higher resolution for depth estimation
-    model_size = (640, 640)
-    model_image = image.resize(model_size, Image.LANCZOS)
-    
-    # Process image for depth estimation
-    inputs = depth_processor(images=model_image, return_tensors="pt")
-    
-    # Run inference
-    with torch.no_grad():
-        outputs = depth_model(**inputs)
-        predicted_depth = outputs.predicted_depth
-    
-    # Process depth map
-    depth = predicted_depth.squeeze().cpu().numpy()
-    depth_map = (depth - depth.min()) / (depth.max() - depth.min())
-    
-    # Create high-res depth map
-    depth_map_highres = cv2.resize(depth_map, (model_size[0], model_size[1]), interpolation=cv2.INTER_CUBIC)
+    # Resize depth map to match image size if needed
+    if depth_map.shape[:2] != image.size[::-1]:
+        depth_map_resized = cv2.resize(depth_map, original_size, interpolation=cv2.INTER_CUBIC)
+    else:
+        depth_map_resized = depth_map
     
     # Invert depth map (closer objects should be less blurred)
-    inverted_depth_map = 1.0 - depth_map_highres
+    inverted_depth_map = 1.0 - depth_map_resized
     
-    # If mask is provided, ensure foreground is not blurred at all
+    # Apply mask to ensure foreground is not blurred
     if mask is not None:
-        # Resize mask to match model size
-        mask_pil = Image.fromarray((mask.astype(np.uint8) * 255))
-        mask_resized = np.array(mask_pil.resize(model_size, Image.LANCZOS)) > 128
+        # Ensure mask has proper dimensions
+        if isinstance(mask, np.ndarray):
+            if mask.ndim == 3:
+                binary_mask = mask[:,:,0] > 0
+            else:
+                binary_mask = mask > 0
+        else:
+            # Convert to numpy if needed
+            binary_mask = np.array(mask) > 0
+            
         # Set depth map to 0 (no blur) for foreground pixels
-        inverted_depth_map = inverted_depth_map * (1 - mask_resized)
+        inverted_depth_map = inverted_depth_map * (1 - binary_mask)
     
     # Convert to numpy array for processing
-    img_array = np.array(model_image)
+    img_array = np.array(image)
     
     # Create a progressive blur effect with multiple levels
     result = np.copy(img_array)
@@ -189,16 +219,7 @@ def apply_depth_based_blur(image, mask=None, max_sigma=15):
         for c in range(3):
             result[:,:,c] = np.where(level_mask, level_blurred[:,:,c], result[:,:,c])
     
-    # Convert result back to PIL and resize to original dimensions
-    result_pil = Image.fromarray(result.astype(np.uint8))
-    result_resized = result_pil.resize(original_size, Image.LANCZOS)
-    
-    # Create colored depth map for visualization
-    depth_map_colored = plt.cm.viridis(depth_map)[:, :, :3]
-    depth_map_viz = Image.fromarray((depth_map_colored * 255).astype(np.uint8))
-    depth_map_viz_resized = depth_map_viz.resize(original_size, Image.LANCZOS)
-    
-    return np.array(depth_map_viz_resized), np.array(result_resized)
+    return result
 
 # Main processing function
 def process_image(input_image, blur_type="Gaussian Blur", blur_intensity=15):
@@ -222,19 +243,16 @@ def process_image(input_image, blur_type="Gaussian Blur", blur_intensity=15):
         # Step 1: Get segmentation mask
         mask_array, mask_viz = segment_image(pil_img)
         
-        # Step 2: Apply appropriate blur effect
+        # Step 2: Always get depth map (for both blur types)
+        depth_viz, depth_map = get_depth_map(pil_img)
+        
+        # Step 3: Apply appropriate blur effect
         if blur_type == "Gaussian Blur":
             # Apply regular Gaussian blur
             result = apply_background_blur(pil_img, mask_array, sigma=blur_intensity)
-            # Create placeholder for depth map (black image)
-            depth_viz = np.zeros_like(img)
-            # Add text saying "Depth map not used for Gaussian blur"
-            font = cv2.FONT_HERSHEY_SIMPLEX
-            cv2.putText(depth_viz, "Depth map not used", (50, 50), font, 1, (255, 255, 255), 2)
-            
         else:  # "Depth-based Lens Blur"
             # Apply depth-based blur
-            depth_viz, result = apply_depth_based_blur(pil_img, mask_array, max_sigma=blur_intensity)
+            result = apply_depth_based_blur(pil_img, mask_array, depth_map, max_sigma=blur_intensity)
         
         return result, mask_viz, depth_viz
         
@@ -296,8 +314,9 @@ with gr.Blocks(title="Image Blur Effects") as demo:
     gr.Markdown("")
     gr.Markdown("### Notes")
     gr.Markdown("- The red areas in the Foreground Mask show what will remain sharp")
-    gr.Markdown("- Depth-based blur creates a more realistic effect with blur increasing with distance")
-    gr.Markdown("- For best results, use images with clear foreground subjects")
+    gr.Markdown("- The Depth Map shows estimated distances (yellow=far, blue=close)")
+    gr.Markdown("- Gaussian Blur applies uniform blur to the background")
+    gr.Markdown("- Depth-based Blur varies blur intensity based on distance")
     gr.Markdown("- Created for EEE 515 Assignment (Problem 2, Part 6)")
 
 # Launch the demo