mlbench123 commited on
Commit
86630a0
·
verified ·
1 Parent(s): af7ebad

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -32
app.py CHANGED
@@ -525,7 +525,7 @@ def remove_bg_u2netp(image: np.ndarray) -> np.ndarray:
525
  logger.error(f"Error in U2NETP background removal: {e}")
526
  raise
527
 
528
- def remove_bg_original(image: np.ndarray) -> np.ndarray:
529
  """Remove background using BiRefNet model for main objects"""
530
  try:
531
  birefnet_model = get_birefnet()
@@ -552,36 +552,6 @@ def remove_bg_original(image: np.ndarray) -> np.ndarray:
552
  except Exception as e:
553
  logger.error(f"Error in BiRefNet background removal: {e}")
554
  raise
555
- def remove_bg(image: np.ndarray, paper_contour: np.ndarray) -> np.ndarray:
556
- """Remove background using either BiRefNet or YOLO-World based on object size"""
557
- # Calculate paper area percentage
558
- paper_area = cv2.contourArea(paper_contour)
559
- image_area = image.shape[0] * image.shape[1]
560
- paper_ratio = paper_area / image_area
561
-
562
- # If paper takes up most of the image (small object case)
563
- if paper_ratio > 0.7:
564
- logger.info("Using YOLO-World for small object detection")
565
- try:
566
- model = get_yolo_world()
567
- results = model.predict(image, conf=0.3, verbose=False)
568
-
569
- # Create blank mask
570
- mask = np.zeros(image.shape[:2], dtype=np.uint8)
571
-
572
- # Draw all detected objects
573
- for box in results[0].boxes.xyxy.cpu().numpy():
574
- x1, y1, x2, y2 = map(int, box)
575
- mask[y1:y2, x1:x2] = 255
576
-
577
- return mask
578
-
579
- except Exception as e:
580
- logger.warning(f"YOLO-World failed, falling back to BiRefNet: {e}")
581
-
582
- # Default case - use BiRefNet/U2NET
583
- logger.info("Using BiRefNet for standard object detection")
584
- return remove_bg_original(image) # Your existing BiRefNet implementation
585
 
586
  # def exclude_paper_area(mask: np.ndarray, paper_contour: np.ndarray, expansion_factor: float = 1.2) -> np.ndarray:
587
  # """
@@ -1275,7 +1245,7 @@ def predict_with_paper(image, paper_size, offset, offset_unit, finger_clearance=
1275
  # Remove background from main objects
1276
  orig_size = image.shape[:2]
1277
  # objects_mask = remove_bg(image)
1278
- objects_mask = remove_bg(image, paper_contour)
1279
  processed_size = objects_mask.shape[:2]
1280
 
1281
  # Resize mask to match original image
 
525
  logger.error(f"Error in U2NETP background removal: {e}")
526
  raise
527
 
528
+ def remove_bg(image: np.ndarray) -> np.ndarray:
529
  """Remove background using BiRefNet model for main objects"""
530
  try:
531
  birefnet_model = get_birefnet()
 
552
  except Exception as e:
553
  logger.error(f"Error in BiRefNet background removal: {e}")
554
  raise
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
555
 
556
  # def exclude_paper_area(mask: np.ndarray, paper_contour: np.ndarray, expansion_factor: float = 1.2) -> np.ndarray:
557
  # """
 
1245
  # Remove background from main objects
1246
  orig_size = image.shape[:2]
1247
  # objects_mask = remove_bg(image)
1248
+ objects_mask = remove_bg(image)
1249
  processed_size = objects_mask.shape[:2]
1250
 
1251
  # Resize mask to match original image