Update app.py

app.py

@@ -368,6 +368,20 @@ def detect_paper_bounds(image: np.ndarray, paper_size: str, output_unit: str = "
             # Use fallback contour detection
             logger.info("Using fallback contour detection for paper")
             paper_contour, _ = detect_paper_contour(image, output_unit)
+
+        # After getting paper_contour, expand it
+        rect = cv2.boundingRect(paper_contour)
+        expansion = int(min(rect[2], rect[3]) * 0.1)  # Expand by 10%
+        
+        x, y, w, h = rect
+        expanded_contour = np.array([
+            [[max(0, x - expansion), max(0, y - expansion)]],
+            [[min(image.shape[1], x + w + expansion), max(0, y - expansion)]],
+            [[min(image.shape[1], x + w + expansion), min(image.shape[0], y + h + expansion)]],
+            [[max(0, x - expansion), min(image.shape[0], y + h + expansion)]]
+        ])
+        
+        paper_contour = expanded_contour
         
         # Calculate scaling factor based on paper size with proper units
         scaling_factor = calculate_paper_scaling_factor(paper_contour, paper_size, output_unit)
@@ -529,24 +543,20 @@ def mask_paper_area_in_image(image: np.ndarray, paper_contour: np.ndarray) -> np
     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 using soft boundaries
-    """
+    """Less aggressive paper area exclusion"""
     # Create paper mask
     paper_mask = np.zeros(mask.shape[:2], dtype=np.uint8)
     cv2.fillPoly(paper_mask, [paper_contour], 255)
     
-    # Create eroded version of paper mask (soft inward boundary)
-    erosion_size = int(min(mask.shape[0], mask.shape[1]) * 0.03)  # 3% of image size
-    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (erosion_size, erosion_size))
-    eroded_paper_mask = cv2.erode(paper_mask, kernel, iterations=2)
+    # Instead of eroding, slightly expand the paper mask
+    rect = cv2.boundingRect(paper_contour)
+    expansion = max(10, int(min(rect[2], rect[3]) * 0.02))  # 2% expansion
     
-    # Apply soft mask: keep pixels in original mask AND in eroded paper area
-    result_mask = cv2.bitwise_and(mask, eroded_paper_mask)
+    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (expansion, expansion))
+    expanded_paper_mask = cv2.dilate(paper_mask, kernel, iterations=1)
     
-    # Gentle cleanup without hard cuts
-    small_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
-    result_mask = cv2.morphologyEx(result_mask, cv2.MORPH_CLOSE, small_kernel, iterations=1)
+    # Keep objects within expanded paper area
+    result_mask = cv2.bitwise_and(mask, expanded_paper_mask)
     
     return result_mask
 
@@ -964,51 +974,31 @@ def predict_with_paper(image, paper_size, offset, offset_unit, finger_clearance=
         raise gr.Error(f"Error processing image: {str(e)}")
 
     try:
-        # Mask paper area in input image first
-        # masked_input_image = mask_paper_area_in_image(image, paper_contour)
+        # Get paper bounds with expansion
+        rect = cv2.boundingRect(paper_contour)
+        expansion = max(20, int(min(rect[2], rect[3]) * 0.05))  # 5% expansion
         
-        # Use YOLOv8 to detect objects
-        yolo_v8 = get_yolo_v8()
-        if yolo_v8 is None:
-            logger.warning("YOLOv8 model not available, proceeding with full image")
-            cropped_image = image
-            crop_offset = (0, 0)
-        else:
-            # YOLOv8 detects all COCO classes by default
-            results = yolo_v8.predict(image, conf=0.1, verbose=False)
-            
-            if not results or len(results) == 0 or not hasattr(results[0], 'boxes') or len(results[0].boxes) == 0:
-                logger.warning("No objects detected by YOLOv8, proceeding with full image")
-                cropped_image = image
-                crop_offset = (0, 0)
-            else:
-                boxes = results[0].boxes.xyxy.cpu().numpy()
-                confidences = results[0].boxes.conf.cpu().numpy()
-                
-                # Filter out very large boxes (likely paper/background)
-                image_area = image.shape[0] * image.shape[1]
-                valid_boxes = []
-                
-                for i, box in enumerate(boxes):
-                    x_min, y_min, x_max, y_max = box
-                    box_area = (x_max - x_min) * (y_max - y_min)
-                    # Keep boxes that are 5% to 40% of image area
-                    if 0.001 * image_area < box_area < 0.6 * image_area:
-                        valid_boxes.append((i, confidences[i]))
-                
-                if not valid_boxes:
-                    logger.warning("No valid objects detected, proceeding with full image")
-                    cropped_image = image
-                    crop_offset = (0, 0)
-                else:
-                    # Get highest confidence valid box
-                    best_idx = max(valid_boxes, key=lambda x: x[1])[0]
-                    x_min, y_min, x_max, y_max = map(int, boxes[best_idx])
+        x, y, w, h = rect
+        x_min = max(0, x - expansion)
+        y_min = max(0, y - expansion)
+        x_max = min(image.shape[1], x + w + expansion)
+        y_max = min(image.shape[0], y + h + expansion)
+        
+        # Process the expanded paper area
+        cropped_image = image[y_min:y_max, x_min:x_max]
+        crop_offset = (x_min, y_min)
         
-        # Remove background from cropped image
-        orig_size = image.shape[:2]
+        # Remove background
         objects_mask = remove_bg(cropped_image)
-        processed_size = objects_mask.shape[:2]
+        
+        # Place back in full image space
+        full_mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
+        full_mask[y_min:y_max, x_min:x_max] = objects_mask
+        
+        # Light filtering only - don't exclude paper area aggressively
+        # Just remove small noise
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
+        objects_mask = cv2.morphologyEx(full_mask, cv2.MORPH_OPEN, kernel)
         
         # Resize mask to match cropped region and place back in original image space
         full_mask = np.zeros((orig_size[0], orig_size[1]), dtype=np.uint8)