Update app.py

app.py

@@ -3,7 +3,6 @@ from ultralytics import YOLO
 import numpy as np
 import fitz  # PyMuPDF
 import spaces
-
 # Load the trained model
 model_path = 'best.pt'  # Replace with the path to your trained .pt file
 model = YOLO(model_path)
@@ -43,12 +42,11 @@ def process_pdf(pdf_file):
     # Calculate the scaling factor
     scale_factor = high_dpi / low_dpi
 
+    # Pre-cache all page pixmaps at low DPI
+    low_res_pixmaps = [page.get_pixmap(dpi=low_dpi) for page in doc]
+    
     # Loop through each page
-    for page_num in range(len(doc)):
-        page = doc.load_page(page_num)
-        
-        # Perform inference at low DPI
-        low_res_pix = page.get_pixmap(dpi=low_dpi)
+    for page_num, low_res_pix in enumerate(low_res_pixmaps):
         low_res_img = np.frombuffer(low_res_pix.samples, dtype=np.uint8).reshape(low_res_pix.height, low_res_pix.width, 3)
         
         # Get bounding boxes from low DPI image
@@ -56,7 +54,7 @@ def process_pdf(pdf_file):
         
         if boxes:
             # Load high DPI image for cropping only if boxes are found
-            high_res_pix = page.get_pixmap(dpi=high_dpi)
+            high_res_pix = doc[page_num].get_pixmap(dpi=high_dpi)
             high_res_img = np.frombuffer(high_res_pix.samples, dtype=np.uint8).reshape(high_res_pix.height, high_res_pix.width, 3)
             
             # Crop images at high DPI