Update app.py

app.py

@@ -5,8 +5,9 @@ import cv2
 import numpy as np
 import fitz  # PyMuPDF
 from PIL import Image
-import spaces
-# Load the trained model
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+# Load the trained model globally
 model_path = 'best.pt'  # Replace with the path to your trained .pt file
 model = YOLO(model_path)
 
@@ -43,23 +44,9 @@ def crop_images_from_boxes(image, boxes, scale_factor):
         cropped_images.append(cropped_image)
     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
-    
-    # Calculate the scaling factor
-    scale_factor = high_dpi / low_dpi
-    
-    # Loop through each page
-    for page_num in range(len(doc)):
-        page = doc.load_page(page_num)
-        
+# Function to process a single page
+def process_single_page(page, low_dpi, high_dpi, scale_factor):
+    try:
         # Perform inference at low DPI
         low_res_pix = page.get_pixmap(dpi=low_dpi)
         low_res_img = Image.frombytes("RGB", [low_res_pix.width, low_res_pix.height], low_res_pix.samples)
@@ -75,8 +62,31 @@ def process_pdf(pdf_file):
 
         # Crop images at high DPI
         cropped_imgs = crop_images_from_boxes(high_res_img, boxes, scale_factor)
-        all_cropped_images.extend(cropped_imgs)
+        
+        return cropped_imgs
+    except Exception as e:
+        print(f"Error processing page: {e}")
+        return []
+
+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
     
+    # Calculate the scaling factor
+    scale_factor = high_dpi / low_dpi
+    
+    # Use ThreadPoolExecutor for batch processing
+    with ThreadPoolExecutor() as executor:
+        futures = [executor.submit(process_single_page, doc.load_page(page_num), low_dpi, high_dpi, scale_factor) for page_num in range(len(doc))]
+        
+        for future in as_completed(futures):
+            all_cropped_images.extend(future.result())
+
     return all_cropped_images
 
 # Create Gradio interface