Update app.py

app.py

@@ -1,11 +1,10 @@
 import gradio as gr
 from ultralytics import YOLO
 import cv2
-import spaces
 import numpy as np
 import fitz  # PyMuPDF
 from PIL import Image
-
+import spaces
 # Load the trained model
 model_path = 'best.pt'  # Replace with the path to your trained .pt file
 model = YOLO(model_path)
@@ -31,7 +30,6 @@ def crop_images_from_boxes(image, boxes, scale_factor):
         for (x1, y1, x2, y2) in boxes
     ]
     return cropped_images
-
 @spaces.GPU
 def process_pdf(pdf_file):
     # Open the PDF file
@@ -45,12 +43,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
@@ -58,7 +55,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