Update app.py

app.py

@@ -4,9 +4,14 @@ import numpy as np
 import fitz  # PyMuPDF
 import spaces
 
-# Load the trained model with FP16 precision
+
+# Load the trained model
+
+
+
 model = YOLO("best.pt")
 
+
 # Define the class indices for figures and tables
 figure_class_index = 3  # class index for figures
 table_class_index = 4   # class index for tables
@@ -42,18 +47,19 @@ 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 in doc:
-        # Get low-resolution pixmap
-        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
         boxes = infer_image_and_get_boxes(low_res_img)
         
         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