Update app.py

app.py

@@ -4,6 +4,7 @@ import numpy as np
 import fitz  # PyMuPDF
 import spaces
 from concurrent.futures import ThreadPoolExecutor
+from multiprocessing import cpu_count
 import cv2
 
 # Load the trained model
@@ -33,7 +34,8 @@ def crop_images_from_boxes(image, boxes, scale_factor):
     return cropped_images
 
 # 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):
+def process_low_res_page(page_num, low_res_pix, scale_factor, doc_path):
+    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
@@ -42,7 +44,8 @@ def process_low_res_page(page_num, low_res_pix, scale_factor, doc):
     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):
+def process_high_res_page(page_num, boxes, scale_factor, doc_path):
+    doc = fitz.open(doc_path)
     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)
     
@@ -55,6 +58,7 @@ def process_high_res_page(page_num, boxes, scale_factor, doc):
 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,16 +72,16 @@ def process_pdf(pdf_file):
     low_res_pixmaps = [page.get_pixmap(dpi=low_dpi) for page in doc]
     
     # 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)]
+    args_low_res = [(page_num, low_res_pix, scale_factor, doc_path) for page_num, low_res_pix in enumerate(low_res_pixmaps)]
 
     # 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))
+        low_res_results = list(executor.map(lambda args: process_low_res_page(*args), 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)
+            cropped_imgs = process_high_res_page(page_num, boxes, scale_factor, doc_path)
             all_cropped_images.extend(cropped_imgs)
     
     return all_cropped_images