Spaces:
Sleeping
Sleeping
| 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) | |