Update app.py

app.py

@@ -70,6 +70,11 @@ class FingerCutOverlapError(Exception):
     def __init__(self, message="There was an overlap with fingercuts... Please try again to generate dxf."):
         super().__init__(message)
 
+class ReferenceBoxNotDetectedError(Exception):
+    """Raised when reference box/paper cannot be detected"""
+    def __init__(self, message="Reference box not detected"):
+        super().__init__(message)
+
 # Global model variables for lazy loading
 paper_detector_global = None
 u2net_global = None
@@ -106,12 +111,16 @@ def get_paper_detector():
     if paper_detector_global is None:
         logger.info("Loading paper detector model...")
         if os.path.exists(paper_model_path):
-            paper_detector_global = YOLO(paper_model_path)
+            try:
+                paper_detector_global = YOLO(paper_model_path)
+                logger.info("Paper detector loaded successfully")
+            except Exception as e:
+                logger.error(f"Failed to load paper detector: {e}")
+                paper_detector_global = None
         else:
             # Fallback to generic object detection for paper-like rectangles
-            logger.warning("Using fallback paper detection")
+            logger.warning("Paper model file not found, using fallback detection")
             paper_detector_global = None
-        logger.info("Paper detector loaded successfully")
     return paper_detector_global
 
 def get_u2net():
@@ -149,47 +158,70 @@ def detect_paper_contour(image: np.ndarray) -> Tuple[np.ndarray, float]:
     Detect paper in the image using contour detection as fallback
     Returns the paper contour and estimated scaling factor
     """
+    logger.info("Using contour-based paper detection")
+    
     # Convert to grayscale
     gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) == 3 else image
     
-    # Apply Gaussian blur
-    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+    # Apply bilateral filter to reduce noise while preserving edges
+    filtered = cv2.bilateralFilter(gray, 9, 75, 75)
+    
+    # Apply adaptive threshold
+    thresh = cv2.adaptiveThreshold(filtered, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
+                                   cv2.THRESH_BINARY, 11, 2)
     
-    # Edge detection
-    edges = cv2.Canny(blurred, 50, 150)
+    # Edge detection with multiple thresholds
+    edges1 = cv2.Canny(filtered, 50, 150)
+    edges2 = cv2.Canny(filtered, 30, 100)
+    edges = cv2.bitwise_or(edges1, edges2)
+    
+    # Morphological operations to connect broken edges
+    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
+    edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
     
     # Find contours
     contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     
     # Filter contours by area and aspect ratio to find paper-like rectangles
     paper_contours = []
-    min_area = (image.shape[0] * image.shape[1]) * 0.1  # At least 10% of image
+    image_area = image.shape[0] * image.shape[1]
+    min_area = image_area * 0.15  # At least 15% of image
+    max_area = image_area * 0.95  # At most 95% of image
     
     for contour in contours:
         area = cv2.contourArea(contour)
-        if area > min_area:
+        if min_area < area < max_area:
             # Approximate contour to polygon
             epsilon = 0.02 * cv2.arcLength(contour, True)
             approx = cv2.approxPolyDP(contour, epsilon, True)
             
-            # Check if it's roughly rectangular (4 corners)
+            # Check if it's roughly rectangular (4 corners) or close to it
             if len(approx) >= 4:
                 # Calculate bounding rectangle
                 rect = cv2.boundingRect(approx)
-                aspect_ratio = rect[2] / rect[3]  # width / height
+                w, h = rect[2], rect[3]
+                aspect_ratio = w / h if h > 0 else 0
                 
                 # Check if aspect ratio matches common paper ratios
                 # A4: 1.414, A3: 1.414, US Letter: 1.294
-                if 0.7 < aspect_ratio < 1.8:  # Allow some tolerance
-                    paper_contours.append((contour, area, aspect_ratio))
+                if 0.6 < aspect_ratio < 2.0:  # More lenient tolerance
+                    # Check if contour area is close to bounding rect area (rectangularity)
+                    rect_area = w * h
+                    if rect_area > 0:
+                        extent = area / rect_area
+                        if extent > 0.7:  # At least 70% rectangular
+                            paper_contours.append((contour, area, aspect_ratio, extent))
     
     if not paper_contours:
-        raise PaperNotDetectedError("Could not detect paper in the image")
+        logger.error("No paper-like contours found")
+        raise ReferenceBoxNotDetectedError("Could not detect paper in the image using contour detection")
     
-    # Select the largest paper-like contour
-    paper_contours.sort(key=lambda x: x[1], reverse=True)
+    # Select the best paper contour based on area and rectangularity
+    paper_contours.sort(key=lambda x: (x[1] * x[3]), reverse=True)  # Sort by area * extent
     best_contour = paper_contours[0][0]
     
+    logger.info(f"Paper detected using contours: area={paper_contours[0][1]}, aspect_ratio={paper_contours[0][2]:.2f}")
+    
     return best_contour, 0.0  # Return 0.0 as placeholder scaling factor
 
 def detect_paper_bounds(image: np.ndarray, paper_size: str) -> Tuple[np.ndarray, float]:
@@ -201,13 +233,24 @@ def detect_paper_bounds(image: np.ndarray, paper_size: str) -> Tuple[np.ndarray,
         
         if paper_detector is not None:
             # Use trained model if available
-            results = paper_detector.predict(image, conf=0.9)
-            if not results or len(results) == 0 or len(results[0].boxes) == 0:
-                logger.warning("Model detection failed, using fallback contour detection")
+            # FIXED: Add verbose=False to suppress prints, and use proper confidence threshold
+            results = paper_detector.predict(image, conf=0.3, verbose=False)  # Lower confidence threshold
+            
+            if not results or len(results) == 0:
+                logger.warning("No results from paper detector")
+                return detect_paper_contour(image)
+                
+            # Check if boxes exist and are not empty
+            if not hasattr(results[0], 'boxes') or results[0].boxes is None or len(results[0].boxes) == 0:
+                logger.warning("No boxes detected by model, using fallback contour detection")
                 return detect_paper_contour(image)
             
             # Get the largest detected paper
-            boxes = results[0].cpu().boxes.xyxy
+            boxes = results[0].boxes.xyxy.cpu().numpy()  # Convert to numpy
+            if len(boxes) == 0:
+                logger.warning("Empty boxes detected, using fallback")
+                return detect_paper_contour(image)
+                
             largest_box = None
             max_area = 0
             
@@ -219,7 +262,8 @@ def detect_paper_bounds(image: np.ndarray, paper_size: str) -> Tuple[np.ndarray,
                     largest_box = box
             
             if largest_box is None:
-                raise PaperNotDetectedError("No paper detected by model")
+                logger.warning("No valid paper box found, using fallback")
+                return detect_paper_contour(image)
                 
             # Convert box to contour-like format
             x_min, y_min, x_max, y_max = map(int, largest_box)
@@ -230,8 +274,11 @@ def detect_paper_bounds(image: np.ndarray, paper_size: str) -> Tuple[np.ndarray,
                 [[x_min, y_max]]
             ])
             
+            logger.info(f"Paper detected by model: {x_min},{y_min} to {x_max},{y_max}")
+            
         else:
             # Use fallback contour detection
+            logger.info("Using fallback contour detection for paper")
             paper_contour, _ = detect_paper_contour(image)
         
         # Calculate scaling factor based on paper size
@@ -241,7 +288,8 @@ def detect_paper_bounds(image: np.ndarray, paper_size: str) -> Tuple[np.ndarray,
         
     except Exception as e:
         logger.error(f"Error in paper detection: {e}")
-        raise PaperNotDetectedError(f"Failed to detect paper: {str(e)}")
+        # Instead of raising PaperNotDetectedError, raise ReferenceBoxNotDetectedError
+        raise ReferenceBoxNotDetectedError(f"Failed to detect paper: {str(e)}")
 
 def calculate_paper_scaling_factor(paper_contour: np.ndarray, paper_size: str) -> float:
     """
@@ -732,6 +780,9 @@ def make_square(img: np.ndarray):
 def predict_with_paper(image, paper_size, offset, offset_unit, edge_radius, 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}, Edge radius: {edge_radius}")
+    
     if offset_unit == "inches":
         offset *= 25.4
 
@@ -743,15 +794,18 @@ def predict_with_paper(image, paper_size, offset, offset_unit, edge_radius, fing
 
     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 with scaling factor: {scaling_factor:.4f} mm/px")
+        logger.info(f"Paper detected successfully with scaling factor: {scaling_factor:.4f} mm/px")
         
-    except PaperNotDetectedError as e:
+    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: