Update app.py

app.py

@@ -1531,19 +1531,79 @@ def get_latex_from_base64(base64_string: str) -> str:
 
 
 
+# def run_yolo_detection_and_count(
+#         image: np.ndarray, model: YOLO, page_num: int, 
+#         current_eq_count: int, current_fig_count: int
+# ) -> Tuple[List[Dict[str, Union[Image.Image, str, Tuple[float,...]]]], int, int]:
+#     """
+#     Performs YOLO detection and returns a list of detected item dictionaries 
+#     and the updated total counters.
+#     """
+    
+#     eq_counter = current_eq_count
+#     fig_counter = current_fig_count
+    
+#     detected_items: List[Dict[str, Union[Image.Image, str, Tuple[float,...]]]] = []
+#     yolo_detections = []
+    
+#     try:
+#         results = model.predict(image, conf=CONF_THRESHOLD, verbose=False)
+#         if results and results[0].boxes:
+#             for box in results[0].boxes.data.tolist():
+#                 x1, y1, x2, y2, conf, cls_id = box
+#                 cls_name = model.names[int(cls_id)]
+#                 if cls_name in TARGET_CLASSES:
+#                     yolo_detections.append({
+#                         'coords': (x1, y1, x2, y2),
+#                         'class': cls_name,
+#                         'conf': conf
+#                     })
+#     except Exception as e:
+#         logging.error(f"ERROR: YOLO inference failed on page {page_num}: {e}")
+#         return [], eq_counter, fig_counter
+
+#     merged_detections = merge_overlapping_boxes(yolo_detections, IOU_MERGE_THRESHOLD)
+#     final_detections = filter_nested_boxes(merged_detections, IOA_SUPPRESSION_THRESHOLD)
+
+#     for det in final_detections:
+#         bbox = det["coords"]
+#         crop_pil = crop_and_convert_to_pil(image, bbox)
+        
+#         item = {
+#             "type": det["class"],
+#             "coords": bbox,
+#             "pil_image": crop_pil,
+#             "page_num": page_num,  # ← ADD THIS LINE
+#         }
+
+#         if det["class"] == "equation":
+#             eq_counter += 1
+#             item["id"] = f"EQUATION{eq_counter}"
+#             item["latex"] = ""
+#         elif det["class"] == "figure":
+#             fig_counter += 1
+#             item["id"] = f"FIGURE{fig_counter}"
+#             item["latex"] = "[FIGURE - No LaTeX]"
+            
+#         detected_items.append(item)
+
+#     return detected_items, eq_counter, fig_counter
+
+
+
+
 def run_yolo_detection_and_count(
         image: np.ndarray, model: YOLO, page_num: int, 
         current_eq_count: int, current_fig_count: int
 ) -> Tuple[List[Dict[str, Union[Image.Image, str, Tuple[float,...]]]], int, int]:
     """
-    Performs YOLO detection and returns a list of detected item dictionaries 
-    and the updated total counters.
+    Performs YOLO detection and returns detected items with counters.
+    This version is for the Gradio script.
     """
-    
     eq_counter = current_eq_count
     fig_counter = current_fig_count
     
-    detected_items: List[Dict[str, Union[Image.Image, str, Tuple[float,...]]]] = []
+    detected_items = []
     yolo_detections = []
     
     try:
@@ -1562,9 +1622,11 @@ def run_yolo_detection_and_count(
         logging.error(f"ERROR: YOLO inference failed on page {page_num}: {e}")
         return [], eq_counter, fig_counter
 
+    # CRITICAL: Use exact same processing order
     merged_detections = merge_overlapping_boxes(yolo_detections, IOU_MERGE_THRESHOLD)
     final_detections = filter_nested_boxes(merged_detections, IOA_SUPPRESSION_THRESHOLD)
 
+    # Create items from final detections
     for det in final_detections:
         bbox = det["coords"]
         crop_pil = crop_and_convert_to_pil(image, bbox)
@@ -1573,7 +1635,7 @@ def run_yolo_detection_and_count(
             "type": det["class"],
             "coords": bbox,
             "pil_image": crop_pil,
-            "page_num": page_num,  # ← ADD THIS LINE
+            "page_num": page_num,
         }
 
         if det["class"] == "equation":
@@ -1592,11 +1654,6 @@ def run_yolo_detection_and_count(
 
 
 
-
-
-
-
-
 # ============================================================================
 # --- MAIN DOCUMENT PROCESSING FUNCTION (Retained Logic) ---
 # ============================================================================