Update app.py

app.py

@@ -1,58 +1,88 @@
 import gradio as gr
 from transformers import pipeline
-from PIL import Image, ImageFilter
+from PIL import Image
 import numpy as np
+from scipy.ndimage import gaussian_filter
+import matplotlib.pyplot as plt
 
-# Load models from Hugging Face
-segmentation_model = pipeline("image-segmentation", model="nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
+# Load the depth estimation model
 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"]
+def process_image(input_image):
+    # Convert Gradio input (numpy array) to PIL Image
+    input_image = Image.fromarray(input_image.astype('uint8'), 'RGB')
 
-    # 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)
+    # Perform depth estimation
+    depth_results = depth_estimator(input_image)
     depth_map = depth_results["depth"]
 
-    # Step 4: Normalize depth map values
+    # Convert depth map to numpy array and normalize to [0, 1]
     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(
+    normalized_depth = (depth_array - np.min(depth_array)) / (np.max(depth_array) - np.min(depth_array))
+
+    # Convert input image to numpy array
+    img_array = np.array(input_image)
+
+    # Create variable blur effect
+    max_blur = 5.0  # Maximum blur radius
+    min_blur = 0.5  # Minimum blur radius to avoid completely sharp areas
+    n_steps = 10  # Number of blur levels
+
+    # Create output array
+    blurred_array = np.zeros_like(img_array, dtype=np.float32)
+
+    # Apply variable blur by processing the image with multiple blur levels
+    for i in range(n_steps):
+        sigma = min_blur + (max_blur - min_blur) * i / (n_steps - 1)
+        # Apply Gaussian blur with current sigma to the whole image
+        blurred_r = gaussian_filter(img_array[:,:,0], sigma=sigma)
+        blurred_g = gaussian_filter(img_array[:,:,1], sigma=sigma)
+        blurred_b = gaussian_filter(img_array[:,:,2], sigma=sigma)
+        blurred_temp = np.stack([blurred_r, blurred_g, blurred_b], axis=2)
+        
+        # Create a mask for this blur level
+        lower_bound = i / n_steps
+        upper_bound = (i + 1) / n_steps
+        mask = (normalized_depth >= lower_bound) & (normalized_depth < upper_bound)
+        mask = mask[..., np.newaxis]  # Add channel dimension
+        
+        # Apply this blur level to the appropriate regions
+        blurred_array = np.where(mask, blurred_temp, blurred_array)
+
+    # Convert back to uint8
+    blurred_image = Image.fromarray(blurred_array.astype('uint8'))
+
+    # Create side-by-side visualization
+    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+    axes[0].imshow(input_image)
+    axes[0].set_title("Input Image")
+    axes[0].axis("off")
+
+    axes[1].imshow(depth_map, cmap="gray")
+    axes[1].set_title("Depth Map")
+    axes[1].axis("off")
+
+    axes[2].imshow(blurred_image)
+    axes[2].set_title("Variable Blur Output")
+    axes[2].axis("off")
+
+    plt.tight_layout()
+
+    # Save the figure to a temporary file or buffer to display in Gradio
+    output_path = "output.png"
+    plt.savefig(output_path, bbox_inches='tight')
+    plt.close()
+
+    return output_path
+
+# Define Gradio interface
+interface = 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."
+    inputs=gr.Image(label="Upload an Image"),
+    outputs=gr.Image(label="Processed Output"),
+    title="Depth-Based Variable Blur App",
+    description="Upload an image to apply a variable blur effect based on depth estimation."
 )
-app.launch()
+
+# Launch the app
+interface.launch()