Update app.py

app.py

@@ -351,15 +351,176 @@ 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"""
+#     if not uploaded_files:
+#         return 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
+
+#         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
+#             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))
+        
+#         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")
+            
+#             # Main processing button
+#             process_btn = gr.Button("🚀 Run Full Pipeline", variant="primary")
+
+#         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")
+
+#     # Wire up the debug button
+#     debug_btn.click(
+#         fn=visualize_detections,
+#         inputs=[file_input, page_selector],
+#         outputs=[detection_preview]
+#     )
+
+#     # Wire up the main processing button
+#     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)
+
+
+
+
+
+
 # ==============================
 # 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 +528,11 @@ def visualize_detections(uploaded_files):
         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 +574,8 @@ def visualize_detections(uploaded_files):
                 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 +589,6 @@ def visualize_detections(uploaded_files):
         traceback.print_exc()
         return None
 
-
 # ==============================
 # GRADIO INTERFACE
 # ==============================
@@ -467,7 +622,7 @@ with gr.Blocks(title="Document Analysis Pipeline") as demo:
             )
 
             # 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 +651,6 @@ with gr.Blocks(title="Document Analysis Pipeline") as demo:
         outputs=[json_output, download_output]
     )
 
-
-
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)