Update app.py

app.py

@@ -10,7 +10,6 @@ from transformers import AutoModelForImageSegmentation, pipeline
 # Global Setup and Model Loading
 # ----------------------------
 
-# Set device (GPU if available, else CPU)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load the segmentation model (RMBG-2.0)
@@ -21,7 +20,7 @@ segmentation_model = AutoModelForImageSegmentation.from_pretrained(
 segmentation_model.to(device)
 segmentation_model.eval()
 
-# Define the image transformation for segmentation (resize to 512x512)
+# Transformation for segmentation (resizes to 512 for the model input)
 image_size = (512, 512)
 segmentation_transform = transforms.Compose([
     transforms.Resize(image_size),
@@ -36,122 +35,120 @@ depth_pipeline = pipeline("depth-estimation", model="depth-anything/Depth-Anythi
 # Processing Functions
 # ----------------------------
 
-def segment_and_blur_background(input_image: Image.Image, blur_radius: int = 15, threshold: float = 0.5) -> Image.Image:
+def segment_and_blur_background(input_image: Image.Image, blur_strength: int = 15, threshold: float = 0.5) -> Image.Image:
     """
-    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.
+    Applies segmentation using the RMBG-2.0 model and composites the original image with
+    a Gaussian-blurred background based on an adjustable mask sensitivity threshold.
     """
-    # Ensure the image is in RGB and get its original dimensions
     image = input_image.convert("RGB")
     orig_width, orig_height = image.size
 
-    # Preprocess image for segmentation
+    # Preprocess image for segmentation (resize only for model inference)
     input_tensor = segmentation_transform(image).unsqueeze(0).to(device)
     
-    # Run inference on the segmentation model
     with torch.no_grad():
         preds = segmentation_model(input_tensor)[-1].sigmoid().cpu()
     pred = preds[0].squeeze()
     
-    # Create a binary mask using the adjustable threshold
+    # Create binary mask with adjustable threshold (mask sensitivity)
     binary_mask = (pred > threshold).float()
     mask_pil = transforms.ToPILImage()(binary_mask).convert("L")
-    # Convert grayscale mask to pure binary (0 or 255)
     mask_pil = mask_pil.point(lambda p: 255 if p > 128 else 0)
-    # Resize mask back to the original image dimensions
     mask_pil = mask_pil.resize((orig_width, orig_height), resample=Image.BILINEAR)
     
-    # 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 background using the mask
+    blurred_image = image.filter(ImageFilter.GaussianBlur(blur_strength))
     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 from Depth-Anything.
-    The max_blur parameter (controlled by a slider) sets the highest blur intensity.
+    Applies a depth-based blur effect using a depth map produced by Depth-Anything.
+    The effect simulates a lens blur where the max_blur parameter controls the maximum blur.
+    This function uses the original input image size.
     """
-    # Resize the input image to 512x512 for the depth estimation model
-    image_resized = input_image.resize((512, 512))
+    # Use the original image for depth estimation (no resizing)
+    image_original = input_image.convert("RGB")
     
-    # Run depth estimation to obtain the depth map (as a PIL image)
-    results = depth_pipeline(image_resized)
+    # Obtain depth map using the pipeline (assumes model accepts variable sizes)
+    results = depth_pipeline(image_original)
     depth_map_image = results['depth']
     
-    # Convert the depth map to a NumPy array and normalize to [0, 1]
     depth_array = np.array(depth_map_image, dtype=np.float32)
     d_min, d_max = depth_array.min(), depth_array.max()
     depth_norm = (depth_array - d_min) / (d_max - d_min + 1e-8)
     if invert_depth:
         depth_norm = 1.0 - depth_norm
 
-    # Convert the resized image to RGBA for compositing
-    orig_rgba = image_resized.convert("RGBA")
+    orig_rgba = image_original.convert("RGBA")
     final_image = orig_rgba.copy()
     
-    # 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]
         band_max = band_edges[i + 1]
-        # Use the midpoint of the band to determine the blur strength.
         mid = (band_min + band_max) / 2.0
         blur_radius_band = (1 - mid) * max_blur
         
-        # Create a blurred version of the image for this band.
         blurred_version = orig_rgba.filter(ImageFilter.GaussianBlur(blur_radius_band))
-        
-        # Create a mask for pixels whose normalized depth falls within this band.
         band_mask = ((depth_norm >= band_min) & (depth_norm < band_max)).astype(np.uint8) * 255
         band_mask_pil = Image.fromarray(band_mask, mode="L")
-        
-        # Composite the blurred version with the current final image using the band mask.
         final_image = Image.composite(blurred_version, final_image, band_mask_pil)
     
-    # Return the final composited image as RGB.
     return final_image.convert("RGB")
 
-
-def process_image(input_image: Image.Image, effect: str, threshold: float, blur_intensity: float) -> Image.Image:
+def process_image(input_image: Image.Image, effect: str, mask_sensitivity: float, blur_strength: float) -> Image.Image:
     """
-    Dispatch function to apply the selected effect:
-      - "Gaussian Blur Background": uses segmentation with an adjustable threshold and blur radius.
+    Applies the selected effect:
+      - "Gaussian Blur Background": uses segmentation with adjustable mask sensitivity and blur strength.
       - "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":
-        # 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)
+        return segment_and_blur_background(input_image, blur_strength=int(blur_strength), threshold=mask_sensitivity)
     elif effect == "Depth-based Lens Blur":
-        # For depth-based blur, use the blur_intensity as the max blur value.
-        return depth_based_lens_blur(input_image, max_blur=blur_intensity)
+        return depth_based_lens_blur(input_image, max_blur=blur_strength)
     else:
         return input_image
 
-
 # ----------------------------
-# Gradio Interface
+# Gradio Blocks Layout
 # ----------------------------
 
-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.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="Interactive Blur Effects Demo",
-    description=(
-        "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."
+with gr.Blocks(title="Interactive Blur Effects Demo") as demo:
+    gr.Markdown(
+        """
+        # Interactive Blur Effects Demo
+        Upload an image and choose an effect below.  
+        For **Gaussian Blur Background**, adjust the mask sensitivity (controls segmentation threshold) 
+        and blur strength (controls Gaussian blur radius).  
+        For **Depth-based Lens Blur**, the blur strength slider sets the maximum blur intensity.
+        """
+    )
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Input Image")
+            effect_choice = gr.Radio(
+                choices=["Gaussian Blur Background", "Depth-based Lens Blur"],
+                label="Select Effect",
+                value="Gaussian Blur Background"
+            )
+            mask_sensitivity_slider = gr.Slider(
+                minimum=0.0, maximum=1.0, value=0.5, step=0.01,
+                label="Mask Sensitivity (for segmentation)"
+            )
+            blur_strength_slider = gr.Slider(
+                minimum=0, maximum=30, value=15, step=1,
+                label="Blur Strength"
+            )
+            run_button = gr.Button("Apply Effect")
+        with gr.Column():
+            output_image = gr.Image(type="pil", label="Output Image")
+    
+    run_button.click(
+        fn=process_image,
+        inputs=[input_image, effect_choice, mask_sensitivity_slider, blur_strength_slider],
+        outputs=output_image
     )
-)
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()