Update app.py

app.py

@@ -3,188 +3,287 @@ import numpy as np
 import torch
 from PIL import Image, ImageFilter
 import cv2
+import os
+import sys
+import traceback
 from transformers import pipeline
 
+# Configure logging to console
+import logging
+logging.basicConfig(level=logging.INFO, 
+                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+                    stream=sys.stdout)
+logger = logging.getLogger(__name__)
+
 # Set device to GPU if available
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
+logger.info(f"Using device: {device}")
+
+# Use smaller, more efficient models for Hugging Face Space
+SEGMENTATION_MODEL = "facebook/sam-vit-base"
+DEPTH_MODEL = "depth-anything/Depth-Anything-V2-Small-hf"
 
-# Load models only once at startup to improve performance
-segmentation_model = "facebook/sam-vit-huge"
-depth_model = "depth-anything/Depth-Anything-V2-Small-hf"
+# Global variables for models
+segmentation_pipe = None
+depth_pipe = None
 
-# Initialize pipelines
-segmentation_pipe = pipeline("image-segmentation", model=segmentation_model)
-depth_pipe = pipeline("depth-estimation", model=depth_model)
+def load_segmentation_model():
+    """Load the segmentation model on demand"""
+    global segmentation_pipe
+    if segmentation_pipe is None:
+        try:
+            logger.info("Loading segmentation model...")
+            segmentation_pipe = pipeline("image-segmentation", model=SEGMENTATION_MODEL)
+            logger.info("Segmentation model loaded successfully")
+        except Exception as e:
+            logger.error(f"Error loading segmentation model: {e}")
+            logger.error(traceback.format_exc())
+            return None
+    return segmentation_pipe
+
+def load_depth_model():
+    """Load the depth model on demand"""
+    global depth_pipe
+    if depth_pipe is None:
+        try:
+            logger.info("Loading depth estimation model...")
+            depth_pipe = pipeline("depth-estimation", model=DEPTH_MODEL)
+            logger.info("Depth estimation model loaded successfully")
+        except Exception as e:
+            logger.error(f"Error loading depth model: {e}")
+            logger.error(traceback.format_exc())
+            return None
+    return depth_pipe
 
 def get_segmentation_mask(input_image):
-    """Get segmentation mask using the pre-loaded segmentation pipeline"""
-    # Resize image to 512x512 for consistent processing
-    input_image = input_image.resize((512, 512)).convert('RGB')
-    
-    # Get the segmentation result
-    result = segmentation_pipe(input_image)
-    
-    # Extract the first mask (assuming it's the most prominent object)
-    if len(result) > 0:
-        # For SAM-like models that return multiple masks
-        mask = result[0]['mask']
-        mask = np.array(mask) * 255  # Scale to [0, 255]
-    else:
-        # Fallback - create empty mask
-        mask = np.zeros((512, 512), dtype=np.uint8)
-    
-    # Convert to PIL Image
-    mask_img = Image.fromarray(mask.astype(np.uint8))
+    """Get segmentation mask using the segmentation pipeline"""
+    try:
+        # Load the model if not already loaded
+        model = load_segmentation_model()
+        if model is None:
+            logger.warning("Failed to load segmentation model, returning empty mask")
+            return Image.new('L', (512, 512), 0), input_image
+        
+        # Process the image
+        input_pil = Image.fromarray(input_image).convert('RGB') if isinstance(input_image, np.ndarray) else input_image.convert('RGB')
+        input_pil = input_pil.resize((512, 512))
+        
+        # Get segmentation results
+        results = model(input_pil)
+        
+        # Extract mask (handling different model outputs)
+        if isinstance(results, list) and len(results) > 0:
+            if 'mask' in results[0]:
+                mask = results[0]['mask']
+                # Convert mask to proper format
+                if isinstance(mask, Image.Image):
+                    mask_array = np.array(mask)
+                else:
+                    mask_array = mask
+                mask_array = (mask_array * 255).astype(np.uint8)
+            else:
+                # Create composite mask from segmentation results
+                mask_array = np.zeros((512, 512), dtype=np.uint8)
+                for segment in results:
+                    if 'segmentation' in segment:
+                        mask_array = np.logical_or(mask_array, segment['segmentation']).astype(np.uint8) * 255
+        else:
+            # Create blank mask as fallback
+            mask_array = np.zeros((512, 512), dtype=np.uint8)
+        
+        # Convert to PIL Image
+        mask_img = Image.fromarray(mask_array)
+        
+        return mask_img, input_pil
     
-    return mask_img, input_image
+    except Exception as e:
+        logger.error(f"Error in segmentation: {e}")
+        logger.error(traceback.format_exc())
+        # Return a blank mask in case of error
+        if isinstance(input_image, np.ndarray):
+            return Image.new('L', (512, 512), 0), Image.fromarray(input_image).resize((512, 512))
+        else:
+            return Image.new('L', (512, 512), 0), input_image.resize((512, 512))
 
 def apply_background_blur(original_image, mask_image, sigma=15):
     """Apply Gaussian blur to the background using a segmentation mask"""
-    # Ensure mask is binary (0 for background, 255 for foreground)
-    mask_array = np.array(mask_image)
-    _, binary_mask = cv2.threshold(mask_array, 127, 255, cv2.THRESH_BINARY)
-    mask_img = Image.fromarray(binary_mask)
-    
-    # Create a blurred version of the original image
-    blurred_img = original_image.filter(ImageFilter.GaussianBlur(radius=sigma))
-    
-    # Convert images to numpy arrays for easier manipulation
-    original_array = np.array(original_image)
-    blurred_array = np.array(blurred_img)
-    mask_array = np.array(mask_img)
-    
-    # Create the composite image: foreground from original, background from blurred
-    result_array = np.zeros_like(original_array)
-    
-    # Where mask is white (255), use original image; where mask is black (0), use blurred image
-    for c in range(3):  # For each color channel (RGB)
-        result_array[:, :, c] = np.where(mask_array == 255, 
-                                         original_array[:, :, c], 
-                                         blurred_array[:, :, c])
-    
-    # Convert back to PIL Image
-    result_img = Image.fromarray(result_array)
+    try:
+        # Convert to PIL Image if needed
+        if isinstance(original_image, np.ndarray):
+            original_image = Image.fromarray(original_image)
+        if isinstance(mask_image, np.ndarray):
+            mask_image = Image.fromarray(mask_image)
+            
+        # Ensure consistent sizes
+        original_image = original_image.resize((512, 512)).convert('RGB')
+        mask_image = mask_image.resize((512, 512)).convert('L')
+        
+        # Ensure mask is binary
+        mask_array = np.array(mask_image)
+        _, binary_mask = cv2.threshold(mask_array, 127, 255, cv2.THRESH_BINARY)
+        
+        # Create blurred version of the original
+        blurred_img = original_image.filter(ImageFilter.GaussianBlur(radius=sigma))
+        
+        # Convert to numpy for processing
+        original_array = np.array(original_image)
+        blurred_array = np.array(blurred_img)
+        
+        # Create mask for all 3 channels
+        mask_3d = np.stack([binary_mask, binary_mask, binary_mask], axis=2)
+        
+        # Combine original foreground with blurred background
+        result_array = np.where(mask_3d == 255, original_array, blurred_array)
+        
+        return Image.fromarray(result_array.astype(np.uint8))
     
-    return result_img
+    except Exception as e:
+        logger.error(f"Error in background blur: {e}")
+        logger.error(traceback.format_exc())
+        # Return original image in case of error
+        return original_image
 
 def get_depth_map(input_image):
-    """Get depth map using the pre-loaded depth estimation pipeline"""
-    # Ensure image is in RGB format and resized to 512x512
-    input_image = input_image.resize((512, 512)).convert('RGB')
-    
-    # Get the depth map
-    result = depth_pipe(input_image)
-    depth_map = result["depth"]
-    
-    # Convert to numpy array for further processing
-    depth_array = np.array(depth_map)
+    """Get depth map using the depth estimation pipeline"""
+    try:
+        # Load model if not already loaded
+        model = load_depth_model()
+        if model is None:
+            logger.warning("Failed to load depth model, returning empty depth map")
+            return Image.new('L', (512, 512), 128), np.ones((512, 512)) * 0.5
+        
+        # Convert to PIL if needed
+        if isinstance(input_image, np.ndarray):
+            input_image = Image.fromarray(input_image)
+            
+        # Ensure consistent format
+        input_image = input_image.resize((512, 512)).convert('RGB')
+        
+        # Get depth estimation
+        result = model(input_image)
+        depth_map = result["depth"]
+        
+        # Convert to numpy for further processing
+        depth_array = np.array(depth_map)
+        
+        return depth_map, depth_array
     
-    return depth_map, depth_array
+    except Exception as e:
+        logger.error(f"Error in depth estimation: {e}")
+        logger.error(traceback.format_exc())
+        # Return default depth in case of error
+        return Image.new('L', (512, 512), 128), np.ones((512, 512)) * 0.5
 
 def apply_depth_based_blur(original_image, depth_array, max_blur=30):
     """Apply variable Gaussian blur based on depth"""
-    # Convert depth array to proper format if needed
-    if len(depth_array.shape) == 3 and depth_array.shape[2] > 1:
-        # If depth map has multiple channels, convert to grayscale
-        depth_array = np.mean(depth_array, axis=2)
-    
-    # Normalize depth values to range [0, 1]
-    depth_min = depth_array.min()
-    depth_max = depth_array.max()
-    normalized_depth = (depth_array - depth_min) / (depth_max - depth_min)
-    
-    # Create a series of increasingly blurred versions of the image
-    blurred_images = []
-    for blur_amount in range(max_blur + 1):
-        blurred_images.append(original_image.filter(ImageFilter.GaussianBlur(radius=blur_amount)))
-    
-    # Convert to numpy arrays for easier processing
-    original_array = np.array(original_image)
-    result_array = np.zeros_like(original_array)
-    
-    # For each pixel, determine the blur level based on depth
-    height, width = normalized_depth.shape
-    for y in range(height):
-        for x in range(width):
-            # Calculate blur radius proportional to depth
-            # Higher normalized_depth = farther object = more blur
-            blur_radius = int(normalized_depth[y, x] * max_blur)
-            result_array[y, x] = np.array(blurred_images[blur_radius])[y, x]
+    try:
+        # Convert to PIL if needed
+        if isinstance(original_image, np.ndarray):
+            original_image = Image.fromarray(original_image)
+            
+        # Ensure consistent size
+        original_image = original_image.resize((512, 512)).convert('RGB')
+        
+        # Handle depth array format
+        if len(depth_array.shape) == 3 and depth_array.shape[2] > 1:
+            depth_array = np.mean(depth_array, axis=2)
+            
+        # Normalize depth values
+        depth_min = np.min(depth_array)
+        depth_max = np.max(depth_array)
+        
+        if depth_max > depth_min:
+            normalized_depth = (depth_array - depth_min) / (depth_max - depth_min)
+        else:
+            # Handle case where depth is constant
+            normalized_depth = np.zeros_like(depth_array)
+        
+        # Create progressively blurred versions of the image
+        blurred_images = []
+        for blur_amount in range(max_blur + 1):
+            blurred_images.append(original_image.filter(ImageFilter.GaussianBlur(radius=blur_amount)))
+        
+        # Create output array
+        result_array = np.zeros((512, 512, 3), dtype=np.uint8)
+        
+        # Apply variable blur based on depth
+        height, width = normalized_depth.shape
+        for y in range(height):
+            for x in range(width):
+                blur_radius = int(normalized_depth[y, x] * max_blur)
+                result_array[y, x] = np.array(blurred_images[blur_radius])[y, x]
+        
+        return Image.fromarray(result_array)
     
-    return Image.fromarray(result_array)
+    except Exception as e:
+        logger.error(f"Error in depth-based blur: {e}")
+        logger.error(traceback.format_exc())
+        # Return original image in case of error
+        return original_image
 
 def process_image(input_image, blur_sigma=15, max_depth_blur=30):
-    """Main function to process the image through all effects"""
-    if input_image is None:
-        return None, None, None, None
-    
-    # Resize input image for consistent processing
-    input_image = Image.fromarray(input_image).convert('RGB')
-    input_image = input_image.resize((512, 512))
-    
-    # Step 1: Get segmentation mask
-    mask, _ = get_segmentation_mask(input_image)
-    
-    # Step 2: Apply background blur
-    blurred_background = apply_background_blur(input_image, mask, sigma=blur_sigma)
-    
-    # Step 3: Get depth map
-    depth_map, depth_array = get_depth_map(input_image)
-    
-    # Step 4: Apply depth-based blur
-    depth_blur = apply_depth_based_blur(input_image, depth_array, max_blur=max_depth_blur)
-    
-    # Convert all PIL images to numpy arrays for Gradio
-    input_np = np.array(input_image)
-    mask_np = np.array(mask)
-    blurred_np = np.array(blurred_background)
-    depth_map_np = np.array(depth_map)
-    depth_blur_np = np.array(depth_blur)
+    """Process the image through all steps with error handling"""
+    try:
+        if input_image is None:
+            logger.warning("No input image provided")
+            return None, None, None, None, None
+        
+        # Step 1: Get segmentation mask
+        mask, resized_image = get_segmentation_mask(input_image)
+        
+        # Step 2: Apply background blur
+        blurred_background = apply_background_blur(resized_image, mask, sigma=blur_sigma)
+        
+        # Step 3: Get depth map
+        depth_map, depth_array = get_depth_map(resized_image)
+        
+        # Step 4: Apply depth-based blur
+        depth_blur = apply_depth_based_blur(resized_image, depth_array, max_blur=max_depth_blur)
+        
+        # Convert to numpy arrays for Gradio
+        input_np = np.array(resized_image)
+        mask_np = np.array(mask)
+        blurred_np = np.array(blurred_background)
+        depth_map_np = np.array(depth_map)
+        depth_blur_np = np.array(depth_blur)
+        
+        return input_np, mask_np, blurred_np, depth_map_np, depth_blur_np
     
-    return input_np, mask_np, blurred_np, depth_map_np, depth_blur_np
+    except Exception as e:
+        logger.error(f"Error in image processing: {e}")
+        logger.error(traceback.format_exc())
+        
+        # Create blank outputs in case of error
+        empty = np.zeros((512, 512, 3), dtype=np.uint8)
+        empty_mask = np.zeros((512, 512), dtype=np.uint8)
+        
+        if input_image is not None and isinstance(input_image, np.ndarray):
+            img_resized = cv2.resize(input_image, (512, 512))
+            return img_resized, empty_mask, empty, empty_mask, empty
+        else:
+            return empty, empty_mask, empty, empty_mask, empty
 
 # Create Gradio Interface
-with gr.Blocks(title="Image Blur Effects - EEE 515 Assignment 3") as demo:
-    gr.Markdown("# Image Blur Effects App")
-    gr.Markdown("Upload an image to apply segmentation-based blur and depth-based lens blur effects")
-    
-    with gr.Row():
-        input_image = gr.Image(label="Upload Image", type="numpy")
-        
-    with gr.Row():
-        blur_sigma = gr.Slider(minimum=1, maximum=30, value=15, step=1, label="Background Blur Strength (σ)")
-        depth_blur_max = gr.Slider(minimum=1, maximum=50, value=30, step=1, label="Max Depth Blur Strength")
-    
-    with gr.Row():
-        process_btn = gr.Button("Process Image")
-    
-    with gr.Tab("Segmentation Results"):
-        with gr.Row():
-            original_output = gr.Image(label="Original Image", type="numpy")
-            mask_output = gr.Image(label="Segmentation Mask", type="numpy")
-        with gr.Row():
-            blurred_output = gr.Image(label="Background Blur Effect", type="numpy")
-    
-    with gr.Tab("Depth Results"):
-        with gr.Row():
-            depth_map_output = gr.Image(label="Depth Map", type="numpy")
-            depth_blur_output = gr.Image(label="Depth-Based Lens Blur", type="numpy")
-    
-    process_btn.click(
-        fn=process_image,
-        inputs=[input_image, blur_sigma, depth_blur_max],
-        outputs=[original_output, mask_output, blurred_output, depth_map_output, depth_blur_output]
-    )
-    
-    gr.Markdown("## How it works")
-    gr.Markdown("""
-    1. **Segmentation-Based Blur**: Uses a segmentation model to identify the foreground object, 
-       then applies Gaussian blur only to the background.
-       
-    2. **Depth-Based Lens Blur**: Uses a monocular depth estimation model to create a depth map,
-       then applies varying levels of blur based on the estimated depth.
-    """)
+demo = gr.Interface(
+    fn=process_image,
+    inputs=[
+        gr.Image(type="numpy", label="Upload Image"),
+        gr.Slider(minimum=1, maximum=30, value=15, step=1, label="Background Blur Strength (σ)"),
+        gr.Slider(minimum=1, maximum=50, value=30, step=1, label="Max Depth Blur Strength")
+    ],
+    outputs=[
+        gr.Image(type="numpy", label="Original Image"),
+        gr.Image(type="numpy", label="Segmentation Mask"),
+        gr.Image(type="numpy", label="Background Blur"),
+        gr.Image(type="numpy", label="Depth Map"),
+        gr.Image(type="numpy", label="Depth-Based Lens Blur")
+    ],
+    title="Image Blur Effects - EEE 515 Assignment 3",
+    description="Upload an image to apply segmentation-based blur and depth-based lens blur effects",
+    examples=[["beach.jpeg", 15, 30]],
+    allow_flagging="never"
+)
 
 # Launch the app
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()