Update working_yolo_pipeline.py

working_yolo_pipeline.py

@@ -25,6 +25,95 @@ import shutil
 from sklearn.feature_extraction.text import CountVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
 
+
+
+
+
+
+#=============================================================================
+#-----EXPERIMENT LATEX
+#=============================================================================
+
+
+# --- NEW IMPORTS ---
+from pix2text import Pix2Text
+import logging
+# -------------------
+
+# ============================================================================
+# --- CONFIGURATION AND CONSTANTS ---
+# ... (Your existing constants like WEIGHTS_PATH, OCR_JSON_OUTPUT_DIR, etc.)
+# ============================================================================
+
+# ============================================================================
+# --- PIX2TEXT INITIALIZATION AND HELPER ---
+# ============================================================================
+# Set up logging to WARNING level to suppress excessive output from model libraries
+logging.basicConfig(level=logging.WARNING)
+logging.getLogger('pix2text').setLevel(logging.WARNING) 
+
+# Initialize Pix2Text model globally (expensive operation, do it once)
+p2t = None
+try:
+    # Use 'yolox_tiny' for potentially faster inference in a pipeline context
+    p2t = Pix2Text(analyzer_config={'model_name': 'yolox_tiny'}) 
+    print("✅ Pix2Text model initialized successfully for equation conversion.")
+except Exception as e:
+    print(f"❌ Error initializing Pix2Text model. Equations will not be converted: {e}")
+    p2t = None
+
+def get_latex_from_base64(base64_string: str) -> str:
+    """
+    Decodes a Base64 image string, uses Pix2Text to recognize the formula, 
+    and returns the LaTeX code, wrapped in $$.
+    """
+    if p2t is None:
+        return "[P2T_ERROR: Model not initialized]"
+
+    try:
+        # 1. Decode Base64 to Image
+        image_data = base64.b64decode(base64_string)
+        image = Image.open(io.BytesIO(image_data))
+        
+        # 2. Recognize text and formulas
+        result = p2t.recognize(image, save_formula_images=False, use_analyzer=True)
+        
+        # 3. Parse the result for LaTeX
+        extracted_latex_parts = []
+        if isinstance(result, list):
+            for item in result:
+                if hasattr(item, 'text'):
+                    extracted_latex_parts.append(item.text)
+                elif isinstance(item, str):
+                    extracted_latex_parts.append(item)
+        elif isinstance(result, str):
+             extracted_latex_parts = [result]
+        
+        extracted_latex = " ".join(extracted_latex_parts).strip()
+        
+        if not extracted_latex:
+             return "[P2T_WARNING: No formula found]"
+
+        # Wrap result in LaTeX delimiters
+        return f"$${extracted_latex}$$"
+
+    except Exception as e:
+        # Catch any unexpected errors
+        print(f"  ❌ Pix2Text Recognition failed: {e}")
+        return f"[P2T_ERROR: Recognition failed: {e}]"
+
+
+
+
+
+
+#=============================================================================
+#-----EXPERIMENT LATEX
+#=============================================================================
+
+
+
+
 # ============================================================================
 # --- CONFIGURATION AND CONSTANTS ---
 # ============================================================================
@@ -1466,10 +1555,50 @@ def get_base64_for_file(filepath: str) -> str:
         return ""
 
 
+# def embed_images_as_base64_in_memory(structured_data: List[Dict[str, Any]], figure_extraction_dir: str) -> List[
+#     Dict[str, Any]]:
+#     print("\n" + "=" * 80)
+#     print("--- 4. STARTING IMAGE EMBEDDING (Base64) ---")
+#     print("=" * 80)
+#     if not structured_data: return []
+#     image_files = glob.glob(os.path.join(figure_extraction_dir, "*.png"))
+#     image_lookup = {}
+#     tag_regex = re.compile(r'(figure|equation)(\d+)', re.IGNORECASE)
+#     for filepath in image_files:
+#         filename = os.path.basename(filepath)
+#         match = re.search(r'_(figure|equation)(\d+)\.png$', filename, re.IGNORECASE)
+#         if match:
+#             key = f"{match.group(1).upper()}{match.group(2)}"
+#             image_lookup[key] = filepath
+#     print(f"  -> Found {len(image_lookup)} image components.")
+#     final_structured_data = []
+#     for item in structured_data:
+#         text_fields = [item.get('question', ''), item.get('passage', '')]
+#         if 'options' in item:
+#             for opt_val in item['options'].values(): text_fields.append(opt_val)
+#         if 'new_passage' in item: text_fields.append(item['new_passage'])
+#         unique_tags_to_embed = set()
+#         for text in text_fields:
+#             if not text: continue
+#             for match in tag_regex.finditer(text):
+#                 tag = match.group(0).upper()
+#                 if tag in image_lookup: unique_tags_to_embed.add(tag)
+#         for tag in sorted(list(unique_tags_to_embed)):
+#             filepath = image_lookup[tag]
+#             base64_code = get_base64_for_file(filepath)
+#             base_key = tag.replace(' ', '').lower()
+#             item[base_key] = base64_code
+#         final_structured_data.append(item)
+#     print(f"✅ Image embedding complete.")
+#     return final_structured_data
+
+
+
+
 def embed_images_as_base64_in_memory(structured_data: List[Dict[str, Any]], figure_extraction_dir: str) -> List[
     Dict[str, Any]]:
     print("\n" + "=" * 80)
-    print("--- 4. STARTING IMAGE EMBEDDING (Base64) ---")
+    print("--- 4. STARTING IMAGE EMBEDDING (Base64) / EQUATION TO LATEX CONVERSION ---")
     print("=" * 80)
     if not structured_data: return []
     image_files = glob.glob(os.path.join(figure_extraction_dir, "*.png"))
@@ -1482,25 +1611,57 @@ def embed_images_as_base64_in_memory(structured_data: List[Dict[str, Any]], figu
             key = f"{match.group(1).upper()}{match.group(2)}"
             image_lookup[key] = filepath
     print(f"  -> Found {len(image_lookup)} image components.")
+    
     final_structured_data = []
+    
     for item in structured_data:
         text_fields = [item.get('question', ''), item.get('passage', '')]
         if 'options' in item:
             for opt_val in item['options'].values(): text_fields.append(opt_val)
         if 'new_passage' in item: text_fields.append(item['new_passage'])
+            
         unique_tags_to_embed = set()
         for text in text_fields:
             if not text: continue
             for match in tag_regex.finditer(text):
                 tag = match.group(0).upper()
                 if tag in image_lookup: unique_tags_to_embed.add(tag)
+                
+        # List of tags that were successfully converted to LaTeX
+        tags_converted_to_latex = set()
+        
         for tag in sorted(list(unique_tags_to_embed)):
             filepath = image_lookup[tag]
+            # Get the base64 code for processing, whether we embed it or convert it to LaTeX
             base64_code = get_base64_for_file(filepath)
+            
+            # --- PIX2TEXT/EQUATION CONVERSION LOGIC START ---
+            if tag.startswith('EQUATION') and p2t is not None:
+                print(f"  -> Converting EQUATION {tag} to LaTeX...")
+                latex_code = get_latex_from_base64(base64_code)
+                
+                # Replace the original tag (e.g., EQUATION1) in the item's text fields with LaTeX
+                for key in ['question', 'passage', 'new_passage']:
+                    if item.get(key) and tag in item[key]:
+                        item[key] = item[key].replace(tag, latex_code)
+                
+                if 'options' in item:
+                    for opt_key, opt_val in item['options'].items():
+                        if tag in opt_val:
+                            item['options'][opt_key] = opt_val.replace(tag, latex_code)
+
+                tags_converted_to_latex.add(tag)
+                # Skip the embedding of the Base64 code for equations
+                continue
+            # --- PIX2TEXT/EQUATION CONVERSION LOGIC END ---
+
+            # Original logic (for figures): Embed the base64 code
             base_key = tag.replace(' ', '').lower()
             item[base_key] = base64_code
+            
         final_structured_data.append(item)
-    print(f"✅ Image embedding complete.")
+    
+    print(f"✅ Image embedding complete. {len(tags_converted_to_latex)} equations converted to LaTeX.")
     return final_structured_data
 
 
@@ -1508,7 +1669,76 @@ def embed_images_as_base64_in_memory(structured_data: List[Dict[str, Any]], figu
 # --- MAIN FUNCTION ---
 # ============================================================================
 
-def run_document_pipeline(input_pdf_path: str, layoutlmv3_model_path: str, label_studio_output_path: str) -> Optional[
+# def run_document_pipeline(input_pdf_path: str, layoutlmv3_model_path: str, label_studio_output_path: str) -> Optional[
+#     List[Dict[str, Any]]]:
+#     if not os.path.exists(input_pdf_path): return None
+
+#     print("\n" + "#" * 80)
+#     print("### STARTING OPTIMIZED FULL DOCUMENT ANALYSIS PIPELINE ###")
+#     print("#" * 80)
+
+#     pdf_name = os.path.splitext(os.path.basename(input_pdf_path))[0]
+#     temp_pipeline_dir = os.path.join(tempfile.gettempdir(), f"pipeline_run_{pdf_name}_{os.getpid()}")
+#     os.makedirs(temp_pipeline_dir, exist_ok=True)
+
+#     preprocessed_json_path = os.path.join(temp_pipeline_dir, f"{pdf_name}_preprocessed.json")
+#     raw_output_path = os.path.join(temp_pipeline_dir, f"{pdf_name}_raw_predictions.json")
+#     structured_intermediate_output_path = os.path.join(temp_pipeline_dir, f"{pdf_name}_structured_intermediate.json")
+
+#     final_result = None
+#     try:
+#         # Phase 1: Preprocessing with YOLO First + Masking
+#         preprocessed_json_path_out = run_single_pdf_preprocessing(input_pdf_path, preprocessed_json_path)
+#         if not preprocessed_json_path_out: return None
+
+#         # Phase 2: Inference
+#         page_raw_predictions_list = run_inference_and_get_raw_words(
+#             input_pdf_path, layoutlmv3_model_path, preprocessed_json_path_out
+#         )
+#         if not page_raw_predictions_list: return None
+
+#         with open(raw_output_path, 'w', encoding='utf-8') as f:
+#             json.dump(page_raw_predictions_list, f, indent=4)
+
+#         # Phase 3: Decoding
+#         structured_data_list = convert_bio_to_structured_json_relaxed(
+#             raw_output_path, structured_intermediate_output_path
+#         )
+#         if not structured_data_list: return None
+#         structured_data_list = correct_misaligned_options(structured_data_list)
+#         structured_data_list = process_context_linking(structured_data_list)
+
+#         try:
+#             convert_raw_predictions_to_label_studio(page_raw_predictions_list, label_studio_output_path)
+#         except Exception as e:
+#             print(f"❌ Error during Label Studio conversion: {e}")
+
+#         # Phase 4: Embedding
+#         final_result = embed_images_as_base64_in_memory(structured_data_list, FIGURE_EXTRACTION_DIR)
+
+#     except Exception as e:
+#         print(f"❌ FATAL ERROR: {e}")
+#         import traceback
+#         traceback.print_exc()
+#         return None
+
+#     finally:
+#         try:
+#             for f in glob.glob(os.path.join(temp_pipeline_dir, '*')):
+#                 os.remove(f)
+#             os.rmdir(temp_pipeline_dir)
+#         except Exception:
+#             pass
+
+#     print("\n" + "#" * 80)
+#     print("### OPTIMIZED PIPELINE EXECUTION COMPLETE ###")
+#     print("#" * 80)
+#     return final_result
+
+
+
+
+def run_document_pipeline(input_pdf_path: str, layoutlmv3_model_path: str, label_studio_output_path: str, raw_predictions_output_path: str) -> Optional[
     List[Dict[str, Any]]]:
     if not os.path.exists(input_pdf_path): return None
 
@@ -1536,9 +1766,17 @@ def run_document_pipeline(input_pdf_path: str, layoutlmv3_model_path: str, label
         )
         if not page_raw_predictions_list: return None
 
+        # --- DEBUG STEP: SAVE RAW PREDICTIONS ---
+        # Save raw predictions to the temporary file
         with open(raw_output_path, 'w', encoding='utf-8') as f:
             json.dump(page_raw_predictions_list, f, indent=4)
 
+        # Explicitly copy/save the raw predictions to the user-specified debug path
+        if raw_predictions_output_path:
+            shutil.copy(raw_output_path, raw_predictions_output_path)
+            print(f"\n✅ DEBUG: Raw predictions saved to: {raw_predictions_output_path}")
+        # ----------------------------------------
+
         # Phase 3: Decoding
         structured_data_list = convert_bio_to_structured_json_relaxed(
             raw_output_path, structured_intermediate_output_path
@@ -1552,7 +1790,7 @@ def run_document_pipeline(input_pdf_path: str, layoutlmv3_model_path: str, label
         except Exception as e:
             print(f"❌ Error during Label Studio conversion: {e}")
 
-        # Phase 4: Embedding
+        # Phase 4: Embedding / Equation to LaTeX Conversion
         final_result = embed_images_as_base64_in_memory(structured_data_list, FIGURE_EXTRACTION_DIR)
 
     except Exception as e:
@@ -1575,19 +1813,61 @@ def run_document_pipeline(input_pdf_path: str, layoutlmv3_model_path: str, label
     return final_result
 
 
+    
+
+# if __name__ == "__main__":
+#     parser = argparse.ArgumentParser(description="Complete Pipeline")
+#     parser.add_argument("--input_pdf", type=str, required=True, help="Input PDF")
+#     parser.add_argument("--layoutlmv3_model_path", type=str, default=DEFAULT_LAYOUTLMV3_MODEL_PATH, help="Model Path")
+#     parser.add_argument("--ls_output_path", type=str, default=None, help="Label Studio Output Path")
+#     args = parser.parse_args()
+
+#     pdf_name = os.path.splitext(os.path.basename(args.input_pdf))[0]
+#     final_output_path = os.path.abspath(f"{pdf_name}_final_output_embedded.json")
+#     ls_output_path = os.path.abspath(
+#         args.ls_output_path if args.ls_output_path else f"{pdf_name}_label_studio_tasks.json")
+
+#     final_json_data = run_document_pipeline(args.input_pdf, args.layoutlmv3_model_path, ls_output_path)
+
+#     if final_json_data:
+#         with open(final_output_path, 'w', encoding='utf-8') as f:
+#             json.dump(final_json_data, f, indent=2, ensure_ascii=False)
+#         print(f"\n✅ Final Data Saved: {final_output_path}")
+#     else:
+#         print("\n❌ Pipeline Failed.")
+#         sys.exit(1)
+
+
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Complete Pipeline")
     parser.add_argument("--input_pdf", type=str, required=True, help="Input PDF")
     parser.add_argument("--layoutlmv3_model_path", type=str, default=DEFAULT_LAYOUTLMV3_MODEL_PATH, help="Model Path")
     parser.add_argument("--ls_output_path", type=str, default=None, help="Label Studio Output Path")
+    # --- ADDED ARGUMENT FOR DEBUGGING ---
+    parser.add_argument("--raw_preds_path", type=str, default='BIO_debug.json',
+                        help="Debug path for raw BIO tag predictions (JSON).")
+    # ------------------------------------
     args = parser.parse_args()
 
     pdf_name = os.path.splitext(os.path.basename(args.input_pdf))[0]
     final_output_path = os.path.abspath(f"{pdf_name}_final_output_embedded.json")
     ls_output_path = os.path.abspath(
         args.ls_output_path if args.ls_output_path else f"{pdf_name}_label_studio_tasks.json")
-
-    final_json_data = run_document_pipeline(args.input_pdf, args.layoutlmv3_model_path, ls_output_path)
+    # --- CALCULATE RAW PREDICTIONS OUTPUT PATH ---
+    raw_predictions_output_path = os.path.abspath(
+        args.raw_preds_path if args.raw_preds_path else f"{pdf_name}_raw_predictions_debug.json")
+    # ---------------------------------------------
+
+    # --- UPDATED FUNCTION CALL ---
+    final_json_data = run_document_pipeline(
+        args.input_pdf, 
+        args.layoutlmv3_model_path, 
+        ls_output_path, 
+        raw_predictions_output_path # Pass the new argument
+    )
+    # -----------------------------
 
     if final_json_data:
         with open(final_output_path, 'w', encoding='utf-8') as f: