davanstrien's picture
davanstrien HF staff
draft
95abc0b
raw
history blame
3.91 kB
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)