Update app.py

app.py

@@ -38,8 +38,9 @@ depth_pipeline = pipeline("depth-estimation", model="depth-anything/Depth-Anythi
 
 def segment_and_blur_background(input_image: Image.Image, blur_radius: int = 15, threshold: float = 0.5) -> Image.Image:
     """
-    Applies segmentation using the RMBG-2.0 model and then uses the segmentation mask
-    to composite a Gaussian-blurred background with a sharp foreground.
+    Uses the RMBG-2.0 segmentation model to create a binary mask,
+    then composites a Gaussian-blurred background with the sharp foreground.
+    The segmentation threshold is adjustable.
     """
     # Ensure the image is in RGB and get its original dimensions
     image = input_image.convert("RGB")
@@ -53,7 +54,7 @@ def segment_and_blur_background(input_image: Image.Image, blur_radius: int = 15,
         preds = segmentation_model(input_tensor)[-1].sigmoid().cpu()
     pred = preds[0].squeeze()
     
-    # Create a binary mask using the threshold
+    # Create a binary mask using the adjustable threshold
     binary_mask = (pred > threshold).float()
     mask_pil = transforms.ToPILImage()(binary_mask).convert("L")
     # Convert grayscale mask to pure binary (0 or 255)
@@ -63,15 +64,15 @@ def segment_and_blur_background(input_image: Image.Image, blur_radius: int = 15,
     
     # Apply Gaussian blur to the entire image for background
     blurred_image = image.filter(ImageFilter.GaussianBlur(blur_radius))
-    # Composite the original image (foreground) with the blurred image (background) using the mask
+    # Composite the original image (foreground) with the blurred background using the mask
     final_image = Image.composite(image, blurred_image, mask_pil)
     return final_image
 
 
 def depth_based_lens_blur(input_image: Image.Image, max_blur: float = 2, num_bands: int = 40, invert_depth: bool = False) -> Image.Image:
     """
-    Applies a depth-based blur effect using a depth map produced by Depth-Anything.
-    The effect simulates a lens blur by applying different blur strengths in depth bands.
+    Applies a depth-based blur effect using a depth map from Depth-Anything.
+    The max_blur parameter (controlled by a slider) sets the highest blur intensity.
     """
     # Resize the input image to 512x512 for the depth estimation model
     image_resized = input_image.resize((512, 512))
@@ -91,7 +92,7 @@ def depth_based_lens_blur(input_image: Image.Image, max_blur: float = 2, num_ban
     orig_rgba = image_resized.convert("RGBA")
     final_image = orig_rgba.copy()
     
-    # Divide the normalized depth range into bands
+    # Divide the normalized depth range into bands and apply variable blur
     band_edges = np.linspace(0, 1, num_bands + 1)
     for i in range(num_bands):
         band_min = band_edges[i]
@@ -114,16 +115,20 @@ def depth_based_lens_blur(input_image: Image.Image, max_blur: float = 2, num_ban
     return final_image.convert("RGB")
 
 
-def process_image(input_image: Image.Image, effect: str) -> Image.Image:
+def process_image(input_image: Image.Image, effect: str, threshold: float, blur_intensity: float) -> Image.Image:
     """
     Dispatch function to apply the selected effect:
-    - "Gaussian Blur Background": uses segmentation and Gaussian blur.
-    - "Depth-based Lens Blur": applies depth-based blur using the estimated depth map.
+      - "Gaussian Blur Background": uses segmentation with an adjustable threshold and blur radius.
+      - "Depth-based Lens Blur": applies depth-based blur with an adjustable maximum blur.
+    The threshold slider is used only for the segmentation effect.
+    The blur_intensity slider controls the blur strength in both effects.
     """
     if effect == "Gaussian Blur Background":
-        return segment_and_blur_background(input_image)
+        # For segmentation, use the threshold and blur_intensity (as blur_radius)
+        return segment_and_blur_background(input_image, blur_radius=int(blur_intensity), threshold=threshold)
     elif effect == "Depth-based Lens Blur":
-        return depth_based_lens_blur(input_image)
+        # For depth-based blur, use the blur_intensity as the max blur value.
+        return depth_based_lens_blur(input_image, max_blur=blur_intensity)
     else:
         return input_image
 
@@ -136,13 +141,15 @@ iface = gr.Interface(
     fn=process_image,
     inputs=[
         gr.Image(type="pil", label="Input Image"),
-        gr.Radio(choices=["Gaussian Blur Background", "Depth-based Lens Blur"], label="Select Effect")
+        gr.Radio(choices=["Gaussian Blur Background", "Depth-based Lens Blur"], label="Select Effect"),
+        gr.Slider(0.0, 1.0, value=0.5, label="Segmentation Threshold (for Gaussian Blur)"),
+        gr.Slider(0, 30, value=15, step=1, label="Blur Intensity (for both effects)")
     ],
     outputs=gr.Image(type="pil", label="Output Image"),
-    title="Blur Effects Demo",
+    title="Interactive Blur Effects Demo",
     description=(
-        "Upload an image and choose an effect: "
-        "apply segmentation + Gaussian blurred background, or a depth-based lens blur effect."
+        "Upload an image and choose an effect. For 'Gaussian Blur Background', adjust the segmentation threshold and blur intensity. "
+        "For 'Depth-based Lens Blur', the blur intensity slider sets the maximum blur based on depth."
     )
 )