Update app.py

app.py

@@ -60,7 +60,7 @@ def stitch_images_in_grid(images: List[Image.Image], num_columns: int, page_orde
             
     return grid_image
 
-def process_pdf(pdf_file, pdf_url, dpi, num_columns, crop_top, crop_bottom, crop_left, crop_right, hide_annotations, page_order, progress=gr.Progress()):
+def process_pdf(pdf_file, pdf_url, dpi, num_columns, num_images, crop_top, crop_bottom, crop_left, crop_right, hide_annotations, page_order, progress=gr.Progress()):
     pdf_input_source = None
     is_bytes = False
     source_name = "document"
@@ -114,26 +114,43 @@ def process_pdf(pdf_file, pdf_url, dpi, num_columns, crop_top, crop_bottom, crop
     else:
         cropped_images = images
 
-    progress(0.7, desc=f"Stitching {len(cropped_images)} images together...")
+    progress(0.7, desc=f"Splitting {len(cropped_images)} pages into {num_images} image(s)...")
+
+    total_pages = len(cropped_images)
+    effective_num_images = min(num_images, total_pages)
     
-    if num_columns > 1:
-        stitched_image = stitch_images_in_grid(cropped_images, num_columns, page_order)
-    else:
-        stitched_image = stitch_images_vertically(cropped_images)
-        
-    if stitched_image is None:
-        raise gr.Error("Image stitching failed.")
-    logger.info("Image stitching complete.")
-
-    progress(0.9, desc="Saving final image...")
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".png", prefix=f"{source_name}_stitched_") as tmp_file:
-        stitched_image.save(tmp_file.name, "PNG")
-        output_path = tmp_file.name
+    chunk_size = math.ceil(total_pages / effective_num_images)
+    image_chunks = [cropped_images[i:i + chunk_size] for i in range(0, total_pages, chunk_size)]
+
+    output_paths = []
+    num_chunks = len(image_chunks)
+
+    for i, chunk in enumerate(image_chunks):
+        progress(0.75 + (0.2 * (i / num_chunks)), desc=f"Stitching image {i+1} of {num_chunks}...")
         
-    logger.info(f"Final image saved to temporary path: {output_path}")
+        if not chunk:
+            continue
+
+        if num_columns > 1:
+            stitched_image = stitch_images_in_grid(chunk, num_columns, page_order)
+        else:
+            stitched_image = stitch_images_vertically(chunk)
+            
+        if stitched_image is None:
+            logger.warning(f"Image stitching failed for chunk {i+1}.")
+            continue
+
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".png", prefix=f"{source_name}_stitched_{i+1}_") as tmp_file:
+            stitched_image.save(tmp_file.name, "PNG")
+            output_paths.append(tmp_file.name)
+    
+    if not output_paths:
+        raise gr.Error("Image stitching failed for all pages.")
+
+    logger.info(f"Final images saved to temporary paths: {output_paths}")
     progress(1, desc="Done!")
     
-    return output_path, output_path
+    return output_paths, output_paths
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
@@ -153,6 +170,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             
             dpi_slider = gr.Slider(minimum=100, maximum=600, step=5, value=200, label="Image Resolution (DPI)")
             columns_slider = gr.Slider(minimum=1, maximum=10, step=1, value=1, label="Number of Columns")
+            num_images_slider = gr.Slider(minimum=1, maximum=10, step=1, value=1, label="Number of Output Images", info="Splits the PDF pages into multiple output images.")
 
             with gr.Accordion("Advanced Options", open=False):
                 hide_annotations_toggle = gr.Checkbox(value=True, label="Hide PDF Annotations (Links/Highlights)", info="Turn this on to remove the colored boxes that can appear around links and references.")
@@ -168,8 +186,8 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
         with gr.Column(scale=2):
             gr.Markdown("## Output")
-            output_image_preview = gr.Image(label="Stitched Image Preview", type="filepath", interactive=False, height=600)
-            output_image_download = gr.File(label="Download Stitched Image", interactive=False)
+            output_image_preview = gr.Gallery(label="Stitched Image(s) Preview", interactive=False, height=600)
+            output_image_download = gr.File(label="Download Stitched Image(s)", interactive=False, file_count="multiple")
 
     submit_btn.click(
         fn=process_pdf,
@@ -177,7 +195,8 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             pdf_file_input, 
             pdf_url_input, 
             dpi_slider, 
-            columns_slider, 
+            columns_slider,
+            num_images_slider,
             crop_top, 
             crop_bottom, 
             crop_left,