Update app.py

app.py

@@ -1,22 +1,34 @@
 
-
-
-
 # import gradio as gr
 # import json
 # import os
 # import tempfile
 # import img2pdf
+# import glob
+# import shutil
 # from img2pdf import Rotation
 # from pathlib import Path
 
+
+
+# print("--- DEBUG: Current Working Directory ---")
+# print(os.getcwd())
+# print("--- DEBUG: Files in Root ---")
+# print(os.listdir('.'))
+
 # # ==============================
 # # PIPELINE IMPORT
 # # ==============================
+# # try:
+# #     from working_yolo_pipeline import run_document_pipeline, DEFAULT_LAYOUTLMV3_MODEL_PATH, WEIGHTS_PATH
+# # except ImportError:
+# #     print("Warning: 'working_yolo_pipeline.py' not found. Using dummy paths.")
 # try:
 #     from working_yolo_pipeline import run_document_pipeline, DEFAULT_LAYOUTLMV3_MODEL_PATH, WEIGHTS_PATH
-# except ImportError:
-#     print("Warning: 'working_yolo_pipeline.py' not found. Using dummy paths.")
+# except Exception as e:  # Catch ALL exceptions
+#     print(f"Warning: Failed to import pipeline: {e}")
+#     import traceback
+#     traceback.print_exc()  # Show the actual error
 #     def run_document_pipeline(*args):
 #         return {"error": "Placeholder pipeline function called."}
 #     DEFAULT_LAYOUTLMV3_MODEL_PATH = "./models/layoutlmv3_model"
@@ -25,13 +37,11 @@
 # def process_file(uploaded_files, layoutlmv3_model_path=None):
 #     """
 #     Robust handler for multiple or single file uploads.
+#     Returns the final JSON and a LIST of all intermediate JSON files (OCR, Predictions, BIO).
 #     """
 #     if uploaded_files is None:
 #         return "❌ Error: No files uploaded.", None
 
-#     # --- THE ROBUST FIX ---
-#     # Gradio sometimes sends a single dict even when set to multiple.
-#     # We force everything into a list so the rest of the logic doesn't break.
 #     if not isinstance(uploaded_files, list):
 #         file_list = [uploaded_files]
 #     else:
@@ -39,7 +49,6 @@
 
 #     if len(file_list) == 0:
 #         return "❌ Error: Empty file list.", None
-#     # ----------------------
 
 #     # 1. Resolve all file paths safely
 #     resolved_paths = []
@@ -62,17 +71,13 @@
 #     is_image = first_file.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp', '.webp', '.tiff']
     
 #     try:
-#         # If it's multiple files or just one image, wrap it in a PDF
 #         if len(resolved_paths) > 1 or is_image:
 #             print(f"📦 Converting {len(resolved_paths)} image(s) to a single PDF...")
 #             temp_pdf = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
 #             with open(temp_pdf.name, "wb") as f_out:
-#                 # f_out.write(img2pdf.convert(resolved_paths))
 #                 f_out.write(img2pdf.convert(resolved_paths, rotation=Rotation.ifvalid))
-
 #             processing_path = temp_pdf.name
 #         else:
-#             # It's a single PDF
 #             processing_path = resolved_paths[0]
 
 #         # 3. Standard Pipeline Checks
@@ -84,267 +89,277 @@
 #         print(f"🚀 Starting pipeline for: {processing_path}")
 #         result = run_document_pipeline(processing_path, final_model_path)
 
-#         if result is None:
-#             return "❌ Error: Pipeline returned None.", None
+#         # 5. SCRAPE FOR INTERMEDIATE FILES
+#         # We look for all .json files in /tmp/ created during this run
+#         base_name = Path(processing_path).stem
+#         # This matches common patterns like /tmp/pipeline_run_... or filenames in /tmp/
+#         search_patterns = [
+#             f"/tmp/pipeline_run_{base_name}*/*.json",
+#             f"/tmp/*{base_name}*.json"
+#         ]
+        
+#         all_intermediate_jsons = []
+#         for pattern in search_patterns:
+#             all_intermediate_jsons.extend(glob.glob(pattern))
+        
+#         # Remove duplicates while preserving order
+#         all_intermediate_jsons = list(dict.fromkeys(all_intermediate_jsons))
 
-#         # 5. Prepare output
-#         temp_output = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json', prefix='analysis_')
-#         with open(temp_output.name, 'w', encoding='utf-8') as f:
-#             json.dump(result, f, indent=2, ensure_ascii=False)
+#         # 6. Prepare Final Output for Display
+#         if result is None or (isinstance(result, list) and len(result) == 0):
+#             display_text = "⚠️ Pipeline failed at Step 3 (BIO Decoding).\nDownload the intermediate JSONs below to inspect OCR and Model Predictions."
+#         else:
+#             display_text = json.dumps(result, indent=2, ensure_ascii=False)
+            
+#             # If the final result succeeded, save it to a temp file so it can be downloaded too
+#             temp_final = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json', prefix='final_result_')
+#             json.dump(result, temp_final, indent=2, ensure_ascii=False)
+#             temp_final.close()
+#             all_intermediate_jsons.append(temp_final.name)
 
-#         return json.dumps(result, indent=2, ensure_ascii=False), temp_output.name
+#         return display_text, all_intermediate_jsons
 
 #     except Exception as e:
 #         import traceback
 #         traceback.print_exc()
 #         return f"❌ Error: {str(e)}", None
 
-# # ==============================
-# # GRADIO INTERFACE
-# # ==============================
-# with gr.Blocks(title="Document Analysis Pipeline") as demo:
-
-#     gr.Markdown("# 📄 Document & Image Analysis Pipeline")
-
-#     with gr.Row():
-#         with gr.Column(scale=1):
-#             file_input = gr.File(
-#                 label="Upload PDFs or Images",
-#                 file_types=[".pdf", ".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"],
-#                 file_count="multiple", # Keep this
-#                 type="filepath"       # Keep this
-#             )
-
-#             model_path_input = gr.Textbox(
-#                 label="Model Path",
-#                 value=DEFAULT_LAYOUTLMV3_MODEL_PATH
-#             )
 
-#             process_btn = gr.Button("🚀 Process Files", variant="primary")
 
-#         with gr.Column(scale=2):
-#             json_output = gr.Code(label="JSON Output", language="json", lines=20)
-#             download_output = gr.File(label="Download JSON")
 
-#     process_btn.click(
-#         fn=process_file,
-#         inputs=[file_input, model_path_input],
-#         outputs=[json_output, download_output]
-#     )
 
-# if __name__ == "__main__":
-#     demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
+# # def visualize_detections(uploaded_files):
+# #     """Shows the first uploaded image with YOLO bounding boxes"""
+# #     if not uploaded_files:
+# #         return None
+    
+# #     # Get first file path
+# #     file_path = uploaded_files[0] if isinstance(uploaded_files, list) else uploaded_files
+# #     if isinstance(file_path, dict):
+# #         file_path = file_path["path"]
+    
+# #     import cv2
+# #     from ultralytics import YOLO
+    
+# #     # Load image
+# #     img = cv2.imread(str(file_path))
+# #     if img is None:
+# #         return None
+    
+# #     # Run YOLO
+# #     model = YOLO(WEIGHTS_PATH)
+# #     results = model.predict(source=img, conf=0.2, imgsz=640, verbose=False)
+    
+# #     # Draw boxes
+# #     for box in results[0].boxes:
+# #         class_id = int(box.cls[0])
+# #         class_name = model.names[class_id]
+# #         if class_name in ['figure', 'equation']:
+# #             x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
+# #             color = (0, 255, 0) if class_name == 'figure' else (255, 0, 0)
+# #             cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
+# #             cv2.putText(img, f"{class_name} {box.conf[0]:.2f}", 
+# #                        (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+    
+# #     # Save and return
+# #     temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".png").name
+# #     cv2.imwrite(temp_path, img)
+# #     return temp_path
 
 
+# # # ==============================
+# # # GRADIO INTERFACE
+# # # ==============================
+# # with gr.Blocks(title="Document Analysis Pipeline") as demo:
 
+# #     gr.Markdown("# 📄 Full Pipeline Analysis")
+# #     gr.Markdown("### 🔍 Intermediate File Recovery Active")
+# #     gr.Markdown("The **Download** box will contain: \n1. OCR JSON (Step 1)\n2. Raw LayoutLMv3 Prediction JSON (Step 2)\n3. Final BIO JSON (Step 3)")
 
+# #     with gr.Row():
+# #         with gr.Column(scale=1):
+# #             file_input = gr.File(
+# #                 label="Upload PDFs or Images",
+# #                 file_types=[".pdf", ".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"],
+# #                 file_count="multiple",
+# #                 type="filepath"
+# #             )
 
-import gradio as gr
-import json
-import os
-import tempfile
-import img2pdf
-import glob
-import shutil
-from img2pdf import Rotation
-from pathlib import Path
+# #             model_path_input = gr.Textbox(
+# #                 label="Model Path",
+# #                 value=DEFAULT_LAYOUTLMV3_MODEL_PATH
+# #             )
 
+# #             process_btn = gr.Button("🚀 Run Pipeline", variant="primary")
 
+# #         with gr.Column(scale=2):
+# #             json_output = gr.Code(label="Final Structured Output", language="json", lines=20)
+# #             # IMPORTANT: file_count="multiple" allows returning the list of all stage files
+# #             download_output = gr.File(label="Download All Pipeline Stages (JSON)", file_count="multiple")
 
-print("--- DEBUG: Current Working Directory ---")
-print(os.getcwd())
-print("--- DEBUG: Files in Root ---")
-print(os.listdir('.'))
+# #     process_btn.click(
+# #         fn=process_file,
+# #         inputs=[file_input, model_path_input],
+# #         outputs=[json_output, download_output]
+# #     )
 
-# ==============================
-# PIPELINE IMPORT
-# ==============================
-# try:
-#     from working_yolo_pipeline import run_document_pipeline, DEFAULT_LAYOUTLMV3_MODEL_PATH, WEIGHTS_PATH
-# except ImportError:
-#     print("Warning: 'working_yolo_pipeline.py' not found. Using dummy paths.")
-try:
-    from working_yolo_pipeline import run_document_pipeline, DEFAULT_LAYOUTLMV3_MODEL_PATH, WEIGHTS_PATH
-except Exception as e:  # Catch ALL exceptions
-    print(f"Warning: Failed to import pipeline: {e}")
-    import traceback
-    traceback.print_exc()  # Show the actual error
-    def run_document_pipeline(*args):
-        return {"error": "Placeholder pipeline function called."}
-    DEFAULT_LAYOUTLMV3_MODEL_PATH = "./models/layoutlmv3_model"
-    WEIGHTS_PATH = "./weights/yolo_weights.pt"
+# # if __name__ == "__main__":
+# #     demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
 
-def process_file(uploaded_files, layoutlmv3_model_path=None):
-    """
-    Robust handler for multiple or single file uploads.
-    Returns the final JSON and a LIST of all intermediate JSON files (OCR, Predictions, BIO).
-    """
-    if uploaded_files is None:
-        return "❌ Error: No files uploaded.", None
 
-    if not isinstance(uploaded_files, list):
-        file_list = [uploaded_files]
-    else:
-        file_list = uploaded_files
 
-    if len(file_list) == 0:
-        return "❌ Error: Empty file list.", None
 
-    # 1. Resolve all file paths safely
-    resolved_paths = []
-    for f in file_list:
-        try:
-            if isinstance(f, dict) and "path" in f:
-                resolved_paths.append(f["path"])
-            elif hasattr(f, 'path'):
-                resolved_paths.append(f.path)
-            else:
-                resolved_paths.append(str(f))
-        except Exception as e:
-            print(f"Error resolving path for {f}: {e}")
 
-    if not resolved_paths:
-        return "❌ Error: Could not resolve file paths.", None
 
-    # 2. Determine if we should merge into a single PDF
-    first_file = Path(resolved_paths[0])
-    is_image = first_file.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp', '.webp', '.tiff']
+# # # ==============================
+# # # VISUAL DEBUG FUNCTION
+# # # ==============================
+# # def visualize_detections(uploaded_files):
+# #     """Shows the first uploaded image with YOLO bounding boxes"""
+# #     if not uploaded_files:
+# #         return None
     
-    try:
-        if len(resolved_paths) > 1 or is_image:
-            print(f"📦 Converting {len(resolved_paths)} image(s) to a single PDF...")
-            temp_pdf = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
-            with open(temp_pdf.name, "wb") as f_out:
-                f_out.write(img2pdf.convert(resolved_paths, rotation=Rotation.ifvalid))
-            processing_path = temp_pdf.name
-        else:
-            processing_path = resolved_paths[0]
-
-        # 3. Standard Pipeline Checks
-        final_model_path = layoutlmv3_model_path or DEFAULT_LAYOUTLMV3_MODEL_PATH
-        if not os.path.exists(final_model_path):
-            return f"❌ Error: Model not found at {final_model_path}", None
+# #     try:
+# #         # Get first file path
+# #         file_path = uploaded_files[0] if isinstance(uploaded_files, list) else uploaded_files
+# #         if isinstance(file_path, dict):
+# #             file_path = file_path["path"]
+# #         elif hasattr(file_path, 'path'):
+# #             file_path = file_path.path
+        
+# #         import cv2
+# #         import numpy as np
 
-        # 4. Call the pipeline
-        print(f"🚀 Starting pipeline for: {processing_path}")
-        result = run_document_pipeline(processing_path, final_model_path)
-
-        # 5. SCRAPE FOR INTERMEDIATE FILES
-        # We look for all .json files in /tmp/ created during this run
-        base_name = Path(processing_path).stem
-        # This matches common patterns like /tmp/pipeline_run_... or filenames in /tmp/
-        search_patterns = [
-            f"/tmp/pipeline_run_{base_name}*/*.json",
-            f"/tmp/*{base_name}*.json"
-        ]
+# #         from ultralytics import YOLO
+# #         import fitz
         
-        all_intermediate_jsons = []
-        for pattern in search_patterns:
-            all_intermediate_jsons.extend(glob.glob(pattern))
+# #         # Handle PDF conversion to image
+# #         if str(file_path).lower().endswith('.pdf'):
+# #             doc = fitz.open(file_path)
+# #             page_idx = int(page_num) - 1
+# #             page = doc.load_page(page_idx) 
+
+# #             pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
+# #             img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n)
+# #             if pix.n == 3:
+# #                 img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+# #             elif pix.n == 4:
+# #                 img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR)
+# #             doc.close()
+# #         else:
+# #             img = cv2.imread(str(file_path))
         
-        # Remove duplicates while preserving order
-        all_intermediate_jsons = list(dict.fromkeys(all_intermediate_jsons))
-
-        # 6. Prepare Final Output for Display
-        if result is None or (isinstance(result, list) and len(result) == 0):
-            display_text = "⚠️ Pipeline failed at Step 3 (BIO Decoding).\nDownload the intermediate JSONs below to inspect OCR and Model Predictions."
-        else:
-            display_text = json.dumps(result, indent=2, ensure_ascii=False)
+# #         if img is None:
+# #             return None
+        
+# #         # Run YOLO detection
+# #         model = YOLO(WEIGHTS_PATH)
+# #         results = model.predict(source=img, conf=0.2, imgsz=640, verbose=False)
+        
+# #         # Draw bounding boxes
+# #         detection_count = {'figure': 0, 'equation': 0}
+# #         for box in results[0].boxes:
+# #             class_id = int(box.cls[0])
+# #             class_name = model.names[class_id]
+# #             if class_name in ['figure', 'equation']:
+# #                 x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
+# #                 conf = float(box.conf[0])
+                
+# #                 # Green for figures, Red for equations
+# #                 color = (0, 255, 0) if class_name == 'figure' else (0, 0, 255)
+# #                 cv2.rectangle(img, (x1, y1), (x2, y2), color, 3)
+                
+# #                 # Add label with confidence
+# #                 label = f"{class_name.upper()} {conf:.2f}"
+# #                 (text_width, text_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
+# #                 cv2.rectangle(img, (x1, y1 - text_height - 10), (x1 + text_width, y1), color, -1)
+# #                 cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
+                
+# #                 detection_count[class_name] += 1
+        
+# #         # Add summary text at top
+# #         summary = f"Detected: {detection_count['figure']} Figures (GREEN), {detection_count['equation']} Equations (RED)"
+# #         cv2.rectangle(img, (10, 10), (10 + len(summary) * 10, 40), (0, 0, 0), -1)
+# #         cv2.putText(img, summary, (15, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
+        
+# #         # Save to temp file
+# #         temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".png").name
+# #         cv2.imwrite(temp_path, img)
+# #         return temp_path
+        
+# #     except Exception as e:
+# #         print(f"Error in visualize_detections: {e}")
+# #         import traceback
+# #         traceback.print_exc()
+# #         return None
+
+
+# # # ==============================
+# # # GRADIO INTERFACE
+# # # ==============================
+# # with gr.Blocks(title="Document Analysis Pipeline") as demo:
+
+# #     gr.Markdown("# 📄 Full Pipeline Analysis")
+# #     gr.Markdown("### 🔍 Intermediate File Recovery Active")
+# #     gr.Markdown("The **Download** box will contain: \n1. OCR JSON (Step 1)\n2. Raw LayoutLMv3 Prediction JSON (Step 2)\n3. Final BIO JSON (Step 3)")
+
+# #     with gr.Row():
+# #         with gr.Column(scale=1):
+# #             file_input = gr.File(
+# #                 label="Upload PDFs or Images",
+# #                 file_types=[".pdf", ".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"],
+# #                 file_count="multiple",
+# #                 type="filepath"
+# #             )
+
+# #             page_selector = gr.Slider(
+# #                 minimum=1,
+# #                 maximum=100,
+# #                 value=1,
+# #                 step=1,
+# #                 label="PDF Page Number (for preview)",
+# #                 visible=True
+# #             )
+
+# #             model_path_input = gr.Textbox(
+# #                 label="Model Path",
+# #                 value=DEFAULT_LAYOUTLMV3_MODEL_PATH
+# #             )
+
+# #             # Debug button for visual inspection
+# #             debug_btn = gr.Button("🔍 Show YOLO Detections (First Page)", variant="secondary")
             
-            # If the final result succeeded, save it to a temp file so it can be downloaded too
-            temp_final = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json', prefix='final_result_')
-            json.dump(result, temp_final, indent=2, ensure_ascii=False)
-            temp_final.close()
-            all_intermediate_jsons.append(temp_final.name)
-
-        return display_text, all_intermediate_jsons
-
-    except Exception as e:
-        import traceback
-        traceback.print_exc()
-        return f"❌ Error: {str(e)}", None
-
-
-
-
-
-# def visualize_detections(uploaded_files):
-#     """Shows the first uploaded image with YOLO bounding boxes"""
-#     if not uploaded_files:
-#         return None
-    
-#     # Get first file path
-#     file_path = uploaded_files[0] if isinstance(uploaded_files, list) else uploaded_files
-#     if isinstance(file_path, dict):
-#         file_path = file_path["path"]
-    
-#     import cv2
-#     from ultralytics import YOLO
-    
-#     # Load image
-#     img = cv2.imread(str(file_path))
-#     if img is None:
-#         return None
-    
-#     # Run YOLO
-#     model = YOLO(WEIGHTS_PATH)
-#     results = model.predict(source=img, conf=0.2, imgsz=640, verbose=False)
-    
-#     # Draw boxes
-#     for box in results[0].boxes:
-#         class_id = int(box.cls[0])
-#         class_name = model.names[class_id]
-#         if class_name in ['figure', 'equation']:
-#             x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
-#             color = (0, 255, 0) if class_name == 'figure' else (255, 0, 0)
-#             cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
-#             cv2.putText(img, f"{class_name} {box.conf[0]:.2f}", 
-#                        (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
-    
-#     # Save and return
-#     temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".png").name
-#     cv2.imwrite(temp_path, img)
-#     return temp_path
-
+# #             # Main processing button
+# #             process_btn = gr.Button("🚀 Run Full Pipeline", variant="primary")
 
-# # ==============================
-# # GRADIO INTERFACE
-# # ==============================
-# with gr.Blocks(title="Document Analysis Pipeline") as demo:
-
-#     gr.Markdown("# 📄 Full Pipeline Analysis")
-#     gr.Markdown("### 🔍 Intermediate File Recovery Active")
-#     gr.Markdown("The **Download** box will contain: \n1. OCR JSON (Step 1)\n2. Raw LayoutLMv3 Prediction JSON (Step 2)\n3. Final BIO JSON (Step 3)")
-
-#     with gr.Row():
-#         with gr.Column(scale=1):
-#             file_input = gr.File(
-#                 label="Upload PDFs or Images",
-#                 file_types=[".pdf", ".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"],
-#                 file_count="multiple",
-#                 type="filepath"
-#             )
+# #         with gr.Column(scale=2):
+# #             # Visual debug output
+# #             detection_preview = gr.Image(label="YOLO Detection Preview (Green=Figure, Red=Equation)", type="filepath")
+            
+# #             # Final JSON output
+# #             json_output = gr.Code(label="Final Structured Output", language="json", lines=20)
+            
+# #             # Download all intermediate files
+# #             download_output = gr.File(label="Download All Pipeline Stages (JSON)", file_count="multiple")
 
-#             model_path_input = gr.Textbox(
-#                 label="Model Path",
-#                 value=DEFAULT_LAYOUTLMV3_MODEL_PATH
-#             )
+# #     # Wire up the debug button
+# #     debug_btn.click(
+# #         fn=visualize_detections,
+# #         inputs=[file_input, page_selector],
+# #         outputs=[detection_preview]
+# #     )
 
-#             process_btn = gr.Button("🚀 Run Pipeline", variant="primary")
+# #     # Wire up the main processing button
+# #     process_btn.click(
+# #         fn=process_file,
+# #         inputs=[file_input, model_path_input],
+# #         outputs=[json_output, download_output]
+# #     )
 
-#         with gr.Column(scale=2):
-#             json_output = gr.Code(label="Final Structured Output", language="json", lines=20)
-#             # IMPORTANT: file_count="multiple" allows returning the list of all stage files
-#             download_output = gr.File(label="Download All Pipeline Stages (JSON)", file_count="multiple")
 
-#     process_btn.click(
-#         fn=process_file,
-#         inputs=[file_input, model_path_input],
-#         outputs=[json_output, download_output]
-#     )
 
-# if __name__ == "__main__":
-#     demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
+# # if __name__ == "__main__":
+# #     demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
 
 
 
@@ -354,12 +369,18 @@ def process_file(uploaded_files, layoutlmv3_model_path=None):
 # # ==============================
 # # VISUAL DEBUG FUNCTION
 # # ==============================
-# def visualize_detections(uploaded_files):
-#     """Shows the first uploaded image with YOLO bounding boxes"""
+# def visualize_detections(uploaded_files, page_num):
+#     """Shows the selected PDF page or image with YOLO bounding boxes"""
 #     if not uploaded_files:
 #         return None
     
 #     try:
+#         import cv2
+#         import numpy as np
+#         import tempfile
+#         from ultralytics import YOLO
+#         import fitz
+        
 #         # Get first file path
 #         file_path = uploaded_files[0] if isinstance(uploaded_files, list) else uploaded_files
 #         if isinstance(file_path, dict):
@@ -367,16 +388,11 @@ def process_file(uploaded_files, layoutlmv3_model_path=None):
 #         elif hasattr(file_path, 'path'):
 #             file_path = file_path.path
         
-#         import cv2
-#         import numpy as np
-
-#         from ultralytics import YOLO
-#         import fitz
-        
 #         # Handle PDF conversion to image
 #         if str(file_path).lower().endswith('.pdf'):
 #             doc = fitz.open(file_path)
-#             page_idx = int(page_num) - 1
+#             # Ensure the selected page exists in the document
+#             page_idx = min(max(int(page_num) - 1, 0), len(doc) - 1)
 #             page = doc.load_page(page_idx) 
 
 #             pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
@@ -418,8 +434,8 @@ def process_file(uploaded_files, layoutlmv3_model_path=None):
 #                 detection_count[class_name] += 1
         
 #         # Add summary text at top
-#         summary = f"Detected: {detection_count['figure']} Figures (GREEN), {detection_count['equation']} Equations (RED)"
-#         cv2.rectangle(img, (10, 10), (10 + len(summary) * 10, 40), (0, 0, 0), -1)
+#         summary = f"Page {page_num} | Detected: {detection_count['figure']} Figures, {detection_count['equation']} Equations"
+#         cv2.rectangle(img, (10, 10), (10 + len(summary) * 11, 40), (0, 0, 0), -1)
 #         cv2.putText(img, summary, (15, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
         
 #         # Save to temp file
@@ -433,7 +449,6 @@ def process_file(uploaded_files, layoutlmv3_model_path=None):
 #         traceback.print_exc()
 #         return None
 
-
 # # ==============================
 # # GRADIO INTERFACE
 # # ==============================
@@ -467,7 +482,7 @@ def process_file(uploaded_files, layoutlmv3_model_path=None):
 #             )
 
 #             # Debug button for visual inspection
-#             debug_btn = gr.Button("🔍 Show YOLO Detections (First Page)", variant="secondary")
+#             debug_btn = gr.Button("🔍 Show YOLO Detections", variant="secondary")
             
 #             # Main processing button
 #             process_btn = gr.Button("🚀 Run Full Pipeline", variant="primary")
@@ -496,8 +511,6 @@ def process_file(uploaded_files, layoutlmv3_model_path=None):
 #         outputs=[json_output, download_output]
 #     )
 
-
-
 # if __name__ == "__main__":
 #     demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
 
@@ -506,6 +519,172 @@ def process_file(uploaded_files, layoutlmv3_model_path=None):
 
 
 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+import gradio as gr
+import json
+import os
+import tempfile
+import img2pdf
+import glob
+import shutil
+from img2pdf import Rotation
+from pathlib import Path
+
+print("--- DEBUG: Current Working Directory ---")
+print(os.getcwd())
+print("--- DEBUG: Files in Root ---")
+print(os.listdir('.'))
+
+# ==============================
+# PIPELINE IMPORT
+# ==============================
+try:
+    from working_yolo_pipeline import run_document_pipeline, DEFAULT_LAYOUTLMV3_MODEL_PATH, WEIGHTS_PATH
+except Exception as e:  # Catch ALL exceptions
+    print(f"Warning: Failed to import pipeline: {e}")
+    import traceback
+    traceback.print_exc()  # Show the actual error
+    def run_document_pipeline(*args):
+        yield {"status": "error", "message": "Placeholder pipeline function called."}
+    DEFAULT_LAYOUTLMV3_MODEL_PATH = "./models/layoutlmv3_model"
+    WEIGHTS_PATH = "./weights/yolo_weights.pt"
+
+
+# ==============================
+# MAIN PROCESSING GENERATOR
+# ==============================
+def process_file(uploaded_files, layoutlmv3_model_path=None):
+    """
+    Robust handler for multiple or single file uploads.
+    Streams the estimation first, then yields the final JSON and intermediate files.
+    """
+    if uploaded_files is None:
+        yield "❌ Error: No files uploaded.", None
+        return
+
+    if not isinstance(uploaded_files, list):
+        file_list = [uploaded_files]
+    else:
+        file_list = uploaded_files
+
+    if len(file_list) == 0:
+        yield "❌ Error: Empty file list.", None
+        return
+
+    # 1. Resolve all file paths safely
+    resolved_paths = []
+    for f in file_list:
+        try:
+            if isinstance(f, dict) and "path" in f:
+                resolved_paths.append(f["path"])
+            elif hasattr(f, 'path'):
+                resolved_paths.append(f.path)
+            else:
+                resolved_paths.append(str(f))
+        except Exception as e:
+            print(f"Error resolving path for {f}: {e}")
+
+    if not resolved_paths:
+        yield "❌ Error: Could not resolve file paths.", None
+        return
+
+    # 2. Determine if we should merge into a single PDF
+    first_file = Path(resolved_paths[0])
+    is_image = first_file.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp', '.webp', '.tiff']
+    
+    try:
+        if len(resolved_paths) > 1 or is_image:
+            print(f"📦 Converting {len(resolved_paths)} image(s) to a single PDF...")
+            temp_pdf = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
+            with open(temp_pdf.name, "wb") as f_out:
+                f_out.write(img2pdf.convert(resolved_paths, rotation=Rotation.ifvalid))
+            processing_path = temp_pdf.name
+        else:
+            processing_path = resolved_paths[0]
+
+        # 3. Standard Pipeline Checks
+        final_model_path = layoutlmv3_model_path or DEFAULT_LAYOUTLMV3_MODEL_PATH
+        if not os.path.exists(final_model_path):
+            yield f"❌ Error: Model not found at {final_model_path}", None
+            return
+
+        # 4. Call the pipeline generator
+        print(f"🚀 Starting pipeline for: {processing_path}")
+        
+        # Iterate through the yields from run_document_pipeline
+        for pipeline_update in run_document_pipeline(processing_path, final_model_path):
+            
+            # --- Handle Estimation Yield ---
+            if pipeline_update.get("status") == "estimating":
+                display_text = "⏱️ ESTIMATING PROCESSING TIME...\n\n" + json.dumps(pipeline_update, indent=2)
+                yield display_text, None 
+            
+            # --- Handle Final Complete Yield ---
+            elif pipeline_update.get("status") == "complete":
+                final_result = pipeline_update.get("result")
+                
+                # SCRAPE FOR INTERMEDIATE FILES
+                base_name = Path(processing_path).stem
+                search_patterns = [
+                    f"/tmp/pipeline_run_{base_name}*/*.json",
+                    f"/tmp/*{base_name}*.json"
+                ]
+                
+                all_intermediate_jsons = []
+                for pattern in search_patterns:
+                    all_intermediate_jsons.extend(glob.glob(pattern))
+                
+                all_intermediate_jsons = list(dict.fromkeys(all_intermediate_jsons))
+
+                # Prepare Final Output for Display
+                if final_result is None or (isinstance(final_result, list) and len(final_result) == 0):
+                    display_text = "⚠️ Pipeline failed at Step 3 (BIO Decoding).\nDownload the intermediate JSONs below to inspect OCR and Model Predictions."
+                else:
+                    display_text = json.dumps(final_result, indent=2, ensure_ascii=False)
+                    
+                    # Save it to a temp file so it can be downloaded too
+                    temp_final = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json', prefix='final_result_')
+                    json.dump(final_result, temp_final, indent=2, ensure_ascii=False)
+                    temp_final.close()
+                    all_intermediate_jsons.append(temp_final.name)
+
+                yield display_text, all_intermediate_jsons
+                
+            # --- Handle Error Yield ---
+            elif pipeline_update.get("status") == "error":
+                yield f"❌ Error: {pipeline_update.get('message')}", None
+
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        yield f"❌ Error: {str(e)}", None
+
+
 # ==============================
 # VISUAL DEBUG FUNCTION
 # ==============================
@@ -631,8 +810,8 @@ with gr.Blocks(title="Document Analysis Pipeline") as demo:
             # Visual debug output
             detection_preview = gr.Image(label="YOLO Detection Preview (Green=Figure, Red=Equation)", type="filepath")
             
-            # Final JSON output
-            json_output = gr.Code(label="Final Structured Output", language="json", lines=20)
+            # Final JSON output (Will update with estimation, then final result)
+            json_output = gr.Code(label="Pipeline Output", language="json", lines=20)
             
             # Download all intermediate files
             download_output = gr.File(label="Download All Pipeline Stages (JSON)", file_count="multiple")
@@ -648,9 +827,11 @@ with gr.Blocks(title="Document Analysis Pipeline") as demo:
     process_btn.click(
         fn=process_file,
         inputs=[file_input, model_path_input],
-        outputs=[json_output, download_output]
+        outputs=[json_output, download_output],
+        api_name="process"  # This enables the streaming endpoint /api/process
     )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
+    # IMPORTANT: .queue() is required for streaming generators to work!
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True)