Update app.py

app.py

@@ -1,12 +1,11 @@
 import gradio as gr
 from ultralytics import YOLO
-import cv2
 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_path = 'runs/detect/train7/weights/best.pt'  # Replace with the path to your trained .pt file
 model = YOLO(model_path)
 
 # Define the class indices for figures and tables
@@ -30,6 +29,7 @@ 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
@@ -43,11 +43,12 @@ 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, low_res_pix in enumerate(low_res_pixmaps):
+    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)
         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
@@ -55,7 +56,7 @@ def process_pdf(pdf_file):
         
         if boxes:
             # Load high DPI image for cropping only if boxes are found
-            high_res_pix = doc[page_num].get_pixmap(dpi=high_dpi)
+            high_res_pix = page.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