Update app.py

app.py

@@ -13,6 +13,7 @@ from transformers import AutoModelForImageSegmentation
 import cv2
 import ezdxf
 import gradio as gr
+from modnet import MODNet
 import gc
 # from scalingtestupdated import calculate_scaling_factor
 from scalingtestupdated import calculate_scaling_factor_with_units, calculate_paper_scaling_factor, convert_units
@@ -42,6 +43,16 @@ PAPER_SIZES = {
     "US Letter": {"width": 215.9, "height": 279.4}
 }
 
+def get_modnet():
+    """Lazy load MODNet model"""
+    global modnet_global
+    if modnet_global is None:
+        logger.info("Loading MODNet model...")
+        modnet_global = MODNet(backbone_pretrained=False)
+        modnet_global.load_state_dict(torch.load('modnet_photographic_portrait_matting.ckpt', map_location='cpu'))
+        modnet_global.to(device).eval()
+    return modnet_global
+    
 def get_yolo_world():
     """Lazy load YOLO-World model"""
     yolo_world = YOLOWorld('yolov8s-world.pt')  # Smaller model for efficiency
@@ -525,6 +536,35 @@ def remove_bg_u2netp(image: np.ndarray) -> np.ndarray:
         logger.error(f"Error in U2NETP background removal: {e}")
         raise
 
+def remove_bg_modnet(image: np.ndarray) -> np.ndarray:
+    """Remove background using MODNet"""
+    try:
+        model = get_modnet()
+        
+        # Preprocess
+        transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+        ])
+        
+        image_pil = Image.fromarray(image)
+        image_tensor = transform(image_pil).unsqueeze(0).to(device)
+        
+        # Predict
+        with torch.no_grad():
+            _, _, matte = model(image_tensor, True)
+        
+        # Post-process
+        matte = matte.squeeze().cpu().numpy()
+        matte = (matte * 255).astype(np.uint8)
+        matte = cv2.resize(matte, (image.shape[1], image.shape[0]))
+        
+        return matte
+        
+    except Exception as e:
+        logger.error(f"MODNet failed: {e}")
+        return remove_bg_u2netp(image)  # Fallback to U2Net
+
 def remove_bg(image: np.ndarray) -> np.ndarray:
     """Remove background using BiRefNet model for main objects"""
     try:
@@ -1245,7 +1285,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)
+        objects_mask = remove_bg_modnet(image)
         processed_size = objects_mask.shape[:2]
         
         # Resize mask to match original image