Update app.py

app.py

@@ -352,6 +352,19 @@
 
 
 
+
+
+
+
+
+
+
+
+
+
+
+
+
 import fitz  # PyMuPDF
 import numpy as np
 import cv2
@@ -532,19 +545,19 @@ def run_yolo_detection_and_count(
 
 
 # ============================================================================
-# --- MAIN DOCUMENT PROCESSING FUNCTION (Revised for Dict Return) ---
+# --- MAIN DOCUMENT PROCESSING FUNCTION (Fixed for JSON serialization) ---
 # ============================================================================
 
-# NOTE: The return signature now includes the equation_counts_per_page dictionary
-def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, float, Dict[int, int], List[str]]:
+# NOTE: The return signature now uses Dict[str, int] for the equation counts
+def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, float, Dict[str, int], List[str]]:
     """
-    Runs the pipeline, returns counts, report, total time, page counts dict, and empty list.
+    Runs the pipeline, returns counts, report, total time, page counts dict (str keys), and empty list.
     """
     global GLOBAL_FIGURE_COUNT, GLOBAL_EQUATION_COUNT
     start_time = time.time()
     log_messages = []
     
-    # NEW: Dictionary to store {page_number: equation_count}
+    # Dictionary to store {page_number (int): equation_count (int)}
     equation_counts_per_page: Dict[int, int] = {} 
 
     # Reset globals
@@ -555,7 +568,6 @@ def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, flo
     t0 = time.time()
     if not os.path.exists(pdf_path):
         report = f"❌ FATAL ERROR: Input PDF not found at {pdf_path}."
-        # Add the new return value (empty dict)
         return 0, 0, 0, report, time.time() - start_time, {}, []
     
     try:
@@ -563,7 +575,6 @@ def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, flo
         logging.warning(f"✅ Loaded YOLO model from: {WEIGHTS_PATH}")
     except Exception as e:
         report = f"❌ ERROR loading YOLO model: {e}\n(Ensure 'best.pt' is available and valid.)"
-        # Add the new return value (empty dict)
         return 0, 0, 0, report, time.time() - start_time, {}, []
     t1 = time.time()
     log_messages.append(f"Model Loading Time: {t1-t0:.4f}s")
@@ -576,7 +587,6 @@ def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, flo
         logging.warning(f"✅ Opened PDF with {doc.page_count} pages")
     except Exception as e:
         report = f"❌ ERROR loading PDF file: {e}"
-        # Add the new return value (empty dict)
         return 0, 0, 0, report, time.time() - start_time, {}, []
     t3 = time.time()
     log_messages.append(f"PDF Initialization Time: {t3-t2:.4f}s")
@@ -605,7 +615,7 @@ def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, flo
         page_equations, _ = run_yolo_detection_and_count(original_img, model, page_num)
         detect_time = time.time() - detect_start
         
-        # NEW: Store the count in the dictionary
+        # Store the count in the dictionary (INT keys)
         equation_counts_per_page[page_num] = page_equations
         
         page_total_time = time.time() - page_start_time
@@ -616,6 +626,11 @@ def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, flo
     detection_loop_time = t5 - t4
     log_messages.append(f"Total Detection Loop Time ({total_pages} pages): {detection_loop_time:.4f}s")
 
+    # FIX APPLIED HERE: Convert integer keys to string keys for JSON serialization
+    equation_counts_per_page_str_keys: Dict[str, int] = {
+        str(k): v for k, v in equation_counts_per_page.items()
+    }
+
     # 4. Final Report Generation
     total_execution_time = t5 - start_time
     
@@ -632,38 +647,37 @@ def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, flo
         f"\n```"
     )
 
-    # Return the new dictionary as the sixth element, and an empty list as the seventh
-    return total_pages, GLOBAL_EQUATION_COUNT, GLOBAL_FIGURE_COUNT, report, total_execution_time, equation_counts_per_page, []
+    # Return the dictionary with string keys
+    return total_pages, GLOBAL_EQUATION_COUNT, GLOBAL_FIGURE_COUNT, report, total_execution_time, equation_counts_per_page_str_keys, []
 
 
 # ============================================================================
 # --- GRADIO INTERFACE FUNCTION (Updated) ---
 # ============================================================================
 
-# NOTE: The return signature now includes the equation_counts_per_page dictionary (Dict[int, int])
-def gradio_process_pdf(pdf_file) -> Tuple[str, str, str, str, Dict[int, int], List[str]]:
+def gradio_process_pdf(pdf_file) -> Tuple[str, str, str, str, Dict[str, int], List[str]]:
     """
     Gradio wrapper function to handle file upload and return results.
     """
     if pdf_file is None:
-        # Return an empty dict for the new JSON output
+        # Return an empty dict with string keys
         return "N/A", "N/A", "N/A", "Please upload a PDF file.", {}, []
     
     pdf_path = pdf_file.name
 
     try:
-        # Unpack the new return value: equation_counts_per_page
+        # Unpack the new return value: equation_counts_per_page (with string keys)
         num_pages, num_equations, num_figures, report, total_time, equation_counts_per_page, _ = run_single_pdf_preprocessing(
             pdf_path
         )
         
-        # Return results (5 items now: 3 textboxes, 1 markdown, 1 JSON, 1 empty list for gallery)
+        # Return results (6 items now)
         return str(num_pages), str(num_equations), str(num_figures), report, equation_counts_per_page, []
             
     except Exception as e:
         error_msg = f"An unexpected error occurred: {e}"
         logging.error(error_msg, exc_info=True)
-        # Return an empty dict for the new JSON output on error
+        # Return an empty dict on error
         return "Error", "Error", "Error", error_msg, {}, []
 
 
@@ -699,13 +713,13 @@ if __name__ == "__main__":
     interface = gr.Interface(
         fn=gradio_process_pdf,
         inputs=input_file,
-        # Outputs now include the JSON component
+        # Outputs list remains the same, but the JSON component now receives string keys.
         outputs=[
             output_pages, 
             output_equations, 
             output_figures, 
             output_report, 
-            output_page_counts, # New
+            output_page_counts, 
             output_gallery
         ],
         title="📊 YOLO Counting with Per-Page Data & Timing",
@@ -716,4 +730,6 @@ if __name__ == "__main__":
     )
 
     print("\nStarting Gradio application...")
-    interface.launch(inbrowser=True)
+    interface.launch(inbrowser=True)
+
+