Update app.py

app.py

@@ -31,7 +31,7 @@ def crop_images_from_boxes(image, boxes, scale_factor):
     return cropped_images
 
 @spaces.GPU
-def process_pdf(pdf_file):
+def process_pdf(pdf_file, chunk_size=10):
     # Open the PDF file
     doc = fitz.open(pdf_file)
     all_cropped_images = []
@@ -43,24 +43,29 @@ def process_pdf(pdf_file):
     # Calculate the scaling factor
     scale_factor = high_dpi / low_dpi
 
-    # Pre-cache all page pixmaps at low DPI
-    low_res_pixmaps = [page.get_pixmap(dpi=low_dpi) for page in doc]
-    
-    # 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)
+    # Process the PDF in chunks
+    num_pages = len(doc)
+    for chunk_start in range(0, num_pages, chunk_size):
+        chunk_end = min(chunk_start + chunk_size, num_pages)
         
-        # Get bounding boxes from low DPI image
-        boxes = infer_image_and_get_boxes(low_res_img)
+        # Pre-cache page pixmaps at low DPI for the current chunk
+        low_res_pixmaps = [doc.load_page(page_num).get_pixmap(dpi=low_dpi) for page_num in range(chunk_start, chunk_end)]
         
-        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)
+        # Process each page in the current chunk
+        for page_num, low_res_pix in enumerate(low_res_pixmaps, start=chunk_start):
+            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)
             
-            # Crop images at high DPI
-            cropped_imgs = crop_images_from_boxes(high_res_img, boxes, scale_factor)
-            all_cropped_images.extend(cropped_imgs)
+            if boxes:
+                # Load high DPI image for cropping only if boxes are found
+                high_res_pix = doc.load_page(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
 
@@ -75,3 +80,4 @@ iface = gr.Interface(
 
 # Launch the app
 iface.launch()
+