File size: 3,910 Bytes
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
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)