import gradio as gr
import pymupdf
import os
import zipfile

def pdf_to_images(pdf_file):

    doc = pymupdf.open(pdf_file)
    images = []
    
    for page_id in range(doc.page_count):
        page = doc[page_id]
        pix = page.get_pixmap(dpi=200)
        img_bytes = pix.tobytes("png")
        images.append((img_bytes, f"{page_id+1}.png"))


    doc.close()


    temp_dir = "temp_images"
    os.makedirs(temp_dir, exist_ok=True)


    for img_bytes, img_name in images:
        with open(os.path.join(temp_dir, img_name), "wb") as f:
            f.write(img_bytes)


    zip_path = "images.zip"
    with zipfile.ZipFile(zip_path, "w") as zipf:
        for img_name in os.listdir(temp_dir):
            zipf.write(os.path.join(temp_dir, img_name), img_name)

  
    for img_name in os.listdir(temp_dir):
        os.remove(os.path.join(temp_dir, img_name))
    os.rmdir(temp_dir)

    return zip_path


iface = gr.Interface(
    fn=pdf_to_images,
    inputs=gr.File(label="Upload PDF File"),
    outputs=gr.File(label="Download ZIP File"),
    title="PDF to Images Converter",
    description="Upload a PDF file and download a ZIP file containing all the pages as images. Host it on huggingface for convenience."
)
iface.launch()

# Source code
# https://juejin.cn/post/7382480523846467595