Update app.py

app.py

@@ -3,7 +3,8 @@ from ultralytics import YOLO
 import numpy as np
 import fitz  # PyMuPDF
 import spaces
-from multiprocessing import Pool, cpu_count
+from concurrent.futures import ThreadPoolExecutor
+import cv2
 
 # Load the trained model
 model_path = 'best.pt'  # Replace with the path to your trained .pt file
@@ -31,31 +32,29 @@ 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)
+# Function to process a single page's low-resolution image and perform inference
+def process_low_res_page(page_num, low_res_pix, scale_factor, doc):
     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 page_num, boxes
+
+# Function to process a single page's high-resolution image for cropping
+def process_high_res_page(page_num, boxes, scale_factor, doc):
+    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
@@ -68,14 +67,18 @@ 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]
     
-    # 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)]
+    # Prepare arguments for threading
+    args_low_res = [(page_num, low_res_pix, scale_factor, doc) 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)
+    # Process low-res pages concurrently using threading to get bounding boxes
+    with ThreadPoolExecutor(max_workers=cpu_count()) as executor:
+        low_res_results = list(executor.map(lambda p: process_low_res_page(*p), args_low_res))
+    
+    # Sequentially process high-res pages to crop images
+    for page_num, boxes in low_res_results:
+        if boxes:
+            cropped_imgs = process_high_res_page(page_num, boxes, scale_factor, doc)
+            all_cropped_images.extend(cropped_imgs)
     
     return all_cropped_images