Update app.py

app.py

@@ -16,242 +16,6 @@ seg_model = AutoModelForSemanticSegmentation.from_pretrained("nvidia/segformer-b
 depth_processor = DPTImageProcessor.from_pretrained("Intel/dpt-large")
 depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
 
-def apply_gaussian_blur(image, mask, sigma=15):
-    """Apply Gaussian blur to the background of an image based on a mask."""
-    # 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 = cv2.resize(binary_mask, (image.shape[1], image.shape[0]))
-    
-    # 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)
-
-def apply_depth_blur(image, depth_map, max_sigma=25):
-    """Apply variable Gaussian blur based on depth map."""
-    # 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 = cv2.resize(depth_norm, (image.shape[1], image.shape[0]))
-    
-    # Create output image
-    result = np.zeros_like(image)
-    
-    # Apply different blur levels based on depth
-    for sigma in range(1, int(max_sigma) + 1, 2):
-        # Create a mask for pixels at this depth level
-        lower_bound = (sigma - 1) / max_sigma
-        upper_bound = (sigma + 1) / max_sigma
-        mask = np.logical_and(depth_norm >= lower_bound, depth_norm <= upper_bound).astype(np.float32)
-        
-        # Skip if no pixels at this depth
-        if not np.any(mask):
-            continue
-        
-        # Blur the image with current sigma
-        blurred = cv2.GaussianBlur(image, (0, 0), sigma)
-        
-        # Create a 3-channel mask if the input mask is single-channel
-        if len(mask.shape) == 2:
-            mask_3ch = np.stack([mask, mask, mask], axis=2)
-        else:
-            mask_3ch = mask
-        
-        # Add the blurred pixels at this depth to the result
-        result += (blurred * mask_3ch).astype(np.uint8)
-    
-    # Fill in any missing pixels (where sum of all masks < 1)
-    total_mask = np.zeros_like(depth_norm)
-    for sigma in range(1, int(max_sigma) + 1, 2):
-        lower_bound = (sigma - 1) / max_sigma
-        upper_bound = (sigma + 1) / max_sigma
-        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
-
-def get_segmentation_mask(image_pil):
-    """Get segmentation mask for person/foreground from an image."""
-    # Resize the image to the size expected by the segmentation model
-    width, height = image_pil.size
-    image_pil_resized = image_pil.resize((512, 512))
-    
-    # Process the image with the segmentation model
-    inputs = seg_processor(images=image_pil_resized, 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=(512, 512),
-        mode="bilinear",
-        align_corners=False,
-    )
-    
-    # Get the predicted segmentation mask
-    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 - considering classes that are likely to be foreground
-    # The ADE20K dataset has 150 classes, so we need to choose which ones to consider as foreground
-    # Common foreground classes: person (12), animal classes, and objects like furniture
-    # This is a simplified approach - you may need to adjust based on your needs
-    foreground_classes = [12, 13, 14, 15, 16, 17, 18, 19, 20]  # Person and some objects
-    foreground_mask = np.zeros_like(mask_np)
-    for cls in foreground_classes:
-        foreground_mask[mask_np == cls] = 1
-    
-    # Resize back to original image size
-    foreground_mask = cv2.resize(foreground_mask, (width, height))
-    
-    return foreground_mask
-
-def get_depth_map(image_pil):
-    """Get depth map from an image."""
-    # 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_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
-    
-    return depth_map
-
-def process_image(input_image, blur_sigma=15, depth_blur_sigma=25):
-    """Main function to process the input image."""
-    try:
-        # Convert to PIL Image if needed
-        if isinstance(input_image, np.ndarray):
-            pil_image = Image.fromarray(input_image)
-        else:
-            pil_image = input_image
-            input_image = np.array(pil_image)
-        
-        # Get segmentation mask
-        seg_mask = get_segmentation_mask(pil_image)
-        
-        # Get depth map
-        depth_map = get_depth_map(pil_image)
-        
-        # Apply gaussian blur to background
-        gaussian_result = apply_gaussian_blur(input_image, seg_mask, sigma=blur_sigma)
-        
-        # Apply 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_visualization = cv2.applyColorMap(depth_visualization, cv2.COLORMAP_INFERNO)
-        
-        # Display segmentation mask
-        seg_visualization = (seg_mask * 255).astype(np.uint8)
-        
-        return [
-            input_image, 
-            seg_visualization, 
-            gaussian_result, 
-            depth_visualization, 
-            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="pil")
-            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
-    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:
@@ -535,11 +299,4 @@ with gr.Blocks(title="Image Blur Effects with Segmentation and Depth Estimation"
     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
-    
-    Try uploading a photo of a person or object against a background to see the effects!
-    """)
-
 demo.launch()