Update app.py

app.py

@@ -234,7 +234,308 @@ with gr.Blocks(title="Image Blur Effects with Segmentation and Depth Estimation"
     
     gr.Markdown("""
     ## How it works
+    import gradio as gr
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+from transformers import AutoImageProcessor, AutoModelForSemanticSegmentation
+from transformers import DPTImageProcessor, DPTForDepthEstimation
+import warnings
+warnings.filterwarnings("ignore")
+
+# Load segmentation model - using SegFormer which is compatible with AutoModelForSemanticSegmentation
+seg_processor = AutoImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
+seg_model = AutoModelForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
+
+# Load depth estimation model
+depth_processor = DPTImageProcessor.from_pretrained("Intel/dpt-large")
+depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
+
+def safe_resize(image, target_size, interpolation=cv2.INTER_LINEAR):
+    """Safely resize an image with validation checks."""
+    if image is None:
+        return None
+    
+    # Ensure image is a proper numpy array
+    if not isinstance(image, np.ndarray):
+        return None
+    
+    # Check that dimensions are valid (non-zero)
+    h, w = target_size
+    if h <= 0 or w <= 0 or image.shape[0] <= 0 or image.shape[1] <= 0:
+        return image  # Return original if target dimensions are invalid
+    
+    # Handle grayscale images differently
+    if len(image.shape) == 2:
+        return cv2.resize(image, (w, h), interpolation=interpolation)
+    else:
+        return cv2.resize(image, (w, h), interpolation=interpolation)
+
+def apply_gaussian_blur(image, mask, sigma=15):
+    """Apply Gaussian blur to the background of an image based on a mask."""
+    try:
+        # Convert mask to binary (0 and 255)
+        if mask.max() <= 1.0:
+            binary_mask = (mask * 255).astype(np.uint8)
+        else:
+            binary_mask = mask.astype(np.uint8)
+        
+        # Create a blurred version of the entire image
+        blurred = cv2.GaussianBlur(image, (0, 0), sigma)
+        
+        # Resize mask to match image dimensions if needed
+        if binary_mask.shape[:2] != image.shape[:2]:
+            binary_mask = safe_resize(binary_mask, (image.shape[0], image.shape[1]))
+        
+        # Create a 3-channel mask if the input mask is single-channel
+        if len(binary_mask.shape) == 2:
+            mask_3ch = np.stack([binary_mask, binary_mask, binary_mask], axis=2)
+        else:
+            mask_3ch = binary_mask
+        
+        # Normalize mask to range [0, 1]
+        mask_3ch = mask_3ch / 255.0
+        
+        # Combine original image (foreground) with blurred image (background) using the mask
+        result = image * mask_3ch + blurred * (1 - mask_3ch)
+        
+        return result.astype(np.uint8)
+    except Exception as e:
+        print(f"Error in apply_gaussian_blur: {e}")
+        return image  # Return original image if there's an error
+
+def apply_depth_blur(image, depth_map, max_sigma=25):
+    """Apply variable Gaussian blur based on depth map."""
+    try:
+        # Normalize depth map to range [0, 1]
+        if depth_map.max() > 1.0:
+            depth_norm = depth_map / depth_map.max()
+        else:
+            depth_norm = depth_map
+        
+        # Resize depth map to match image dimensions if needed
+        if depth_norm.shape[:2] != image.shape[:2]:
+            depth_norm = safe_resize(depth_norm, (image.shape[0], image.shape[1]))
+        
+        # Create output image
+        result = np.zeros_like(image)
+        
+        # Instead of many small blurs, use fewer blur levels for efficiency
+        blur_levels = 5
+        step = max_sigma / blur_levels
+        
+        for i in range(blur_levels):
+            sigma = (i + 1) * step
+            
+            # Calculate depth range for this blur level
+            lower_bound = i / blur_levels
+            upper_bound = (i + 1) / blur_levels
+            
+            # Create mask for pixels in this depth range
+            mask = np.logical_and(depth_norm >= lower_bound, depth_norm <= upper_bound).astype(np.float32)
+            
+            # Skip if no pixels in this range
+            if not np.any(mask):
+                continue
+            
+            # Apply blur for this level
+            blurred = cv2.GaussianBlur(image, (0, 0), sigma)
+            
+            # Create 3-channel mask
+            mask_3ch = np.stack([mask, mask, mask], axis=2) if len(mask.shape) == 2 else mask
+            
+            # Add to result
+            result += (blurred * mask_3ch).astype(np.uint8)
+        
+        # Check if there are any pixels not covered and fill with original
+        total_mask = np.zeros_like(depth_norm)
+        for i in range(blur_levels):
+            lower_bound = i / blur_levels
+            upper_bound = (i + 1) / blur_levels
+            mask = np.logical_and(depth_norm >= lower_bound, depth_norm <= upper_bound).astype(np.float32)
+            total_mask += mask
+        
+        missing_mask = (total_mask < 0.5).astype(np.float32)
+        if np.any(missing_mask):
+            missing_mask_3ch = np.stack([missing_mask, missing_mask, missing_mask], axis=2)
+            result += (image * missing_mask_3ch).astype(np.uint8)
+        
+        return result
+    except Exception as e:
+        print(f"Error in apply_depth_blur: {e}")
+        return image  # Return original image if there's an error
+
+def get_segmentation_mask(image_pil):
+    """Get segmentation mask for person/foreground from an image."""
+    try:
+        # Process the image with the segmentation model
+        inputs = seg_processor(images=image_pil, return_tensors="pt")
+        with torch.no_grad():
+            outputs = seg_model(**inputs)
+        
+        # Get the predicted segmentation mask
+        logits = outputs.logits
+        upsampled_logits = torch.nn.functional.interpolate(
+            logits,
+            size=image_pil.size[::-1],  # Resize directly to original size
+            mode="bilinear",
+            align_corners=False,
+        )
+        
+        # Get the predicted class for each pixel
+        predicted_mask = upsampled_logits.argmax(dim=1)[0]
+        
+        # Convert the mask to a numpy array
+        mask_np = predicted_mask.cpu().numpy()
+        
+        # Create a foreground mask - human and common foreground objects
+        # Classes based on ADE20K dataset
+        foreground_classes = [12]  # Person class (you can add more classes as needed)
+        
+        # Create a binary mask for foreground classes
+        foreground_mask = np.zeros_like(mask_np)
+        for cls in foreground_classes:
+            foreground_mask[mask_np == cls] = 1
+        
+        return foreground_mask
+    except Exception as e:
+        print(f"Error in get_segmentation_mask: {e}")
+        # Return a default mask (all ones) in case of error
+        return np.ones((image_pil.size[1], image_pil.size[0]), dtype=np.uint8)
+
+def get_depth_map(image_pil):
+    """Get depth map from an image."""
+    try:
+        # Process the image with the depth estimation model
+        inputs = depth_processor(images=image_pil, return_tensors="pt")
+        with torch.no_grad():
+            outputs = depth_model(**inputs)
+            predicted_depth = outputs.predicted_depth
+        
+        # Interpolate to original size
+        prediction = torch.nn.functional.interpolate(
+            predicted_depth.unsqueeze(1),
+            size=image_pil.size[::-1],
+            mode="bicubic",
+            align_corners=False,
+        )
+        
+        # Convert to numpy array
+        depth_map = prediction.squeeze().cpu().numpy()
+        
+        # Normalize depth map
+        depth_min = depth_map.min()
+        depth_max = depth_map.max()
+        if depth_max > depth_min:
+            depth_map = (depth_map - depth_min) / (depth_max - depth_min)
+        else:
+            depth_map = np.zeros_like(depth_map)
+        
+        return depth_map
+    except Exception as e:
+        print(f"Error in get_depth_map: {e}")
+        # Return a default depth map (gradient from top to bottom) in case of error
+        h, w = image_pil.size[1], image_pil.size[0]
+        default_depth = np.zeros((h, w), dtype=np.float32)
+        for i in range(h):
+            default_depth[i, :] = i / h
+        return default_depth
+
+def process_image(input_image, blur_sigma=15, depth_blur_sigma=25):
+    """Main function to process the input image."""
+    try:
+        # Input validation
+        if input_image is None:
+            print("No input image provided")
+            return [None, None, None, None, None]
+        
+        # Convert to PIL Image if needed
+        if isinstance(input_image, np.ndarray):
+            # Make sure we have a valid image with at least 2 dimensions
+            if input_image.ndim < 2 or input_image.shape[0] <= 0 or input_image.shape[1] <= 0:
+                print("Invalid input image dimensions")
+                return [None, None, None, None, None]
+            pil_image = Image.fromarray(input_image)
+        else:
+            pil_image = input_image
+            input_image = np.array(pil_image)
+        
+        # Get segmentation mask
+        print("Getting segmentation mask...")
+        seg_mask = get_segmentation_mask(pil_image)
+        
+        # Get depth map
+        print("Getting depth map...")
+        depth_map = get_depth_map(pil_image)
+        
+        # Apply gaussian blur to background
+        print("Applying gaussian blur...")
+        gaussian_result = apply_gaussian_blur(input_image, seg_mask, sigma=blur_sigma)
+        
+        # Apply depth-based blur
+        print("Applying depth-based blur...")
+        depth_result = apply_depth_blur(input_image, depth_map, max_sigma=depth_blur_sigma)
+        
+        # Display depth map as an image
+        depth_visualization = (depth_map * 255).astype(np.uint8)
+        depth_colored = cv2.applyColorMap(depth_visualization, cv2.COLORMAP_INFERNO)
+        
+        # Display segmentation mask
+        seg_visualization = (seg_mask * 255).astype(np.uint8)
+        
+        print("Processing complete!")
+        return [
+            input_image, 
+            seg_visualization, 
+            gaussian_result, 
+            depth_colored, 
+            depth_result
+        ]
+    except Exception as e:
+        print(f"Error processing image: {e}")
+        return [None, None, None, None, None]
+
+# Create Gradio interface
+with gr.Blocks(title="Image Blur Effects with Segmentation and Depth Estimation") as demo:
+    gr.Markdown("# Image Blur Effects App")
+    gr.Markdown("This app demonstrates two types of blur effects: background blur using segmentation and depth-based lens blur.")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label="Upload an image", type="numpy")
+            blur_sigma = gr.Slider(minimum=1, maximum=50, value=15, step=1, label="Background Blur Intensity")
+            depth_blur_sigma = gr.Slider(minimum=1, maximum=50, value=25, step=1, label="Depth Blur Max Intensity")
+            process_btn = gr.Button("Process Image")
+        
+        with gr.Column():
+            with gr.Tab("Original Image"):
+                output_original = gr.Image(label="Original Image")
+            with gr.Tab("Segmentation Mask"):
+                output_segmentation = gr.Image(label="Segmentation Mask")
+            with gr.Tab("Background Blur"):
+                output_gaussian = gr.Image(label="Background Blur Result")
+            with gr.Tab("Depth Map"):
+                output_depth = gr.Image(label="Depth Map")
+            with gr.Tab("Depth-based Lens Blur"):
+                output_depth_blur = gr.Image(label="Depth-based Lens Blur Result")
+    
+    process_btn.click(
+        fn=process_image,
+        inputs=[input_image, blur_sigma, depth_blur_sigma],
+        outputs=[output_original, output_segmentation, output_gaussian, output_depth, output_depth_blur]
+    )
     
+    gr.Markdown("""
+    ## How it works
+    
+    1. **Background Blur**: Uses a SegFormer model to identify foreground objects (like people) and blurs only the background
+    2. **Depth-based Lens Blur**: Uses a DPT depth estimation model to apply variable blur based on estimated distance
+    
+    Try uploading a photo of a person against a background to see the effects!
+    """)
+
+demo.launch()
     1. **Background Blur**: Uses a segmentation model to identify foreground objects and blurs only the background
     2. **Depth-based Lens Blur**: Uses a depth estimation model to apply variable blur based on estimated distance