Resize changes

app.py

@@ -19,7 +19,6 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 seg_model = AutoModelForImageSegmentation.from_pretrained(
     "briaai/RMBG-2.0", trust_remote_code=True
 )
-# Set higher precision for matmul if desired
 torch.set_float32_matmul_precision(["high", "highest"][0])
 seg_model.to(device)
 seg_model.eval()
@@ -47,13 +46,8 @@ 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.
-    
-    Returns:
-      - final segmented and blurred image (PIL Image)
-      - segmentation mask (PIL Image)
-      - blurred background image (PIL Image) [optional display]
     """
-    # Resize input for segmentation processing
+    # Resize input image for segmentation processing
     imageResized = input_image.resize(seg_image_size)
     input_tensor = seg_transform(imageResized).unsqueeze(0).to(device)
     
@@ -61,11 +55,12 @@ 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 resize to original input size
+    # Convert predicted mask to a PIL image and ensure it matches imageResized's size
     pred_pil = transforms.ToPILImage()(pred)
-    mask = pred_pil.resize(input_image.size)
+    # Resize mask to match imageResized to avoid size mismatch in OpenCV operations
+    mask = pred_pil.resize(imageResized.size)
     
-    # Create a binary mask (convert to grayscale, then threshold)
+    # 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)
     
@@ -74,20 +69,19 @@ def segmentation_blur_effect(input_image: Image.Image):
     # 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 to 3 channels
+    # 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 separately
+    # Extract the foreground and background using the mask
     foreground = cv2.bitwise_and(img, cv2.bitwise_not(maskInv3))
     background = cv2.bitwise_and(blurredBg, maskInv3)
     
-    # Combine the two components
+    # 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)
-    blurredBg_pil = Image.fromarray(cv2.cvtColor(blurredBg, cv2.COLOR_BGR2RGB))
     
-    return finalImg_pil, mask, blurredBg_pil
+    return finalImg_pil, mask
 
 # -----------------------------
 # Define the Depth-Based Lens Blur Effect
@@ -95,15 +89,9 @@ def segmentation_blur_effect(input_image: Image.Image):
 def lens_blur_effect(input_image: Image.Image):
     """
     Uses DepthPro to estimate a depth map and applies a dynamic lens blur effect
-    by precomputing three versions of the image (foreground, middleground, background)
-    with increasing blur. Regions are blended based on the estimated depth.
-    
-    Returns:
-      - Depth map (PIL Image)
-      - Final lens-blurred image (PIL Image)
-      - Foreground mask (PIL Image)
-      - Middleground mask (PIL Image)
-      - Background mask (PIL Image)
+    by blending three versions of the image (foreground, middleground, background)
+    with increasing blur levels. Returns the depth map, the final lens-blurred image,
+    and the depth masks.
     """
     # Process the image with the depth estimation model
     inputs = depth_processor(images=input_image, return_tensors="pt").to(device)
@@ -124,7 +112,7 @@ def lens_blur_effect(input_image: Image.Image):
     # Convert input image to OpenCV BGR format
     img = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
     
-    # Precompute three blurred versions of the image
+    # 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)
@@ -133,17 +121,17 @@ def lens_blur_effect(input_image: Image.Image):
     threshold1 = 255 / 3      # ~85
     threshold2 = 2 * 255 / 3  # ~170
     
-    # Create masks for the three regions based on depth
+    # 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 to match image dimensions
+    # 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)
     
-    # Combine the images using the masks (vectorized blending)
+    # Blend the images using the masks (vectorized operation)
     final_img = (img_foreground * mask_fg_3 +
                  img_middleground * mask_mg_3 +
                  img_background * mask_bg_3).astype(np.uint8)
@@ -151,7 +139,7 @@ def lens_blur_effect(input_image: Image.Image):
     final_img_rgb = cv2.cvtColor(final_img, cv2.COLOR_BGR2RGB)
     lensBlurImage = Image.fromarray(final_img_rgb)
     
-    # Create mask images (scaled to 0-255)
+    # 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))
@@ -170,7 +158,7 @@ def process_image(input_image: Image.Image):
       4. Depth-based lens blur effect.
       5. Depth-based masks for foreground, middleground, and background.
     """
-    seg_blur, seg_mask, _ = segmentation_blur_effect(input_image)
+    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)
     
     return (
@@ -188,7 +176,7 @@ description = (
     "Upload an image to apply two distinct effects:\n\n"
     "1. A segmentation-based Gaussian blur that blurs the background (using RMBG-2.0).\n"
     "2. A depth-based lens blur effect that simulates realistic lens blur based on depth (using DepthPro).\n\n"
-    "Outputs include the blurred image, segmentation mask, depth map, lens-blurred image, and depth masks."
+    "Outputs include the blurred image, segmentation mask, depth map, lens-blurred image, and individual depth masks."
 )
 
 demo = gr.Interface(