huge

app.py

@@ -2,6 +2,8 @@ import spaces
 import gradio as gr
 import surya.detection as detection
 import surya.layout as layout
+import os
+import base64
 
 
 # Monkey patch to prevent spawning processes
@@ -46,27 +48,39 @@ layout.batch_layout_detection = batch_layout_detection
 from marker.convert import convert_single_pdf
 from marker.models import load_all_models
 
-model_lst = load_all_models()
+model_list = load_all_models()
 
 
 @spaces.GPU
-def convert(file_path):
-    print(file_path)
-    global model_lst
-
-    full_text, images, meta = convert_single_pdf(
-        file_path,
-        model_lst,
-        max_pages=None,
-        langs=None,
-        batch_multiplier=16,
-    )
+def convert(pdf_file, extract_images):
+    global model_list
+
+    full_text, images, out_meta = convert_single_pdf(pdf_file, model_list)
+    image_data = {}
+    if extract_images:
+        for filename, image in images.items():
+            image.save(filename, "PNG")
+
+            with open(filename, "rb") as f:
+                image_bytes = f.read()
+
+            image_base64 = base64.b64encode(image_bytes).decode("utf-8")
+            image_data[filename] = image_base64
+
+            os.remove(filename)
 
-    return full_text
+    return full_text, out_meta, image_data
 
 
 gr.Interface(
     convert,
-    inputs=gr.File(label="PDF file", type="filepath"),
-    outputs=gr.Markdown(label="Markdown"),
+    inputs=[
+        gr.File(label="Upload PDF", type="filepath"),
+        gr.Checkbox(label="Extract Images"),
+    ],
+    outputs=[
+        gr.Text(label="Markdown"),
+        gr.JSON(label="Metadata"),
+        gr.JSON(label="Images"),
+    ],
 ).launch()