Update app.py

app.py

@@ -303,7 +303,6 @@
 
 
 
-
 import fitz  # PyMuPDF
 import numpy as np
 import cv2
@@ -316,7 +315,6 @@ import logging
 import gradio as gr
 import shutil
 import tempfile
-from PIL import Image
 import io
 
 # ============================================================================
@@ -351,11 +349,9 @@ GLOBAL_FIGURE_COUNT = 0
 GLOBAL_EQUATION_COUNT = 0
 
 # ============================================================================
-# --- BOX COMBINATION LOGIC (Retained) ---
+# --- BOX COMBINATION LOGIC ---
 # ============================================================================
 
-# (calculate_iou, filter_nested_boxes, merge_overlapping_boxes functions remain unchanged)
-
 def calculate_iou(box1, box2):
     x1_a, y1_a, x2_a, y2_a = box1
     x1_b, y1_b, x2_b, y2_b = box2
@@ -423,7 +419,7 @@ def merge_overlapping_boxes(detections, iou_threshold):
     return merged_detections
 
 # ============================================================================
-# --- UTILITY FUNCTIONS (Modified to capture coordinates) ---
+# --- UTILITY FUNCTIONS ---
 # ============================================================================
 
 def pixmap_to_numpy(pix: fitz.Pixmap) -> np.ndarray:
@@ -443,7 +439,7 @@ def run_yolo_detection_and_count(
 ) -> Tuple[int, int, List[Dict[str, Any]]]:
     """
     Runs YOLO inference, applies NMS/filtering, and updates global counters.
-    Returns counts AND a list of equation detection results.
+    Returns counts AND a list of equation detection results (PDF coordinates).
     """
     global GLOBAL_FIGURE_COUNT, GLOBAL_EQUATION_COUNT
     
@@ -452,8 +448,6 @@ def run_yolo_detection_and_count(
     page_figures = 0
     
     try:
-        # Run prediction
-        # Setting device to 'cpu' is a safety measure if CUDA isn't available
         results = model.predict(image, conf=CONF_THRESHOLD, verbose=False)
         
         if results and results[0].boxes:
@@ -471,7 +465,7 @@ def run_yolo_detection_and_count(
         logging.error(f"YOLO inference failed on page {page_num}: {e}")
         return 0, 0, []
 
-    # Apply NMS/Merging/Filtering based on your provided logic
+    # Apply NMS/Merging/Filtering
     merged_detections = merge_overlapping_boxes(yolo_detections, IOU_MERGE_THRESHOLD)
     final_detections = filter_nested_boxes(merged_detections, IOA_SUPPRESSION_THRESHOLD)
 
@@ -502,16 +496,16 @@ def run_yolo_detection_and_count(
 
 
 # ============================================================================
-# --- MAIN DOCUMENT PROCESSING FUNCTION (Modified to handle image cropping) ---
+# --- MAIN DOCUMENT PROCESSING FUNCTION ---
 # ============================================================================
 
-def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, List[str]]:
+def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, List[bytes]]:
     """
-    Runs the pipeline, returns counts, report, and a list of paths/bytes for
-    the cropped equation images.
+    Runs the pipeline, returns counts, report, and a list of cropped equation images (as bytes).
     """
     global GLOBAL_FIGURE_COUNT, GLOBAL_EQUATION_COUNT
 
+    # Reset globals
     GLOBAL_FIGURE_COUNT = 0
     GLOBAL_EQUATION_COUNT = 0
 
@@ -519,21 +513,25 @@ def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, Lis
         report = f"❌ FATAL ERROR: Input PDF not found at {pdf_path}."
         return 0, 0, 0, report, []
     
+    # Model Loading
     try:
         model = YOLO(WEIGHTS_PATH)
+        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.)"
         return 0, 0, 0, report, []
 
+    # PDF Loading
     try:
         doc = fitz.open(pdf_path)
         total_pages = doc.page_count
+        logging.warning(f"✅ Opened PDF with {doc.page_count} pages")
     except Exception as e:
         report = f"❌ ERROR loading PDF file: {e}"
         return 0, 0, 0, report, []
 
     mat = fitz.Matrix(SCALE_FACTOR, SCALE_FACTOR)
-    all_equation_images = [] # List to store cropped image data (base64 or bytes)
+    all_equation_images = [] # Stores PNG image bytes for Gradio gallery
     
     for page_num_0_based in range(doc.page_count):
         fitz_page = doc.load_page(page_num_0_based)
@@ -553,25 +551,24 @@ def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, Lis
         
         # --- Image Cropping and Saving for Debugging ---
         for eq in equation_results_page:
-            # bbox_pdf is in PyMuPDF's Rect format (x0, y0, x1, y1)
             bbox = eq['bbox_pdf'] 
             
             try:
-                # Crop the equation using the bounding box on the fitz page
-                # We use a slight border (e.g., 5 points) for better visualization
-                rect = fitz.Rect(bbox).prerotate(fitz_page.rotation)
+                # FIX APPLIED: Removed .prerotate() call on Rect object
+                rect = fitz.Rect(bbox) 
                 clip_rect = rect + (0, 0, 5, 5) # Add small padding
                 
-                # Get the pixmap for the cropped area
+                # 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 a format Gradio can display (PNG bytes)
+                # Convert the pixmap to PNG bytes
                 img_bytes = eq_pix.tobytes("png")
                 
                 all_equation_images.append(img_bytes)
                 
             except Exception as e:
-                logging.error(f"Error cropping equation on page {page_num}: {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()
 
@@ -583,15 +580,14 @@ def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, Lis
         f"**3) Total Figures Detected:** **{GLOBAL_FIGURE_COUNT}**"
     )
 
-    # Note the return type change to include the list of image bytes
     return total_pages, GLOBAL_EQUATION_COUNT, GLOBAL_FIGURE_COUNT, report, all_equation_images
 
 
 # ============================================================================
-# --- GRADIO INTERFACE FUNCTION (Modified for image output) ---
+# --- GRADIO INTERFACE FUNCTION ---
 # ============================================================================
 
-def gradio_process_pdf(pdf_file) -> Tuple[str, str, str, str, List[str]]:
+def gradio_process_pdf(pdf_file) -> Tuple[str, str, str, str, List[bytes]]:
     """
     Gradio wrapper function to handle file upload and return all results + images.
     """
@@ -629,7 +625,8 @@ if __name__ == "__main__":
     output_equations = gr.Textbox(label="Total Equations Detected", interactive=False)
     output_figures = gr.Textbox(label="Total Figures Detected", interactive=False)
     output_report = gr.Markdown(label="Processing Summary")
-    # NEW: Gradio Gallery to display the list of cropped images
+    
+    # Gradio Gallery to display the list of cropped images
     output_gallery = gr.Gallery(
         label="Detected Equations for Debugging", 
         columns=5, 
@@ -644,8 +641,8 @@ if __name__ == "__main__":
         outputs=[output_pages, output_equations, output_figures, output_report, output_gallery],
         title="🎯 Minimalist YOLO Counting & Equation Debugger",
         description=(
-            "Upload a PDF to run YOLO detection. The counts are displayed, and a gallery "
-            "of **all detected equation images** is shown below for debugging the detection accuracy."
+            "Upload a PDF to run YOLO detection using your **`best.pt`** model. "
+            "The counts are displayed, and a gallery of **all detected equation images** is shown for debugging."
         ),
     )