Update app.py

app.py

@@ -301,8 +301,6 @@
 
 
 
-
-
 import fitz  # PyMuPDF
 import numpy as np
 import cv2
@@ -328,7 +326,6 @@ def patched_torch_load(*args, **kwargs):
     return _original_torch_load(*args, **kwargs)
 torch.load = patched_torch_load
 
-# Suppress warnings during normal operation
 logging.basicConfig(level=logging.WARNING)
 
 # ============================================================================
@@ -496,12 +493,12 @@ def run_yolo_detection_and_count(
 
 
 # ============================================================================
-# --- MAIN DOCUMENT PROCESSING FUNCTION ---
+# --- MAIN DOCUMENT PROCESSING FUNCTION (Fixed for Gradio Gallery) ---
 # ============================================================================
 
-def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, List[bytes]]:
+def run_single_pdf_preprocessing(pdf_path: str, temp_output_dir: str) -> Tuple[int, int, int, str, List[str]]:
     """
-    Runs the pipeline, returns counts, report, and a list of cropped equation images (as bytes).
+    Runs the pipeline, returns counts, report, and a list of paths to cropped equation images.
     """
     global GLOBAL_FIGURE_COUNT, GLOBAL_EQUATION_COUNT
 
@@ -531,7 +528,7 @@ def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, Lis
         return 0, 0, 0, report, []
 
     mat = fitz.Matrix(SCALE_FACTOR, SCALE_FACTOR)
-    all_equation_images = [] # Stores PNG image bytes for Gradio gallery
+    all_equation_images = [] # Stores file paths (strings) for Gradio gallery
     
     for page_num_0_based in range(doc.page_count):
         fitz_page = doc.load_page(page_num_0_based)
@@ -554,20 +551,24 @@ def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, Lis
             bbox = eq['bbox_pdf'] 
             
             try:
-                # FIX APPLIED: Removed .prerotate() call on Rect object
+                # FIX 1 (PyMuPDF Rect): Only instantiate Rect, no .prerotate()
                 rect = fitz.Rect(bbox) 
                 clip_rect = rect + (0, 0, 5, 5) # Add small padding
                 
                 # Get the pixmap for the cropped area (high-res render)
                 eq_pix = fitz_page.get_pixmap(matrix=fitz.Matrix(3.0, 3.0), clip=clip_rect)
                 
-                # Convert the pixmap to PNG bytes
+                # FIX 2 (Gradio Gallery): Save to a temporary file path
                 img_bytes = eq_pix.tobytes("png")
+                filename = f"eq_{GLOBAL_EQUATION_COUNT}_p{page_num}.png"
+                output_path = os.path.join(temp_output_dir, filename)
+
+                with open(output_path, 'wb') as f:
+                    f.write(img_bytes)
                 
-                all_equation_images.append(img_bytes)
+                all_equation_images.append(output_path)
                 
             except Exception as e:
-                # This error means the image generation failed for a specific box.
                 logging.error(f"Error cropping equation on page {page_num} with bbox {bbox}: {e}")
 
     doc.close()
@@ -580,33 +581,40 @@ def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, Lis
         f"**3) Total Figures Detected:** **{GLOBAL_FIGURE_COUNT}**"
     )
 
+    # Return the list of file paths (strings)
     return total_pages, GLOBAL_EQUATION_COUNT, GLOBAL_FIGURE_COUNT, report, all_equation_images
 
 
 # ============================================================================
-# --- GRADIO INTERFACE FUNCTION ---
+# --- GRADIO INTERFACE FUNCTION (Fixed for temp dir management) ---
 # ============================================================================
 
-def gradio_process_pdf(pdf_file) -> Tuple[str, str, str, str, List[bytes]]:
+def gradio_process_pdf(pdf_file) -> Tuple[str, str, str, str, List[str]]:
     """
-    Gradio wrapper function to handle file upload and return all results + images.
+    Gradio wrapper function to handle file upload, manage temporary directory, and return file paths.
     """
     if pdf_file is None:
         return "N/A", "N/A", "N/A", "Please upload a PDF file.", []
     
     pdf_path = pdf_file.name
+    temp_output_dir = tempfile.mkdtemp() # Create temp directory
 
     try:
-        # Run the core logic
-        num_pages, num_equations, num_figures, report, equation_images = run_single_pdf_preprocessing(pdf_path)
+        # Run the core logic, passing the temp directory
+        num_pages, num_equations, num_figures, report, equation_images = run_single_pdf_preprocessing(
+            pdf_path, temp_output_dir
+        )
         
-        # Return results as formatted strings and the list of image bytes
+        # Return results and the list of image file paths
         return str(num_pages), str(num_equations), str(num_figures), report, equation_images
             
     except Exception as e:
         error_msg = f"An unexpected error occurred: {e}"
         logging.error(error_msg, exc_info=True)
         return "Error", "Error", "Error", error_msg, []
+    finally:
+        # CRUCIAL: Clean up the temporary directory containing the image files
+        shutil.rmtree(temp_output_dir, ignore_errors=True)
 
 
 # ============================================================================
@@ -626,7 +634,7 @@ if __name__ == "__main__":
     output_figures = gr.Textbox(label="Total Figures Detected", interactive=False)
     output_report = gr.Markdown(label="Processing Summary")
     
-    # Gradio Gallery to display the list of cropped images
+    # Gradio Gallery expects a list of file paths (strings)
     output_gallery = gr.Gallery(
         label="Detected Equations for Debugging", 
         columns=5,