Update app.py

app.py

@@ -1,25 +1,51 @@
 import gradio as gr
-import argparse
-import torch
-from donut import DonutModel
-from PIL import Image
-from poppler import load_from_file, PageRenderer
-
-def demo_process(pdf_file):
-    pdf_document = load_from_file(pdf_file.name)
-    page_1 = pdf_document.create_page(0)
-
-    renderer = PageRenderer()
-    image = renderer.render_page(page_1)
-    image_data = image.data
-    input_img = Image.fromarray(image_data)
-    output = model.inference(image=input_img, prompt=task_prompt)["predictions"][0]
-    return output
-    
-model = DonutModel.from_pretrained("uartimcs/donut-invoice-extract")
-model.eval()
-task_name = "SGSInvoice"
-task_prompt = f"<s_{task_name}>"
-
-demo = gr.Interface(fn=demo_process,inputs=gr.File(file_count='single', file_types=['.pdf']),outputs="json", title=f"Donut 🍩 demonstration for `{task_name}` task",)
-demo.launch()
+import fitz
+import os
+import zipfile
+
+def pdf_to_images(pdf_file):
+
+    doc = fitz.open(pdf_file)
+    images = []
+
+    for page_id in range(doc.page_count):
+        page = doc[page_id]
+
+        pix = page.get_pixmap()
+        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
+
+def main():
+    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."
+    )
+    iface.launch()