Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-import fitz  # PyMuPDF
+from pdf2image import convert_from_path
 from transformers import AutoModel, AutoTokenizer
 from PIL import Image
 import numpy as np
@@ -31,29 +31,33 @@ def image_to_base64(image):
     image.save(buffered, format="PNG")
     return base64.b64encode(buffered.getvalue()).decode()
 
-def pdf_to_images(pdf_path):
-    images = []
-    pdf_document = fitz.open(pdf_path)
-    for page_num in range(len(pdf_document)):
-        page = pdf_document.load_page(page_num)
-        pix = page.get_pixmap()
-        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-        images.append(img)
-    return images
+def convert_pdf_to_images(pdf_path, output_folder):
+    # Ensure the output folder exists
+    if not os.path.exists(output_folder):
+        os.makedirs(output_folder)
+
+    # Convert PDF to images
+    images = convert_from_path(pdf_path)
+
+    # Save each image to the output folder
+    image_paths = []
+    for i, image in enumerate(images):
+        image_path = os.path.join(output_folder, f"page_{i + 1}.png")
+        image.save(image_path, 'JPEG')
+        image_paths.append(image_path)
+        print(f"Saved {image_path}")
+    return image_paths
 
 def run_GOT(pdf_file):
     unique_id = str(uuid.uuid4())
     pdf_path = os.path.join(UPLOAD_FOLDER, f"{unique_id}.pdf")
     shutil.copy(pdf_file, pdf_path)
     
-    images = pdf_to_images(pdf_path)
+    images = convert_pdf_to_images(pdf_path, UPLOAD_FOLDER)
     results = []
     
     try:
-        for i, image in enumerate(images):
-            image_path = os.path.join(UPLOAD_FOLDER, f"{unique_id}_page_{i+1}.jpg")
-            image.save(image_path)
-            
+        for i, image_path in enumerate(images):
             result_path = os.path.join(RESULTS_FOLDER, f"{unique_id}_page_{i+1}.html")
             
             res = model.chat_crop(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)