Update app.py

app.py

@@ -3,6 +3,7 @@ from ultralytics import YOLO
 import numpy as np
 import fitz  # PyMuPDF
 import spaces
+from multiprocessing import Pool, cpu_count
 
 # Load the trained model
 model_path = 'best.pt'  # Replace with the path to your trained .pt file
@@ -30,10 +31,31 @@ def crop_images_from_boxes(image, boxes, scale_factor):
     ]
     return cropped_images
 
+# Function to process a single page
+def process_page(args):
+    page_num, low_res_pix, scale_factor, doc_path = args
+    doc = fitz.open(doc_path)
+    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)
+    
+    cropped_imgs = []
+    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)
+        
+        # Crop images at high DPI
+        cropped_imgs = crop_images_from_boxes(high_res_img, boxes, scale_factor)
+        
+    return cropped_imgs
+
 @spaces.GPU
 def process_pdf(pdf_file):
     # Open the PDF file
     doc = fitz.open(pdf_file)
+    doc_path = pdf_file.name
     all_cropped_images = []
 
     # Set the DPI for inference and high resolution for cropping
@@ -46,21 +68,14 @@ def process_pdf(pdf_file):
     # 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:
-            # 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)
+    # Prepare arguments for multiprocessing
+    args = [(page_num, low_res_pix, scale_factor, doc_path) for page_num, low_res_pix in enumerate(low_res_pixmaps)]
+
+    # Process pages concurrently using multiprocessing
+    with Pool(cpu_count()) as pool:
+        results = pool.map(process_page, args)
+        for result in results:
+            all_cropped_images.extend(result)
     
     return all_cropped_images