Update app.py

Added Code to the application

app.py

@@ -1,21 +1,59 @@
 import gradio as gr
+from transformers import pipeline
+from PIL import Image, ImageFilter
+import numpy as np
 
-def process_image(image):
-    # Perform segmentation and apply Gaussian blur (steps above)
-    # Return output images for display
-    
-    segmented_output = ... # Segmented output with blurred background
-    depth_map_output = ... # Depth map visualization
-    variable_blur_output = ... # Variable Gaussian blur
-    
-    return segmented_output, depth_map_output, variable_blur_output
+# 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"),
-    outputs=[gr.Image(type="pil"), gr.Image(type="pil"), gr.Image(type="pil")],
-    title="Vision Transformer Segmentation & Depth Estimation",
-    description="Upload an image to apply segmentation and lens blur effects."
+    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()