draft

app.py

@@ -0,0 +1,120 @@
+import gradio as gr
+import os
+import random
+import tempfile
+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):
+    if not os.path.exists(temp_dir):
+        os.makedirs(temp_dir)
+
+    all_images = []
+    for pdf_file in pdf_files:
+        pdf_path = pdf_file.name
+        pdf = PdfReader(pdf_path)
+        total_pages = len(pdf.pages)
+
+        # Determine the number of pages to convert
+        pages_to_convert = (
+            total_pages if sample_size == 0 else min(sample_size, total_pages)
+        )
+
+        # Select random pages if sampling
+        if sample_size > 0 and sample_size < total_pages:
+            selected_pages = sorted(
+                random.sample(range(1, total_pages + 1), pages_to_convert)
+            )
+        else:
+            selected_pages = range(1, total_pages + 1)
+
+        # Convert selected PDF pages to images
+        for page_num in selected_pages:
+            images = convert_from_path(
+                pdf_path, first_page=page_num, last_page=page_num
+            )
+            for image in images:
+                image_path = os.path.join(
+                    temp_dir, f"{os.path.basename(pdf_path)}_page_{page_num}.jpg"
+                )
+                image.save(image_path, "JPEG")
+                all_images.append(image_path)
+
+    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):
+    if not pdf_files:
+        return None, "No PDF files uploaded."
+
+    if oauth_token is None:
+        return None, "Please log in to upload to Hugging Face."
+
+    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:
+                    api = HfApi(token=oauth_token.token)
+                    api.create_repo(
+                        hf_repo,
+                        repo_type="dataset",
+                    )
+                    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
+    except Exception as e:
+        return None, f"An error occurred: {str(e)}"
+
+
+# Define the Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# PDF to Image Converter")
+    gr.Markdown(
+        "Upload PDF(s), convert pages to images, and optionally upload them to a Hugging Face repo. If a sample size is specified, random pages will be selected."
+    )
+
+    with gr.Row():
+        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)",
+        )
+        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")
+
+    submit_button = gr.Button("Process PDFs")
+    submit_button.click(
+        process_pdfs,
+        inputs=[pdf_files, sample_size, hf_repo],
+        outputs=[output_gallery, status_text],
+    )
+
+# Launch the app
+demo.launch(debug=True)