Update app.py

app.py

@@ -3,9 +3,6 @@ from ultralytics import YOLO
 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
 model_path = 'best.pt'  # Replace with the path to your trained .pt file
@@ -33,32 +30,10 @@ 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_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
-    boxes = infer_image_and_get_boxes(low_res_img)
-    
-    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_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)
-    
-    # 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
@@ -71,17 +46,20 @@ 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 threading
-    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 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:
+    # 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:
-            cropped_imgs = process_high_res_page(page_num, boxes, scale_factor, doc_path)
+            # 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