Update app.py

app.py

@@ -43,11 +43,6 @@ depth_model.eval()
 # Define the Segmentation-Based Blur Effect
 # -----------------------------
 def segmentation_blur_effect(input_image: Image.Image):
-    """
-    Creates a segmentation mask using RMBG-2.0 and applies a Gaussian blur (sigma=15)
-    to the background while keeping the foreground sharp.
-    """
-    # Resize input image for segmentation processing
     imageResized = input_image.resize(seg_image_size)
     input_tensor = seg_transform(imageResized).unsqueeze(0).to(device)
     
@@ -55,54 +50,28 @@ def segmentation_blur_effect(input_image: Image.Image):
         preds = seg_model(input_tensor)[-1].sigmoid().cpu()
     pred = preds[0].squeeze()
     
-    # Convert predicted mask to a PIL image and ensure it matches imageResized's size
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize(imageResized.size)
     
-    # Convert mask to grayscale and threshold to create a binary mask
     mask_np = np.array(mask.convert("L"))
     _, maskBinary = cv2.threshold(mask_np, 127, 255, cv2.THRESH_BINARY)
     
-    # Convert the resized image to an OpenCV BGR array
     img = cv2.cvtColor(np.array(imageResized), cv2.COLOR_RGB2BGR)
-    # Apply Gaussian blur (sigmaX=15, sigmaY=15)
     blurredBg = cv2.GaussianBlur(np.array(imageResized), (0, 0), sigmaX=15, sigmaY=15)
     
-    # Create the inverse mask and convert it to 3 channels
     maskInv = cv2.bitwise_not(maskBinary)
     maskInv3 = cv2.cvtColor(maskInv, cv2.COLOR_GRAY2BGR)
     
-    # Extract the foreground and background using the mask
     foreground = cv2.bitwise_and(img, cv2.bitwise_not(maskInv3))
     background = cv2.bitwise_and(blurredBg, maskInv3)
     
-    # Combine foreground and background; convert back to RGB for display
     finalImg = cv2.add(cv2.cvtColor(foreground, cv2.COLOR_BGR2RGB), background)
     finalImg_pil = Image.fromarray(finalImg)
     
     return finalImg_pil, mask
 
-# -----------------------------
-# Define the Depth-Based Lens Blur Effect with Slider-Controlled Thresholds
-# -----------------------------
 def lens_blur_effect(input_image: Image.Image, fg_threshold: float = 85, mg_threshold: float = 170):
-    """
-    Uses DepthPro to estimate a depth map and applies a dynamic lens blur effect
-    by blending three versions of the image with increasing blur levels.
-    
-    Parameters:
-      input_image: The original PIL image.
-      fg_threshold: Foreground threshold (0-255). Pixels with depth below this are considered foreground.
-      mg_threshold: Middleground threshold (0-255). Pixels with depth between fg_threshold and mg_threshold are middleground.
-    
-    Returns:
-      depthImg: The computed depth map (PIL Image).
-      lensBlurImage: The final lens-blurred image (PIL Image).
-      mask_fg_img: Foreground depth mask.
-      mask_mg_img: Middleground depth mask.
-      mask_bg_img: Background depth mask.
-    """
-    # Process the image with the depth estimation model
+
     inputs = depth_processor(images=input_image, return_tensors="pt").to(device)
     with torch.no_grad():
         outputs = depth_model(**inputs)
@@ -111,39 +80,32 @@ def lens_blur_effect(input_image: Image.Image, fg_threshold: float = 85, mg_thre
     )
     depth = post_processed_output[0]["predicted_depth"]
     
-    # Normalize depth to [0, 255]
     depth = (depth - depth.min()) / (depth.max() - depth.min())
     depth = depth * 255.
     depth = depth.detach().cpu().numpy()
     depth_map = depth.astype(np.uint8)
     depthImg = Image.fromarray(depth_map)
     
-    # Convert input image to OpenCV BGR format
     img = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
     
-    # Precompute blurred versions for different depth regions
     img_foreground = img.copy()  # No blur for foreground
     img_middleground = cv2.GaussianBlur(img, (0, 0), sigmaX=7, sigmaY=7)
     img_background = cv2.GaussianBlur(img, (0, 0), sigmaX=15, sigmaY=15)
 
     print(depth_map)
-    depth_map /= depth_map.max()
+    depth_map = depth_map.astype(np.float32) / depth_map.max()
     
-    # Use slider values as thresholds
-    threshold1 = fg_threshold   # e.g., default 85
-    threshold2 = mg_threshold   # e.g., default 170
+    threshold1 = fg_threshold
+    threshold2 = mg_threshold
     
-    # Create masks for foreground, middleground, and background based on depth
     mask_fg = (depth_map < threshold1).astype(np.float32)
     mask_mg = ((depth_map >= threshold1) & (depth_map < threshold2)).astype(np.float32)
     mask_bg = (depth_map >= threshold2).astype(np.float32)
     
-    # Expand masks to 3 channels
     mask_fg_3 = np.stack([mask_fg]*3, axis=-1)
     mask_mg_3 = np.stack([mask_mg]*3, axis=-1)
     mask_bg_3 = np.stack([mask_bg]*3, axis=-1)
     
-    # Blend the images using the masks
     final_img = (img_foreground * mask_fg_3 +
                  img_middleground * mask_mg_3 +
                  img_background * mask_bg_3).astype(np.uint8)
@@ -151,27 +113,14 @@ def lens_blur_effect(input_image: Image.Image, fg_threshold: float = 85, mg_thre
     final_img_rgb = cv2.cvtColor(final_img, cv2.COLOR_BGR2RGB)
     lensBlurImage = Image.fromarray(final_img_rgb)
     
-    # Create mask images for display (scaled to 0-255)
     mask_fg_img = Image.fromarray((mask_fg * 255).astype(np.uint8))
     mask_mg_img = Image.fromarray((mask_mg * 255).astype(np.uint8))
     mask_bg_img = Image.fromarray((mask_bg * 255).astype(np.uint8))
     
     return depthImg, lensBlurImage, mask_fg_img, mask_mg_img, mask_bg_img
 
-# -----------------------------
-# Gradio App: Process Image and Display Multiple Effects
-# -----------------------------
 def process_image(input_image: Image.Image, fg_threshold: float, mg_threshold: float):
-    """
-    Processes the uploaded image to generate:
-      1. Segmentation-based Gaussian blur effect.
-      2. Segmentation mask.
-      3. Depth map.
-      4. Depth-based lens blur effect.
-      5. Depth masks for foreground, middleground, and background.
-    
-    The depth thresholds for foreground and middleground regions are adjustable via sliders.
-    """
+
     seg_blur, seg_mask = segmentation_blur_effect(input_image)
     depth_map_img, lens_blur_img, mask_fg_img, mask_mg_img, mask_bg_img = lens_blur_effect(
         input_image, fg_threshold, mg_threshold