Update app.py

app.py

@@ -40,6 +40,7 @@ logging.basicConfig(level=logging.WARNING)
 
 WEIGHTS_PATH = 'best.pt' 
 SCALE_FACTOR = 2.0 
+OUTPUT_DIR = "yolo_extracted_regions"
 
 # Detection parameters
 CONF_THRESHOLD = 0.2
@@ -229,6 +230,26 @@ def crop_and_convert_to_base64(image: np.ndarray, bbox: Tuple[float, float, floa
 
 
 
+def crop_and_save(image: np.ndarray, bbox, label: str, index: int) -> str:
+    """Crop bounding box and save to disk. Return file path."""
+    x1, y1, x2, y2 = map(int, bbox)
+
+    h, w, _ = image.shape
+    x1 = max(0, x1)
+    y1 = max(0, y1)
+    x2 = min(w, x2)
+    y2 = min(h, y2)
+
+    crop = image[y1:y2, x1:x2]
+    filename = f"{label}{index}.png"
+    filepath = os.path.join(OUTPUT_DIR, filename)
+
+    cv2.imwrite(filepath, crop)
+
+    return filepath
+
+
+
 
 
 
@@ -237,15 +258,14 @@ def run_yolo_detection_and_count(
         image: np.ndarray, model: YOLO, page_num: int
 ) -> Tuple[int, int, List[str]]:
     """
-    Runs YOLO inference and returns:
-    equations, figures, list of base64 images
+    Runs YOLO inference, saves crops, and returns file paths.
     """
     global GLOBAL_FIGURE_COUNT, GLOBAL_EQUATION_COUNT
-    
+
     yolo_detections = []
     page_equations = 0
     page_figures = 0
-    cropped_images_base64 = []
+    saved_images = []
 
     try:
         results = model.predict(image, conf=CONF_THRESHOLD, verbose=False)
@@ -269,22 +289,25 @@ def run_yolo_detection_and_count(
     final_detections = filter_nested_boxes(merged_detections, IOA_SUPPRESSION_THRESHOLD)
 
     for det in final_detections:
-        bbox = det['coords']
-
-        if det['class'] == 'figure':
-            GLOBAL_FIGURE_COUNT += 1
-            page_figures += 1
+        bbox = det["coords"]
 
-        elif det['class'] == 'equation':
+        if det["class"] == "equation":
             GLOBAL_EQUATION_COUNT += 1
             page_equations += 1
 
-            # ✅ Only crop equations (you can include figures too if needed)
-            b64_img = crop_and_convert_to_base64(image, bbox)
-            cropped_images_base64.append(b64_img)
+            path = crop_and_save(image, bbox, "EQUATION", GLOBAL_EQUATION_COUNT)
+            saved_images.append(path)
+
+        elif det["class"] == "figure":
+            GLOBAL_FIGURE_COUNT += 1
+            page_figures += 1
+
+            path = crop_and_save(image, bbox, "FIGURE", GLOBAL_FIGURE_COUNT)
+            saved_images.append(path)
 
     logging.warning(f"  -> Page {page_num}: EQs={page_equations}, Figs={page_figures}")
-    return page_equations, page_figures, cropped_images_base64
+    return page_equations, page_figures, saved_images
+