Update app.py

app.py

@@ -1,48 +1,24 @@
 import gradio as gr
-import fitz
-import os
-import zipfile
-from donut import DonutModel
+import fitz  # PyMuPDF
+from PIL import Image
 
-def demo_process(input_pdf):
-    
-    # Conversion of PDF to JPG images
-    pdf = fitz.open(input_pdf)
-    first_page = pdf[0]
-    pix = first_page.get_pixmap()
-    image_bytes = pix.tobytes("png")
-    pdf.close()
-
-    temp_dir = "images"
-    basename = os.path.basename(input_pdf).split('.')[0]
-    image_name = basename + ".jpg"
-    os.makedirs(temp_dir, exist_ok=True)
-    
-    with open(os.path.join(temp_dir, image_name), "wb") as f:
-        f.write(image_bytes)
-
-    # zip_path = "images.zip"
-    image_path = os.path.join(temp_dir, image_name)
-    # output = model.inference(image=image_name, prompt=task_prompt)["predictions"][0]    
-    # with zipfile.ZipFile(zip_path, "w") as zipf:
-    #     zipf.write(os.path.join(temp_dir, image_name), image_name)    
-    # os.remove(image_path)
-    # os.rmdir(temp_dir)
-    image_path = Image.fromarray(image_path)
-    
+def pdf_to_jpg(pdf_file):
+    # Open the PDF file
+    document = fitz.open(pdf_file.name)
+    # Get the first page
+    page = document.load_page(0)
+    # Render the page to a pixmap
+    pix = page.get_pixmap()
+    # Save the pixmap as an image
+    image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+    image_path = "first_page.jpg"
+    image.save(image_path)
     return image_path
-    
 
-task_name = "SGSInvoice"
-task_prompt = f"<s_{task_name}>"
-model = DonutModel.from_pretrained("uartimcs/donut-invoice-extract")
-model.eval()
-
-iface = gr.Interface(
-    fn=demo_process,
-    inputs=gr.File(label="Upload PDF File"),
-    outputs="image",
-    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."
+interface = gr.Interface(
+    fn=pdf_to_jpg,
+    inputs=gr.inputs.File(label="Upload PDF"),
+    outputs=gr.outputs.Image(label="First Page JPG")
 )
-iface.launch()
+
+interface.launch()