File size: 4,617 Bytes
95abc0b 64cd544 95abc0b 64cd544 95abc0b 64cd544 95abc0b 64cd544 95abc0b 64cd544 95abc0b 64cd544 95abc0b 64cd544 95abc0b 64cd544 95abc0b 64cd544 95abc0b 64cd544 95abc0b 64cd544 95abc0b 64cd544 95abc0b 64cd544 95abc0b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
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
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 progress.tqdm(pdf_files, desc="Converting PDFs"):
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,
progress=gr.Progress(),
):
if not pdf_files:
return None, None, "No PDF files uploaded."
if oauth_token is None:
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:
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:
if "temp_dir" in locals():
shutil.rmtree(temp_dir)
return None, 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)", 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.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, download_button, status_text],
)
# Launch the app
demo.launch(debug=True)
|