Update app.py

app.py

@@ -2,6 +2,7 @@ import os
 import tempfile
 import logging
 from typing import List
+import math
 
 import gradio as gr
 import requests
@@ -31,10 +32,35 @@ def stitch_images_vertically(images: List[Image.Image]) -> Image.Image:
         
     return stitched_image
 
-def process_pdf(pdf_file, pdf_url, dpi, progress=gr.Progress()):
+def stitch_images_in_grid(images: List[Image.Image], num_columns: int) -> Image.Image:
     """
-    Processes a PDF from a file upload or a URL, converts it to images with a specified DPI,
-    and stitches them together vertically.
+    Stitches images into a grid with a specified number of columns.
+    """
+    if not images:
+        return None
+
+    num_images = len(images)
+    num_rows = math.ceil(num_images / num_columns)
+    
+    columns = [images[i::num_columns] for i in range(num_columns)]
+    
+    stitched_columns = [stitch_images_vertically(col) for col in columns]
+    
+    max_height = max(col.height for col in stitched_columns if col)
+    total_width = sum(col.width for col in stitched_columns if col)
+    
+    grid_image = Image.new('RGB', (total_width, max_height), (255, 255, 255))
+    
+    current_x = 0
+    for col_img in stitched_columns:
+        if col_img:
+            grid_image.paste(col_img, (current_x, 0))
+            current_x += col_img.width
+            
+    return grid_image
+
+    """
+    Processes a PDF, converts it to images, and stitches them into a single image or a grid.
     """
     pdf_input_source = None
     is_bytes = False
@@ -79,7 +105,11 @@ def process_pdf(pdf_file, pdf_url, dpi, progress=gr.Progress()):
 
     progress(0.7, desc=f"Stitching {len(images)} images together...")
     
-    stitched_image = stitch_images_vertically(images)
+    if num_columns > 1:
+        stitched_image = stitch_images_in_grid(images, num_columns)
+    else:
+        stitched_image = stitch_images_vertically(images)
+        
     if stitched_image is None:
         raise gr.Error("Image stitching failed.")
         
@@ -101,7 +131,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         """
         # PDF Page Stitcher 📄 ➡️ 🖼️
         Upload a PDF file or provide a URL. This tool will convert every page of the PDF into an image 
-        and then append them beneath each other to create a single, tall image that you can download.
+        and then append them to create a single image that you can download.
         """
     )
 
@@ -115,7 +145,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                         label="PDF URL",
                         placeholder="e.g., https://arxiv.org/pdf/1706.03762.pdf"
                     )
-
+            
             dpi_slider = gr.Slider(
                 minimum=100,
                 maximum=600,
@@ -124,6 +154,15 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 label="Image Resolution (DPI)",
                 info="Higher DPI results in a clearer image but increases processing time and file size."
             )
+
+            columns_slider = gr.Slider(
+                minimum=1,
+                maximum=10,
+                step=1,
+                value=1,
+                label="Number of Columns",
+                info="Increase to make the final image wider and less tall."
+            )
             
             submit_btn = gr.Button("Stitch PDF Pages", variant="primary")
 
@@ -142,7 +181,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     submit_btn.click(
         fn=process_pdf,
-        inputs=[pdf_file_input, pdf_url_input, dpi_slider],
+        inputs=[pdf_file_input, pdf_url_input, dpi_slider, columns_slider],
         outputs=[output_image_preview, output_image_download]
     )