Update app.py

app.py

@@ -402,36 +402,92 @@ def remove_bg(image: np.ndarray) -> np.ndarray:
         logger.error(f"Error in BiRefNet background removal: {e}")
         raise
 
+# def exclude_paper_area(mask: np.ndarray, paper_contour: np.ndarray, expansion_factor: float = 1.2) -> np.ndarray:
+#     """
+#     Remove paper area from the mask to focus only on objects
+#     """
+#     # Create paper mask with slight expansion to ensure complete removal
+#     paper_mask = np.zeros(mask.shape[:2], dtype=np.uint8)
+    
+#     # # Expand paper contour slightly
+#     # epsilon = expansion_factor * cv2.arcLength(paper_contour, True)
+#     # expanded_contour = cv2.approxPolyDP(paper_contour, epsilon, True)
+    
+#     # cv2.fillPoly(paper_mask, [expanded_contour], 255)
+#     # Create a more aggressive inward shrinking of paper bounds
+#     rect = cv2.boundingRect(paper_contour)
+#     shrink_pixels = int(min(rect[2], rect[3]) * 0.05)  # Shrink by 5% of smaller dimension
+    
+#     # Create shrunken rectangle
+#     x, y, w, h = rect
+#     shrunken_contour = np.array([
+#         [[x + shrink_pixels, y + shrink_pixels]],
+#         [[x + w - shrink_pixels, y + shrink_pixels]],
+#         [[x + w - shrink_pixels, y + h - shrink_pixels]],
+#         [[x + shrink_pixels, y + h - shrink_pixels]]
+#     ])
+    
+#     cv2.fillPoly(paper_mask, [shrunken_contour], 255)
+    
+#     # Invert paper mask and apply to object mask
+#     paper_mask_inv = cv2.bitwise_not(paper_mask)
+#     result_mask = cv2.bitwise_and(mask, paper_mask_inv)
+    
+#     return result_mask
+def mask_paper_area_in_image(image: np.ndarray, paper_contour: np.ndarray) -> np.ndarray:
+    """
+    Black out paper area in the input image before sending to BiRefNet
+    """
+    masked_image = image.copy()
+    
+    # Create more aggressive paper mask
+    rect = cv2.boundingRect(paper_contour)
+    shrink_pixels = int(min(rect[2], rect[3]) * 0.08)  # 8% shrink
+    
+    x, y, w, h = rect
+    # Create mask for everything OUTSIDE the inner paper area
+    outer_mask = np.ones(image.shape[:2], dtype=np.uint8) * 255
+    
+    inner_contour = np.array([
+        [[x + shrink_pixels, y + shrink_pixels]],
+        [[x + w - shrink_pixels, y + shrink_pixels]],
+        [[x + w - shrink_pixels, y + h - shrink_pixels]],
+        [[x + shrink_pixels, y + h - shrink_pixels]]
+    ])
+    
+    # Black out everything outside inner paper bounds
+    cv2.fillPoly(outer_mask, [inner_contour], 0)
+    
+    # Apply mask to image
+    masked_image[outer_mask == 255] = [0, 0, 0]  # Black out paper areas
+    
+    return masked_image
+    
 def exclude_paper_area(mask: np.ndarray, paper_contour: np.ndarray, expansion_factor: float = 1.2) -> np.ndarray:
     """
     Remove paper area from the mask to focus only on objects
     """
-    # Create paper mask with slight expansion to ensure complete removal
+    # Create paper mask - this will be the area to EXCLUDE
     paper_mask = np.zeros(mask.shape[:2], dtype=np.uint8)
     
-    # # Expand paper contour slightly
-    # epsilon = expansion_factor * cv2.arcLength(paper_contour, True)
-    # expanded_contour = cv2.approxPolyDP(paper_contour, epsilon, True)
-    
-    # cv2.fillPoly(paper_mask, [expanded_contour], 255)
     # Create a more aggressive inward shrinking of paper bounds
     rect = cv2.boundingRect(paper_contour)
     shrink_pixels = int(min(rect[2], rect[3]) * 0.05)  # Shrink by 5% of smaller dimension
     
-    # Create shrunken rectangle
+    # Create shrunken rectangle (area INSIDE paper bounds)
     x, y, w, h = rect
-    shrunken_contour = np.array([
+    inner_contour = np.array([
         [[x + shrink_pixels, y + shrink_pixels]],
         [[x + w - shrink_pixels, y + shrink_pixels]],
         [[x + w - shrink_pixels, y + h - shrink_pixels]],
         [[x + shrink_pixels, y + h - shrink_pixels]]
     ])
     
-    cv2.fillPoly(paper_mask, [shrunken_contour], 255)
+    # Fill the INNER area as white (255) - this is where objects should be
+    cv2.fillPoly(paper_mask, [inner_contour], 255)
     
-    # Invert paper mask and apply to object mask
-    paper_mask_inv = cv2.bitwise_not(paper_mask)
-    result_mask = cv2.bitwise_and(mask, paper_mask_inv)
+    # Apply mask: keep only pixels that are both in original mask AND inside paper bounds
+    result_mask = cv2.bitwise_and(mask, paper_mask)
     
     return result_mask
 
@@ -831,8 +887,16 @@ def predict_with_paper(image, paper_size, offset,offset_unit, finger_clearance=F
 
     try:
         # Remove background from main objects
+        # orig_size = image.shape[:2]
+        # objects_mask = remove_bg(image)
+        
+        # Mask paper area in input image first
+        masked_input_image = mask_paper_area_in_image(image, paper_contour)
+        
+        # Remove background from main objects using masked image
         orig_size = image.shape[:2]
-        objects_mask = remove_bg(image)
+        objects_mask = remove_bg(masked_input_image)
+        
         processed_size = objects_mask.shape[:2]
         
         # Resize mask to match original image