Update working_yolo_pipeline.py

working_yolo_pipeline.py

@@ -1390,12 +1390,25 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
     # ====================================================================
     items_to_sort = []
 
+    # for ocr_word in raw_ocr_output:
+    #     is_suppressed = False
+    #     for component in component_metadata:
+    #         # Do not include words that are inside figure/equation boxes
+    #         ioa = calculate_ioa(ocr_word['bbox'], component['bbox'])
+    #         if ioa > IOA_SUPPRESSION_THRESHOLD:
+    #             is_suppressed = True
+    #             break
+    #     if not is_suppressed:
+    #         items_to_sort.append(ocr_word)
+
     for ocr_word in raw_ocr_output:
         is_suppressed = False
         for component in component_metadata:
-            # Do not include words that are inside figure/equation boxes
-            ioa = calculate_ioa(ocr_word['bbox'], component['bbox'])
-            if ioa > IOA_SUPPRESSION_THRESHOLD:
+            ioa_ocr_in_eq = calculate_ioa(ocr_word['bbox'], component['bbox'])
+            ioa_eq_in_ocr = calculate_ioa(component['bbox'], ocr_word['bbox'])
+            # Suppress if either: OCR word is mostly inside equation box,
+            # OR equation box is mostly inside OCR word box
+            if ioa_ocr_in_eq > IOA_SUPPRESSION_THRESHOLD or ioa_eq_in_ocr > IOA_SUPPRESSION_THRESHOLD:
                 is_suppressed = True
                 break
         if not is_suppressed: