Update app.py

app.py

@@ -5,9 +5,8 @@ import cv2
 import numpy as np
 import fitz  # PyMuPDF
 from PIL import Image
-from concurrent.futures import ThreadPoolExecutor, as_completed
 import spaces
-# Load the trained model globally
+# Load the trained model
 model_path = 'best.pt'  # Replace with the path to your trained .pt file
 model = YOLO(model_path)
 
@@ -44,9 +43,23 @@ def crop_images_from_boxes(image, boxes, scale_factor):
         cropped_images.append(cropped_image)
     return cropped_images
 
-# Function to process a single page
-def process_single_page(page, low_dpi, high_dpi, scale_factor):
-    try:
+@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)
+        
         # 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)
@@ -62,32 +75,8 @@ def process_single_page(page, low_dpi, high_dpi, scale_factor):
 
         # Crop images at high DPI
         cropped_imgs = crop_images_from_boxes(high_res_img, boxes, scale_factor)
-        
-        return cropped_imgs
-    except Exception as e:
-        print(f"Error processing page: {e}")
-        return []
-
-@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
+        all_cropped_images.extend(cropped_imgs)
     
-    # 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
@@ -101,3 +90,4 @@ iface = gr.Interface(
 
 # Launch the app
 iface.launch()
+