Update app.py

app.py

@@ -13,7 +13,6 @@ from transformers import AutoModelForImageSegmentation
 import cv2
 import ezdxf
 import gradio as gr
-from MODNet.src.models.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,21 +41,6 @@ PAPER_SIZES = {
     "A3": {"width": 297, "height": 420},
     "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
-    return yolo_world
     
 # Custom Exception Classes
 class TimeoutReachedError(Exception):
@@ -536,34 +520,6 @@ 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"""
@@ -1103,143 +1059,6 @@ def make_square(img: np.ndarray):
     
     return padded
 
-# def predict_with_paper(image, paper_size, offset,offset_unit, finger_clearance=False):
-#     """Main prediction function using paper as reference"""
-    
-#     logger.info(f"Starting prediction with image shape: {image.shape}")
-#     logger.info(f"Paper size: {paper_size}")
-    
-#     if offset_unit == "inches":
-#         offset_mm = offset * 25.4  # Convert to mm for internal calculations
-#     else:
-#         offset_mm = offset
-        
-#     edge_radius = None
-#     if edge_radius is None or edge_radius == 0:
-#         edge_radius = 0.0001
-
-#     if offset < 0:
-#         raise gr.Error("Offset Value Can't be negative")
-
-#     try:
-#         # Detect paper bounds and calculate scaling factor
-#         logger.info("Starting paper detection...")
-#         paper_contour, scaling_factor = detect_paper_bounds(image, paper_size)
-#         logger.info(f"Paper detected successfully with scaling factor: {scaling_factor:.4f} mm/px")
-        
-#     except ReferenceBoxNotDetectedError as e:
-#         logger.error(f"Paper detection failed: {e}")
-#         return (
-#             None, None, None, None,
-#             f"Error: {str(e)}"
-#         )
-#     except Exception as e:
-#         logger.error(f"Unexpected error in paper detection: {e}")
-#         raise gr.Error(f"Error processing image: {str(e)}")
-
-#     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
-#         orig_size = image.shape[:2]
-#         objects_mask = remove_bg(image)
-#         processed_size = objects_mask.shape[:2]
-        
-#         # Resize mask to match original image
-#         objects_mask = cv2.resize(objects_mask, (image.shape[1], image.shape[0]))
-        
-#         # Remove paper area from mask to focus only on objects
-#         objects_mask = exclude_paper_area(objects_mask, paper_contour)
-        
-#         # Check if we actually have object pixels after paper exclusion
-#         object_pixels = np.count_nonzero(objects_mask)
-#         if object_pixels < 1000:  # Minimum threshold
-#             raise NoObjectDetectedError("No significant object detected after excluding paper area")
-        
-#         # Validate single object
-#         validate_single_object(objects_mask, paper_contour)
-        
-#     except (MultipleObjectsError, NoObjectDetectedError) as e:
-#         return (
-#             None, None, None, None,
-#             f"Error: {str(e)}"
-#         )
-#     except Exception as e:
-#         raise gr.Error(f"Error in object detection: {str(e)}")
-
-#     # Apply edge rounding if specified
-#     rounded_mask = objects_mask.copy()
-    
-#     # Apply dilation for offset - more precise calculation
-#     if offset_mm > 0:
-#         offset_pixels = int(round(float(offset_mm) / scaling_factor))
-#         if offset_pixels > 0:
-#             kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (offset_pixels*2+1, offset_pixels*2+1))
-#             dilated_mask = cv2.dilate(rounded_mask, kernel, iterations=1)
-#         else:
-#             dilated_mask = rounded_mask.copy()
-
-#     # Save original dilated mask for output
-#     Image.fromarray(dilated_mask).save("./outputs/scaled_mask_original.jpg")
-#     dilated_mask_orig = dilated_mask.copy()
-
-#     # Extract contours
-#     outlines, contours = extract_outlines(dilated_mask)
-
-#     try:
-#         # Generate DXF
-#         dxf, finger_polygons, original_polygons = save_dxf_spline(
-#             contours,
-#             scaling_factor,
-#             processed_size[0],
-#             finger_clearance=(finger_clearance == "On")
-#         )
-#     except FingerCutOverlapError as e:
-#         raise gr.Error(str(e))
-
-#     # Create annotated image
-#     shrunked_img_contours = image.copy()
-
-#     if finger_clearance == "On":
-#         outlines = np.full_like(dilated_mask, 255)
-#         for poly in finger_polygons:
-#             try:
-#                 coords = np.array([
-#                     (int(x / scaling_factor), int(processed_size[0] - y / scaling_factor))
-#                     for x, y in poly.exterior.coords
-#                 ], np.int32).reshape((-1, 1, 2))
-
-#                 cv2.drawContours(shrunked_img_contours, [coords], -1, (0, 255, 0), thickness=2)
-#                 cv2.drawContours(outlines, [coords], -1, 0, thickness=2)
-#             except Exception as e:
-#                 logger.warning(f"Failed to draw finger cut: {e}")
-#                 continue
-#     else:
-#         outlines = np.full_like(dilated_mask, 255)
-#         cv2.drawContours(shrunked_img_contours, contours, -1, (0, 255, 0), thickness=2)
-#         cv2.drawContours(outlines, contours, -1, 0, thickness=2)
-    
-#     # # Draw paper bounds on annotated image
-#     # cv2.drawContours(shrunked_img_contours, [paper_contour], -1, (255, 0, 0), thickness=3)
-    
-#     # # Add paper size text
-#     # paper_text = f"Paper: {paper_size}"
-#     # cv2.putText(shrunked_img_contours, paper_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
-    
-#     cleanup_models()
-
-#     return (
-#         shrunked_img_contours,
-#         outlines,
-#         dxf,
-#         dilated_mask_orig,
-#         f"Scale: {scaling_factor:.4f} mm/px | Paper: {paper_size}"
-#     )
 def predict_with_paper(image, paper_size, offset, offset_unit, finger_clearance=False):
     """Main prediction function using paper as reference"""
     
@@ -1286,7 +1105,7 @@ def predict_with_paper(image, paper_size, offset, offset_unit, finger_clearance=
         orig_size = image.shape[:2]
         # objects_mask = remove_bg(image)
         # objects_mask = remove_bg(image)
-        objects_mask = remove_bg_modnet(image)
+        objects_mask = remove_bg_u2netp(image)
         processed_size = objects_mask.shape[:2]
         
         # Resize mask to match original image