Update app.py

app.py

@@ -1,22 +1,19 @@
 import gradio as gr
 import numpy as np
-import fitz  # PyMuPDF
-import spaces
+from pdf2image import convert_from_path
+from PIL import Image
 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(image)
+    results = model.predict(np.array(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
@@ -25,46 +22,39 @@ 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[int(y1 * scale_factor):int(y2 * scale_factor), int(x1 * scale_factor):int(x2 * scale_factor)]
+        image.crop((int(x1 * scale_factor), int(y1 * scale_factor), int(x2 * scale_factor), int(y2 * 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
 
-    # 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)
-        
+    for page_num, low_res_img in enumerate(low_res_images):
         # 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 = 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)
+            # 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]
             
             # 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
@@ -78,3 +68,4 @@ iface = gr.Interface(
 
 # Launch the app
 iface.launch()
+