Update app.py

app.py

@@ -1,19 +1,22 @@
 import gradio as gr
 import numpy as np
-from pdf2image import convert_from_path
-from PIL import Image
-from ultralytics import YOLOv10
+import fitz  # PyMuPDF
 import spaces
+from ultralytics import YOLOv10
+
 # Load the trained model
+
 model = YOLOv10("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
 
 # Function to perform inference on an image and return bounding boxes for figures and tables
+
 def infer_image_and_get_boxes(image, confidence_threshold=0.6):
-    results = model.predict(np.array(image))
+    results = model.predict(image)
     boxes = [
         (int(box.xyxy[0][0]), int(box.xyxy[0][1]), int(box.xyxy[0][2]), int(box.xyxy[0][3]))
         for result in results for box in result.boxes
@@ -22,39 +25,46 @@ def infer_image_and_get_boxes(image, confidence_threshold=0.6):
     return boxes
 
 # Function to crop images from the boxes
+
 def crop_images_from_boxes(image, boxes, scale_factor):
     cropped_images = [
-        image.crop((int(x1 * scale_factor), int(y1 * scale_factor), int(x2 * scale_factor), int(y2 * scale_factor)))
+        image[int(y1 * scale_factor):int(y2 * scale_factor), int(x1 * scale_factor):int(x2 * scale_factor)]
         for (x1, y1, x2, y2) in boxes
     ]
     return cropped_images
 
 @spaces.GPU
 def process_pdf(pdf_file):
+    # Open the PDF file
+    doc = fitz.open(pdf_file)
     all_cropped_images = []
 
     # Set the DPI for inference and high resolution for cropping
     low_dpi = 50
     high_dpi = 300
 
-    # Convert PDF pages to images at low DPI
-    low_res_images = convert_from_path(pdf_file.name, dpi=low_dpi)
-    
     # Calculate the scaling factor
     scale_factor = high_dpi / low_dpi
 
-    for page_num, low_res_img in enumerate(low_res_images):
+    # 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):
+        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:
-            # Convert the specific page to high DPI only if boxes are found
-            high_res_img = convert_from_path(pdf_file.name, dpi=high_dpi, first_page=page_num+1, last_page=page_num+1)[0]
+            # Load high DPI image for cropping only if boxes are found
+            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
             cropped_imgs = crop_images_from_boxes(high_res_img, boxes, scale_factor)
             all_cropped_images.extend(cropped_imgs)
-
+    
     return all_cropped_images
 
 # Create Gradio interface
@@ -69,3 +79,4 @@ iface = gr.Interface(
 # Launch the app
 iface.launch()
 
+