Update app.py

app.py

@@ -42,6 +42,11 @@ PAPER_SIZES = {
     "US Letter": {"width": 215.9, "height": 279.4}
 }
 
+def get_yolo_world():
+    """Lazy load YOLO-World model"""
+    yolo_world = YOLOWorld('yolov8s-world.pt')  # Smaller model for efficiency
+    return yolo_world
+    
 # Custom Exception Classes
 class TimeoutReachedError(Exception):
     pass
@@ -520,7 +525,7 @@ def remove_bg_u2netp(image: np.ndarray) -> np.ndarray:
         logger.error(f"Error in U2NETP background removal: {e}")
         raise
 
-def remove_bg(image: np.ndarray) -> np.ndarray:
+def remove_bg_original(image: np.ndarray) -> np.ndarray:
     """Remove background using BiRefNet model for main objects"""
     try:
         birefnet_model = get_birefnet()
@@ -547,6 +552,36 @@ def remove_bg(image: np.ndarray) -> np.ndarray:
     except Exception as e:
         logger.error(f"Error in BiRefNet background removal: {e}")
         raise
+def remove_bg(image: np.ndarray, paper_contour: np.ndarray) -> np.ndarray:
+    """Remove background using either BiRefNet or YOLO-World based on object size"""
+    # Calculate paper area percentage
+    paper_area = cv2.contourArea(paper_contour)
+    image_area = image.shape[0] * image.shape[1]
+    paper_ratio = paper_area / image_area
+    
+    # If paper takes up most of the image (small object case)
+    if paper_ratio > 0.7:  
+        logger.info("Using YOLO-World for small object detection")
+        try:
+            model = get_yolo_world()
+            results = model.predict(image, conf=0.3, verbose=False)
+            
+            # Create blank mask
+            mask = np.zeros(image.shape[:2], dtype=np.uint8)
+            
+            # Draw all detected objects
+            for box in results[0].boxes.xyxy.cpu().numpy():
+                x1, y1, x2, y2 = map(int, box)
+                mask[y1:y2, x1:x2] = 255
+                
+            return mask
+            
+        except Exception as e:
+            logger.warning(f"YOLO-World failed, falling back to BiRefNet: {e}")
+    
+    # Default case - use BiRefNet/U2NET
+    logger.info("Using BiRefNet for standard object detection")
+    return remove_bg_original(image)  # Your existing BiRefNet implementation
 
 # def exclude_paper_area(mask: np.ndarray, paper_contour: np.ndarray, expansion_factor: float = 1.2) -> np.ndarray:
 #     """
@@ -1239,7 +1274,8 @@ def predict_with_paper(image, paper_size, offset, offset_unit, finger_clearance=
         
         # Remove background from main objects
         orig_size = image.shape[:2]
-        objects_mask = remove_bg(image)
+        # objects_mask = remove_bg(image)
+        objects_mask = remove_bg(image, paper_contour)
         processed_size = objects_mask.shape[:2]
         
         # Resize mask to match original image