reqs and app

app.py

@@ -1,18 +1,21 @@
-import gradio as gr
 import os
 import random
+import shutil
 import tempfile
+import zipfile
+
+import gradio as gr
+from huggingface_hub import HfApi
 from pdf2image import convert_from_path
 from PyPDF2 import PdfReader
-from huggingface_hub import create_repo, upload_folder, HfApi
 
 
-def pdf_to_images(pdf_files, sample_size, temp_dir):
+def pdf_to_images(pdf_files, sample_size, temp_dir, progress=gr.Progress()):
     if not os.path.exists(temp_dir):
         os.makedirs(temp_dir)
-
+    progress(0, desc="Starting conversion")
     all_images = []
-    for pdf_file in pdf_files:
+    for pdf_file in progress.tqdm(pdf_files, desc="Converting PDFs"):
         pdf_path = pdf_file.name
         pdf = PdfReader(pdf_path)
         total_pages = len(pdf.pages)
@@ -45,42 +48,60 @@ def pdf_to_images(pdf_files, sample_size, temp_dir):
     return all_images, f"Saved {len(all_images)} images to temporary directory"
 
 
-def process_pdfs(pdf_files, sample_size, hf_repo, oauth_token: gr.OAuthToken | None):
+def process_pdfs(
+    pdf_files,
+    sample_size,
+    hf_repo,
+    oauth_token: gr.OAuthToken | None,
+    progress=gr.Progress(),
+):
     if not pdf_files:
-        return None, "No PDF files uploaded."
+        return None, None, "No PDF files uploaded."
 
     if oauth_token is None:
-        return None, "Please log in to upload to Hugging Face."
+        gr.Info("Please log in to upload to Hugging Face.")
+        return (
+            None,
+            None,
+            "Not logged in to Hugging Face, please log in to upload to a Hugging Face dataset.",
+        )
 
     try:
-        with tempfile.TemporaryDirectory() as temp_dir:
-            images_dir = os.path.join(temp_dir, "images")
-            os.makedirs(images_dir)
-
-            images, message = pdf_to_images(pdf_files, sample_size, images_dir)
-
-            if hf_repo:
-                try:
-                    hf_api = HfApi(token=oauth_token.token)
-                    hf_api.create_repo(
-                        hf_repo,
-                        repo_type="dataset",
-                    )
-                    hf_api.upload_folder(
-                        folder_path=images_dir,
-                        repo_id=hf_repo,
-                        repo_type="dataset",
-                        path_in_repo="images",
-                    )
-                    message += (
-                        f"\nUploaded images to Hugging Face repo: {hf_repo}/images"
-                    )
-                except Exception as e:
-                    message += f"\nFailed to upload to Hugging Face: {str(e)}"
-
-            return images, message
+        temp_dir = tempfile.mkdtemp()
+        images_dir = os.path.join(temp_dir, "images")
+        os.makedirs(images_dir)
+
+        images, message = pdf_to_images(pdf_files, sample_size, images_dir)
+
+        # Create a zip file of the images
+        zip_path = os.path.join(temp_dir, "converted_images.zip")
+        with zipfile.ZipFile(zip_path, "w") as zipf:
+            progress(0, desc="Zipping images")
+            for image in progress.tqdm(images, desc="Zipping images"):
+                zipf.write(image, os.path.basename(image))
+
+        if hf_repo:
+            try:
+                hf_api = HfApi(token=oauth_token.token)
+                hf_api.create_repo(
+                    hf_repo,
+                    repo_type="dataset",
+                )
+                hf_api.upload_folder(
+                    folder_path=images_dir,
+                    repo_id=hf_repo,
+                    repo_type="dataset",
+                    path_in_repo="images",
+                )
+                message += f"\nUploaded images to Hugging Face repo: {hf_repo}/images"
+            except Exception as e:
+                message += f"\nFailed to upload to Hugging Face: {str(e)}"
+
+        return images, zip_path, message
     except Exception as e:
-        return None, f"An error occurred: {str(e)}"
+        if "temp_dir" in locals():
+            shutil.rmtree(temp_dir)
+        return None, None, f"An error occurred: {str(e)}"
 
 
 # Define the Gradio interface
@@ -94,26 +115,27 @@ with gr.Blocks() as demo:
         gr.LoginButton(size="sm")
 
     with gr.Row():
-        pdf_files = gr.File(file_count="multiple", label="Upload PDF(s)")
-        sample_size = gr.Slider(
-            minimum=0,
-            maximum=50,
-            step=1,
-            value=0,
-            label="Sample Size (0 for all pages)",
+        pdf_files = gr.File(
+            file_count="multiple", label="Upload PDF(s)", file_types=["*.pdf"]
+        )
+    with gr.Row():
+        sample_size = gr.Number(
+            value=None,
+            label="Number of sample pages (0 will return all pages)",
         )
         hf_repo = gr.Textbox(
             label="Hugging Face Repo", placeholder="username/repo-name"
         )
 
     output_gallery = gr.Gallery(label="Converted Images")
-    status_text = gr.Textbox(label="Status")
+    status_text = gr.Markdown(label="Status")
+    download_button = gr.File(label="Download Converted Images")
 
     submit_button = gr.Button("Process PDFs")
     submit_button.click(
         process_pdfs,
         inputs=[pdf_files, sample_size, hf_repo],
-        outputs=[output_gallery, status_text],
+        outputs=[output_gallery, download_button, status_text],
     )
 
 # Launch the app

requirements.in

@@ -1,5 +1,6 @@
-gradio[oauth]==4.44.0
-PyPDF2
-pdf2image
+
 Pillow
-huggingface_hub
+PyPDF2
+gradio[oauth]==4.44.0
+huggingface_hub
+pdf2image