Spaces:
Runtime error
Runtime error
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) | |