Testing

app.py

@@ -1,61 +1,229 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Thu Mar 27 13:56:42 2025
+@author: perghect
+"""
 import gradio as gr
-from transformers import pipeline
-from PIL import Image, ImageFilter
+import requests
+import io
+import torch
 import numpy as np
+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]
+    )
 
-print("Hello")
-
-# 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/zoedepth-nyu-kitti")
-
-def process_image(image, blur_type, sigma):
-    # Step 1: Perform segmentation
-    segmentation_results = segmentation_model(image)
-    foreground_mask = segmentation_results[-1]["mask"]
-
-    # Step 2: Apply Gaussian blur to background
-    blurred_background = image.filter(ImageFilter.GaussianBlur(sigma))
-    segmented_output = Image.composite(image, blurred_background, foreground_mask)
-
-    # Step 3: Perform depth estimation
-    depth_results = depth_estimator(image)
-    depth_map = depth_results["depth"]
-
-    # Step 4: Normalize depth map values
-    depth_array = np.array(depth_map)
-    normalized_depth = (depth_array - np.min(depth_array)) / (np.max(depth_array) - np.min(depth_array)) * 255
-    normalized_depth_image = Image.fromarray(normalized_depth.astype('uint8'))
-
-    # Step 5: Apply variable Gaussian blur based on depth map (Lens Blur)
-    if blur_type == "Lens Blur":
-        variable_blur_image = image.copy()
-        for x in range(variable_blur_image.width):
-            for y in range(variable_blur_image.height):
-                blur_intensity = normalized_depth[y, x] / 255 * sigma  # Scale blur intensity by depth value
-                pixel_value = image.getpixel((x, y))
-                variable_blur_image.putpixel((x, y), tuple(int(p * blur_intensity) for p in pixel_value))
-        output_image = variable_blur_image
-    else:
-        output_image = segmented_output
-
-    return segmented_output, normalized_depth_image, output_image
-
-# Create Gradio interface
-app = 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="Segmented Output with Background Blur"),
-        gr.Image(type="pil", label="Depth Map Visualization"),
-        gr.Image(type="pil", label="Final Output with Selected Blur")
-    ],
-    title="Vision Transformer Segmentation & Depth-Based Blur Effects",
-    description="Upload an image and select the type of blur effect (Gaussian or Lens). Adjust the blur intensity using the slider."
-)
-
-app.launch()
+# Launch the app
+demo.launch()