Update app.py

app.py

@@ -171,6 +171,76 @@ def get_birefnet():
         logger.info("BiRefNet model loaded successfully")
     return birefnet
 
+# 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 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 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 = []
+#     image_area = image.shape[0] * image.shape[1]
+#     min_area = image_area * 0.20  # At least 15% of image
+#     max_area = image_area * 0.85  # At most 95% of image
+    
+#     for contour in contours:
+#         area = cv2.contourArea(contour)
+#         if min_area < area < max_area:
+#             # Approximate contour to polygon
+#             epsilon = 0.015 * cv2.arcLength(contour, True)
+#             approx = cv2.approxPolyDP(contour, epsilon, True)
+            
+#             # Check if it's roughly rectangular (4 corners) or close to it
+#             if len(approx) >= 4:
+#                 # Calculate bounding rectangle
+#                 rect = cv2.boundingRect(approx)
+#                 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 1.3 < aspect_ratio < 1.5:  # 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.85:  # At least 70% rectangular
+#                             paper_contours.append((contour, area, aspect_ratio, extent))
+    
+#     if not paper_contours:
+#         logger.error("No paper-like contours found")
+#         raise ReferenceBoxNotDetectedError("Could not detect paper in the image using contour detection")
+    
+#     # 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_contour(image: np.ndarray, output_unit: str = "mm") -> Tuple[np.ndarray, float]:
     """
     Detect paper in the image using contour detection as fallback
@@ -422,6 +492,38 @@ def remove_bg(image: np.ndarray) -> np.ndarray:
         logger.error(f"Error in BiRefNet background removal: {e}")
         raise
 
+# def exclude_paper_area(mask: np.ndarray, paper_contour: np.ndarray, expansion_factor: float = 1.2) -> np.ndarray:
+#     """
+#     Remove paper area from the mask to focus only on objects
+#     """
+#     # Create paper mask with slight expansion to ensure complete removal
+#     paper_mask = np.zeros(mask.shape[:2], dtype=np.uint8)
+    
+#     # # Expand paper contour slightly
+#     # epsilon = expansion_factor * cv2.arcLength(paper_contour, True)
+#     # expanded_contour = cv2.approxPolyDP(paper_contour, epsilon, True)
+    
+#     # cv2.fillPoly(paper_mask, [expanded_contour], 255)
+#     # Create a more aggressive inward shrinking of paper bounds
+#     rect = cv2.boundingRect(paper_contour)
+#     shrink_pixels = int(min(rect[2], rect[3]) * 0.05)  # Shrink by 5% of smaller dimension
+    
+#     # Create shrunken rectangle
+#     x, y, w, h = rect
+#     shrunken_contour = np.array([
+#         [[x + shrink_pixels, y + shrink_pixels]],
+#         [[x + w - shrink_pixels, y + shrink_pixels]],
+#         [[x + w - shrink_pixels, y + h - shrink_pixels]],
+#         [[x + shrink_pixels, y + h - shrink_pixels]]
+#     ])
+    
+#     cv2.fillPoly(paper_mask, [shrunken_contour], 255)
+    
+#     # Invert paper mask and apply to object mask
+#     paper_mask_inv = cv2.bitwise_not(paper_mask)
+#     result_mask = cv2.bitwise_and(mask, paper_mask_inv)
+    
+#     return result_mask
 def mask_paper_area_in_image(image: np.ndarray, paper_contour: np.ndarray) -> np.ndarray:
     """Less aggressive masking to preserve corner objects"""
     masked_image = image.copy()
@@ -1000,7 +1102,7 @@ def predict_with_paper(image, paper_size, offset, offset_unit, finger_clearance=
         scale_info
     )
 
-def predict_full_paper(image, paper_size, offset_value_mm = 0.02,offset_unit='mm', enable_finger_cut='Off', selected_outputs):
+def predict_full_paper(image, paper_size, offset_value_mm,offset_unit, enable_finger_cut, selected_outputs):
     finger_flag = "On" if enable_finger_cut == "On" else "Off"
     
     # Always get all outputs from predict_with_paper
@@ -1026,29 +1128,18 @@ if __name__ == "__main__":
     os.makedirs("./outputs", exist_ok=True)
 
     with gr.Blocks(title="Paper-Based DXF Generator", theme=gr.themes.Soft()) as demo:
-         # Example gallery
         gr.Markdown("""
-            # Paper-Based DXF Generator
-            
-            Upload an image with a single object placed on paper (A4, A3, or US Letter).
-            The paper serves as a size reference for accurate DXF generation.
-            """)
-        with gr.Row():
-            gr.Markdown("""
-            **Instructions:**
-            1. Place a single object on paper
-            2. Select the correct paper size
-            3. Configure options as needed
-            4. Click Submit to generate DXF
-            """)
-            gr.Markdown("""
-            ### Tips for Best Results:
-            - Ensure good lighting and clear paper edges
-            - Place object completely at the center of the paper
-            - Avoid shadows that might interfere with detection
-            - Use high contrast between object and paper
-            """)
+        # Paper-Based DXF Generator
         
+        Upload an image with a single object placed on paper (A4, A3, or US Letter).
+        The paper serves as a size reference for accurate DXF generation.
+        
+        **Instructions:**
+        1. Place a single object on paper
+        2. Select the correct paper size
+        3. Configure options as needed
+        4. Click Submit to generate DXF
+        """)
         
         with gr.Row():
             with gr.Column():
@@ -1065,31 +1156,31 @@ if __name__ == "__main__":
                     info="Select the paper size used in your image"
                 )
                 
-                # with gr.Group():
-                #     gr.Markdown("### Contour Offset")
-                #     with gr.Row():
-                #         offset_value_mm = gr.Number(
-                #             value=0.02,
-                #             label="Offset",
-                #             info="Expand contours outward by this amount",
-                #             precision=3,
-                #             minimum=0,
-                #             maximum=50
-                #         )
-                #         offset_unit = gr.Dropdown(
-                #             choices=["mm", "inches"],
-                #             value="mm",
-                #             label="Unit"
-                #         )
-
-                # with gr.Group():
-                #     gr.Markdown("### Finger Cuts")
-                #     enable_finger_cut = gr.Radio(
-                #         choices=["On", "Off"],
-                #         value="Off",
-                #         label="Enable Finger Cuts",
-                #         info="Add circular cuts for easier handling"
-                #     )
+                with gr.Group():
+                    gr.Markdown("### Contour Offset")
+                    with gr.Row():
+                        offset_value_mm = gr.Number(
+                            value=0.02,
+                            label="Offset",
+                            info="Expand contours outward by this amount",
+                            precision=3,
+                            minimum=0,
+                            maximum=50
+                        )
+                        offset_unit = gr.Dropdown(
+                            choices=["mm", "inches"],
+                            value="mm",
+                            label="Unit"
+                        )
+
+                with gr.Group():
+                    gr.Markdown("### Finger Cuts")
+                    enable_finger_cut = gr.Radio(
+                        choices=["On", "Off"],
+                        value="Off",
+                        label="Enable Finger Cuts",
+                        info="Add circular cuts for easier handling"
+                    )
                 
                 output_options = gr.CheckboxGroup(
                     choices=["Annotated Image", "Outlines", "Mask"],
@@ -1131,14 +1222,22 @@ if __name__ == "__main__":
             inputs=[
                 input_image, 
                 paper_size, 
-                0.02,
-                'mm',
-                'Off', 
+                offset_value_mm,
+                offset_unit,
+                enable_finger_cut, 
                 output_options
             ],
             outputs=[dxf_file, output_image, outlines_image, mask_image, scale_info]
         )
         
-       
+        # Example gallery
+        with gr.Row():
+            gr.Markdown("""
+            ### Tips for Best Results:
+            - Ensure good lighting and clear paper edges
+            - Place object completely on the paper
+            - Avoid shadows that might interfere with detection
+            - Use high contrast between object and paper
+            """)
 
     demo.launch(share=True)