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, chunk_size=10):
+def process_pdf(pdf_file):
     # Open the PDF file
     doc = fitz.open(pdf_file)
     all_cropped_images = []
@@ -43,29 +43,24 @@ def process_pdf(pdf_file, chunk_size=10):
     # Calculate the scaling factor
     scale_factor = high_dpi / low_dpi
 
-    # 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)
+    # 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)
         
-        # 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)]
+        # Get bounding boxes from low DPI image
+        boxes = infer_image_and_get_boxes(low_res_img)
         
-        # 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)
+        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)
             
-            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)
+            # 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