Update app.py

app.py

@@ -1,229 +1,78 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 """
-Created on Thu Mar 27 13:56:42 2025
-@author: perghect
+@author: Nikhil Kunjoor
 """
 import gradio as gr
-import requests
-import io
-import torch
-import numpy as np
+from transformers import pipeline
 from PIL import Image, ImageFilter
-from torchvision import transforms
-from transformers import AutoModelForImageSegmentation, AutoImageProcessor, AutoModelForDepthEstimation
-
-# Set device and precision
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-torch.set_float32_matmul_precision('high')
-
-# Load models at startup
-rmbg_model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-2.0", trust_remote_code=True).to(device).eval()
-depth_processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
-depth_model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf").to(device)
-
-def load_image_from_link(url: str) -> Image.Image:
-    """Downloads an image from a URL and returns a Pillow Image."""
-    response = requests.get(url)
-    response.raise_for_status()
-    image = Image.open(io.BytesIO(response.content)).convert("RGB")
-    return image
-
-# Gaussian Blur Functions
-def run_rmbg(image: Image.Image, threshold=0.5):
-    """Runs the RMBG-2.0 model on the image and returns a binary mask."""
-    try:
-        image_size = (1024, 1024)
-        transform_image = transforms.Compose([
-            transforms.Resize(image_size),
-            transforms.ToTensor(),
-            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-        ])
-
-        input_images = transform_image(image).unsqueeze(0).to(device)
-
-        with torch.no_grad():
-            preds = rmbg_model(input_images)
-            if isinstance(preds, list):
-                mask_logits = preds[-1]
-            else:
-                raise ValueError(f"Unexpected output format: {type(preds)}")
-
-        mask_prob = mask_logits.sigmoid().cpu()[0].squeeze()
-        pred_pil = transforms.ToPILImage()(mask_prob)
-        mask_pil = pred_pil.resize(image.size, resample=Image.BILINEAR)
-
-        mask_np = np.array(mask_pil, dtype=np.uint8) / 255.0
-        binary_mask = (mask_np > threshold).astype(np.uint8)
-        return binary_mask
-    except Exception as e:
-        raise Exception(f"Error in background removal: {str(e)}")
-
-def apply_background_blur(image: Image.Image, mask: np.ndarray, sigma: float = 15):
-    """Applies a Gaussian blur to the background while keeping the foreground sharp."""
-    image_np = np.array(image)
-    mask_np = mask.astype(np.uint8)
-
-    blurred_image = image.filter(ImageFilter.GaussianBlur(radius=sigma))
-    blurred_np = np.array(blurred_image)
-
-    output_np = np.where(mask_np[..., None] == 1, image_np, blurred_np)
-    output_image = Image.fromarray(output_np.astype(np.uint8))
-    return output_image
-
-# Lens Blur Functions
-def run_depth_estimation(image: Image.Image, target_size=(512, 512)):
-    """Runs the Depth-Anything-V2-Small model and returns the depth map."""
-    try:
-        image_resized = image.resize(target_size, resample=Image.BILINEAR)
-        inputs = depth_processor(images=image_resized, return_tensors="pt").to(device)
-
-        with torch.no_grad():
-            outputs = depth_model(**inputs)
-            predicted_depth = outputs.predicted_depth
-
-        prediction = torch.nn.functional.interpolate(
-            predicted_depth.unsqueeze(1),
-            size=image.size[::-1],
-            mode="bicubic",
-            align_corners=False,
-        )
-
-        depth_map = prediction.squeeze().cpu().numpy()
-        depth_max = depth_map.max()
-        depth_min = depth_map.min()
-        if depth_max == depth_min:
-            depth_max = depth_min + 1e-6  # Avoid division by zero
-        depth_map = (depth_map - depth_min) / (depth_max - depth_min)
-        depth_map = 1 - depth_map  # Invert: higher values = farther
-        return depth_map
-    except Exception as e:
-        raise Exception(f"Error in depth estimation: {str(e)}")
-
-def apply_depth_based_blur(image: Image.Image, depth_map: np.ndarray, max_radius: float = 15, foreground_percentile: float = 30):
-    """Applies a variable Gaussian blur based on the depth map."""
-    image_np = np.array(image)
-
-    if depth_map.shape != image_np.shape[:2]:
-        depth_map = np.array(Image.fromarray(depth_map).resize(image.size, resample=Image.BILINEAR))
-
-    foreground_threshold = np.percentile(depth_map.flatten(), foreground_percentile)
-
-    output_np = np.zeros_like(image_np)
-    mask_foreground = (depth_map <= foreground_threshold)
-    output_np[mask_foreground] = image_np[mask_foreground]
-
-    depth_max = depth_map.max()
-    depth_range = depth_max - foreground_threshold
-    if depth_range == 0:
-        depth_range = 1e-6
-    normalized_depth = np.zeros_like(depth_map)
-    mask_above_foreground = (depth_map > foreground_threshold)
-    normalized_depth[mask_above_foreground] = (depth_map[mask_above_foreground] - foreground_threshold) / depth_range
-    normalized_depth = np.clip(normalized_depth, 0, 1)
-
-    depth_levels = np.linspace(0, 1, 20)
-    for i in range(len(depth_levels) - 1):
-        depth_min = depth_levels[i]
-        depth_max = depth_levels[i + 1]
-        mask = (normalized_depth >= depth_min) & (normalized_depth < depth_max) & (depth_map > foreground_threshold)
-        if not np.any(mask):
-            continue
-
-        avg_depth = (depth_min + depth_max) / 2
-        blur_radius = max_radius * avg_depth
-
-        blurred_image = image.filter(ImageFilter.GaussianBlur(radius=blur_radius))
-        blurred_np = np.array(blurred_image)
-        output_np[mask] = blurred_np[mask]
-
-    mask_farthest = (normalized_depth >= depth_levels[-1]) & (depth_map > foreground_threshold)
-    if np.any(mask_farthest):
-        blurred_max = image.filter(ImageFilter.GaussianBlur(radius=max_radius))
-        output_np[mask_farthest] = np.array(blurred_max)[mask_farthest]
-
-    output_image = Image.fromarray(output_np.astype(np.uint8))
-    return output_image
-
-# Main Processing Function for Gradio
-def process_image(image, blur_type, sigma=15, max_radius=15, foreground_percentile=30):
-    """Processes the image based on the selected blur type."""
-    if image is None:
-        return None, "Please upload an image."
-
-    try:
-        image = Image.fromarray(image).convert("RGB")
-    except Exception as e:
-        return None, f"Error processing image: {str(e)}"
-
-    # Resize image if too large
-    max_size = (1024, 1024)
-    if image.size[0] > max_size[0] or image.size[1] > max_size[1]:
-        image.thumbnail(max_size, Image.Resampling.LANCZOS)
-
-    try:
-        if blur_type == "Gaussian Blur":
-            mask = run_rmbg(image, threshold=0.5)
-            output_image = apply_background_blur(image, mask, sigma=sigma)
-            title = f"Gaussian Blur (sigma={sigma})"
-        else:  # Lens Blur
-            depth_map = run_depth_estimation(image, target_size=(512, 512))
-            output_image = apply_depth_based_blur(image, depth_map, max_radius=max_radius, foreground_percentile=foreground_percentile)
-            title = f"Lens Blur (max_radius={max_radius}, foreground_percentile={foreground_percentile})"
-    except Exception as e:
-        return None, f"Error applying blur: {str(e)}"
-
-    return output_image, title
-
-# Gradio Interface with Conditional Parameter Display
-with gr.Blocks() as demo:
-    gr.Markdown("# Image Blur Effects with Gaussian and Lens Blur")
-    gr.Markdown("""
-    This app applies blur effects to your images. Follow these steps to use it:
-    **Note**: This app is hosted on Hugging Face Spaces’ free tier and may go to "Sleeping" mode after 48 hours of inactivity. If it doesn’t load immediately, please wait a few seconds while it wakes up.
-    1. **Upload an Image**: Click the "Upload Image" box to upload an image from your device.
-    2. **Choose a Blur Type**:
-       - **Gaussian Blur**: Applies a uniform blur to the background, keeping the foreground sharp. Adjust the sigma parameter to control blur intensity.
-       - **Lens Blur**: Applies a depth-based blur, simulating a depth-of-field effect (closer objects are sharp, farther objects are blurred). Adjust the max radius and foreground percentile to fine-tune the effect.
-    3. **Adjust Parameters**:
-       - For Gaussian Blur, use the "Gaussian Blur Sigma" slider to control blur intensity (higher values = more blur).
-       - For Lens Blur, use the "Max Blur Radius" slider to control the maximum blur intensity and the "Foreground Percentile" slider to adjust the depth threshold for the foreground.
-    4. **Apply the Blur**: Click the "Apply Blur" button to process the image.
-    5. **View the Result**: The processed image will appear in the "Output Image" box, along with a description of the effect applied.
-    **Example**: Try uploading an image with a clear foreground and background (e.g., a person in front of a landscape) to see the effects in action.
-    """)
-
-    with gr.Row():
-        image_input = gr.Image(label="Upload Image", type="numpy")
-        with gr.Column():
-            blur_type = gr.Radio(choices=["Gaussian Blur", "Lens Blur"], label="Blur Type", value="Gaussian Blur")
-            sigma = gr.Slider(minimum=1, maximum=50, step=1, value=15, label="Gaussian Blur Sigma", visible=True)
-            max_radius = gr.Slider(minimum=1, maximum=50, step=1, value=15, label="Max Lens Blur Radius", visible=False)
-            foreground_percentile = gr.Slider(minimum=1, maximum=50, step=1, value=30, label="Foreground Percentile", visible=False)
-
-    # Update visibility of parameters based on blur type
-    def update_visibility(blur_type):
-        if blur_type == "Gaussian Blur":
-            return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
-        else:  # Lens Blur
-            return gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
-
-    blur_type.change(
-        fn=update_visibility,
-        inputs=blur_type,
-        outputs=[sigma, max_radius, foreground_percentile]
-    )
-
-    process_button = gr.Button("Apply Blur")
-    with gr.Row():
-        output_image = gr.Image(label="Output Image")
-        output_text = gr.Textbox(label="Effect Applied")
-
-    process_button.click(
-        fn=process_image,
-        inputs=[image_input, blur_type, sigma, max_radius, foreground_percentile],
-        outputs=[output_image, output_text]
-    )
+import numpy as np
+import torch
 
-# Launch the app
-demo.launch()
+# Load models from Hugging Face
+segmentation_model = pipeline("image-segmentation", model="nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
+depth_estimator = pipeline("depth-estimation", model="Intel/dpt-large")
+
+def apply_gaussian_blur(image, mask, sigma):
+    blurred = image.filter(ImageFilter.GaussianBlur(sigma))
+    return Image.composite(image, blurred, mask)
+
+def apply_lens_blur(image, depth_map, sigma):
+    depth_array = np.array(depth_map)
+    normalized_depth = (depth_array - np.min(depth_array)) / (np.max(depth_array) - np.min(depth_array))
+    
+    blurred = image.copy()
+    for x in range(image.width):
+        for y in range(image.height):
+            blur_intensity = normalized_depth[y, x] * sigma
+            local_blur = image.crop((x-1, y-1, x+2, y+2)).filter(ImageFilter.GaussianBlur(blur_intensity))
+            blurred.putpixel((x, y), local_blur.getpixel((1, 1)))
+    return blurred
+
+def process_image(image, blur_type, sigma):
+    # Perform segmentation
+    segmentation_results = segmentation_model(image)
+    person_mask = None
+    for segment in segmentation_results:
+        if segment['label'] == 'person':
+            person_mask = Image.fromarray((segment['mask'] * 255).astype(np.uint8))
+            break
+    
+    if person_mask is None:
+        person_mask = Image.new('L', image.size, 255)  # Create a white mask if no person is detected
+    
+    # Perform depth estimation
+    depth_results = depth_estimator(image)
+    depth_map = depth_results["depth"]
+    
+    # Normalize depth map for visualization
+    depth_array = np.array(depth_map)
+    normalized_depth = (depth_array - np.min(depth_array)) / (np.max(depth_array) - np.min(depth_array)) * 255
+    depth_visualization = Image.fromarray(normalized_depth.astype(np.uint8))
+    
+    # Apply selected blur effect
+    if blur_type == "Gaussian Blur":
+        output_image = apply_gaussian_blur(image, person_mask, sigma)
+    else:  # Lens Blur
+        output_image = apply_lens_blur(image, depth_map, sigma)
+    
+    return person_mask, depth_visualization, output_image
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=process_image,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Radio(["Gaussian Blur", "Lens Blur"], label="Blur Type", value="Gaussian Blur"),
+        gr.Slider(0, 50, step=1, label="Blur Intensity (Sigma)", value=15)
+    ],
+    outputs=[
+        gr.Image(type="pil", label="Segmentation Mask"),
+        gr.Image(type="pil", label="Depth Map"),
+        gr.Image(type="pil", label="Output Image")
+    ],
+    title="Vision Transformer Segmentation & Depth-Based Blur Effects",
+    description="Upload an image to apply segmentation and lens blur effects. Adjust the blur type and intensity using the controls below."
+)
+
+iface.launch()