Update app.py

app.py

@@ -5,54 +5,32 @@ from transformers import AutoImageProcessor, AutoModelForDepthEstimation
 from PIL import Image, ImageFilter
 import matplotlib.pyplot as plt
 import matplotlib.cm as cm
-
-# ---------------------------
-# Depth Estimation Utilities
-# ---------------------------
 def compute_depth_map(image: Image.Image, scale_factor: float) -> np.ndarray:
-    """
-    Loads the LiheYoung/depth-anything-large-hf model and computes a depth map.
-    The depth map is normalized, inverted (so that near=0 and far=1),
-    and multiplied by the given scale_factor.
-    """
-    # Load model and processor from pretrained weights
     image_processor = AutoImageProcessor.from_pretrained("LiheYoung/depth-anything-large-hf")
     model = AutoModelForDepthEstimation.from_pretrained("LiheYoung/depth-anything-large-hf")
-    
-    # Prepare image for the model
     inputs = image_processor(images=image, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
         predicted_depth = outputs.predicted_depth
 
-    # Interpolate predicted depth map to match image size
     prediction = torch.nn.functional.interpolate(
         predicted_depth.unsqueeze(1),
-        size=image.size[::-1],  # PIL image size is (width, height)
+        size=image.size[::-1],  
         mode="bicubic",
         align_corners=False,
     )
-    # Normalize for visualization
+   
     depth_min = prediction.min()
     depth_max = prediction.max()
     depth_vis = (prediction - depth_min) / (depth_max - depth_min + 1e-8)
     depth_map = depth_vis.squeeze().cpu().numpy()
-    # Invert so that near=0 and far=1, then scale
     depth_map_inverted = 1.0 - depth_map
     depth_map_inverted *= scale_factor
     return depth_map_inverted
 
-# ---------------------------
-# Depth-Based Blur Functions
-# ---------------------------
 def layered_blur(image: Image.Image, depth_map: np.ndarray, num_layers: int, max_blur: float) -> Image.Image:
-    """
-    Creates multiple blurred versions of the image (using Gaussian blur with radii from 0 to max_blur)
-    and composites them using masks generated from bins of the normalized depth map.
-    """
     blur_radii = np.linspace(0, max_blur, num_layers)
     blur_versions = [image.filter(ImageFilter.GaussianBlur(radius)) for radius in blur_radii]
-    # Use a fixed range (0 to 1) since the depth map is normalized
     thresholds = np.linspace(0, 1, num_layers + 1)
     final_image = blur_versions[-1]
     for i in range(num_layers - 1, -1, -1):
@@ -65,51 +43,27 @@ def layered_blur(image: Image.Image, depth_map: np.ndarray, num_layers: int, max
     return final_image
 
 def process_depth_blur(uploaded_image, max_blur_value, scale_factor, num_layers):
-    """
-    Resizes the uploaded image to 512x512, computes its depth map,
-    and applies layered blur based on the depth.
-    """
     if not isinstance(uploaded_image, Image.Image):
         uploaded_image = Image.open(uploaded_image)
     image = uploaded_image.convert("RGB").resize((512, 512))
     depth_map = compute_depth_map(image, scale_factor)
     final_image = layered_blur(image, depth_map, int(num_layers), max_blur_value)
     return final_image
-
-# ---------------------------
-# Depth Heatmap Functions
-# ---------------------------
 def create_heatmap(depth_map: np.ndarray, intensity: float) -> Image.Image:
-    """
-    Applies a colormap to the normalized depth map.
-    The 'intensity' slider multiplies the normalized depth values (clipped to [0,1])
-    before applying the "inferno" colormap.
-    """
-    # Multiply depth map by intensity and clip to 0-1 range
     normalized = np.clip(depth_map * intensity, 0, 1)
     colormap = cm.get_cmap("inferno")
-    colored = colormap(normalized)  # Returns an RGBA image in [0, 1]
-    heatmap = (colored[:, :, :3] * 255).astype(np.uint8)  # drop alpha and convert to [0,255]
+    colored = colormap(normalized) 
+    heatmap = (colored[:, :, :3] * 255).astype(np.uint8) 
     return Image.fromarray(heatmap)
 
 def process_depth_heatmap(uploaded_image, intensity):
-    """
-    Resizes the uploaded image to 512x512, computes its depth map (with scale factor 1.0),
-    and returns a heatmap visualization.
-    """
     if not isinstance(uploaded_image, Image.Image):
         uploaded_image = Image.open(uploaded_image)
     image = uploaded_image.convert("RGB").resize((512, 512))
     depth_map = compute_depth_map(image, scale_factor=1.0)
     heatmap_img = create_heatmap(depth_map, intensity)
     return heatmap_img
-
-# --- Segmentation-Based Blur using BEN2 ---
 def load_segmentation_model():
-    """
-    Loads and caches the segmentation model from BEN2.
-    Ensure you have ben2 installed and accessible in your path.
-    """
     global seg_model, seg_device
     if "seg_model" not in globals():
         from ben2 import BEN_Base  # Import BEN2
@@ -119,18 +73,11 @@ def load_segmentation_model():
     return seg_model, seg_device
 
 def process_segmentation_blur(uploaded_image, seg_blur_radius: float):
-    """
-    Processes the image with segmentation-based blur.
-    The image is resized to 512x512. A Gaussian blur with the specified radius is applied,
-    then the segmentation mask is computed to composite the sharp foreground over the blurred background.
-    """
     if not isinstance(uploaded_image, Image.Image):
         uploaded_image = Image.open(uploaded_image)
     image = uploaded_image.convert("RGB").resize((512, 512))
     seg_model, seg_device = load_segmentation_model()
     blurred_image = image.filter(ImageFilter.GaussianBlur(seg_blur_radius))
-    
-    # Generate segmentation mask (foreground)
     foreground = seg_model.inference(image, refine_foreground=False)
     foreground_rgba = foreground.convert("RGBA")
     _, _, _, alpha = foreground_rgba.split()
@@ -138,39 +85,29 @@ def process_segmentation_blur(uploaded_image, seg_blur_radius: float):
     final_image = Image.composite(image, blurred_image, binary_mask)
     return final_image
 
-# --- Merged Gradio Interface ---
+
 with gr.Blocks() as demo:
-    gr.Markdown("# Depth-Based vs Segmentation-Based Blur")
+    gr.Markdown("#Gaussian Blur & Lens Blur Effect")
     with gr.Tabs():
-        with gr.Tab("Depth Blur"):
+        with gr.Tab("Gaussian Blur"):
+            seg_img = gr.Image(type="pil", label="Upload Image")
+            seg_blur = gr.Slider(5, 30, value=15, step=1, label="Segmentation Blur Radius")
+            seg_out = gr.Image(label="Gaussian Blurred Image")
+            seg_button = gr.Button("Process Gaussian Blur")
+            seg_button.click(process_segmentation_blur, inputs=[seg_img, seg_blur], outputs=seg_out)
+        with gr.Tab("Lens Blur"):
             img_input = gr.Image(type="pil", label="Upload Image")
             blur_slider = gr.Slider(1, 50, value=6, label="Maximum Blur Radius")
             scale_slider = gr.Slider(0.1, 2.0, value=0.72, label="Depth Scale Factor")
             layers_slider = gr.Slider(2, 10, value=2.91, label="Number of Layers")
-            blur_output = gr.Image(label="Depth Blur Result")
-            blur_button = gr.Button("Process Depth Blur")
+            blur_output = gr.Image(label="Lens Blur Result")
+            blur_button = gr.Button("Process Blur")
             blur_button.click(
                 process_depth_blur, 
                 inputs=[img_input, blur_slider, scale_slider, layers_slider], 
                 outputs=blur_output
             )
-        with gr.Tab("Depth Heatmap"):
-            img_input2 = gr.Image(type="pil", label="Upload Image")
-            intensity_slider = gr.Slider(0.5, 5.0, value=1.0, label="Heatmap Intensity")
-            heatmap_output = gr.Image(label="Depth Heatmap")
-            heatmap_button = gr.Button("Generate Depth Heatmap")
-            heatmap_button.click(
-                process_depth_heatmap, 
-                inputs=[img_input2, intensity_slider], 
-                outputs=heatmap_output
-            )
-        with gr.Tab("Segmentation-Based Blur (BEN2)"):
-            seg_img = gr.Image(type="pil", label="Upload Image")
-            seg_blur = gr.Slider(5, 30, value=15, step=1, label="Segmentation Blur Radius")
-            seg_out = gr.Image(label="Segmentation-Based Blurred Image")
-            seg_button = gr.Button("Process Segmentation Blur")
-            seg_button.click(process_segmentation_blur, inputs=[seg_img, seg_blur], outputs=seg_out)
+       
 
 if __name__ == "__main__":
-    # Optionally, set share=True to generate a public link.
     demo.launch(share=True)