File size: 2,581 Bytes
c8a32e7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import os
import base64
from marker.convert import convert_single_pdf
from marker.models import load_all_models
from marker.settings import Settings
import gradio as gr
model_list = load_all_models()
def parse_pdf_and_return_markdown(pdf_file: bytes , extract_images: bool):
full_text, images, out_meta = convert_single_pdf(pdf_file, model_list)
image_data = {}
if extract_images:
for i, (filename, image) in enumerate(images.items()):
image_filepath = f"image_{i+1}.png"
image.save(image_filepath, "PNG")
with open(image_filepath, "rb") as f:
image_bytes = f.read()
image_base64 = base64.b64encode(image_bytes).decode('utf-8')
image_data[f'image_{i+1}'] = image_base64
os.remove(image_filepath)
return full_text, out_meta, image_data
# @app.post("/convert")
# async def convert_pdf_to_markdown(pdf_file: UploadFile = File(...), extract_images: bool = True):
# if extract_images == False:
# Settings.EXTRACT_IMAGES = False
# print("Print EXTRACT_IMAGES set to False")
# else:
# Settings.EXTRACT_IMAGES = True
# if pdf_file.content_type != "application/pdf":
# raise HTTPException(
# status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE,
# detail=f'File {pdf_file.filename} has unsupported extension type',
# )
# markdown_text, metadata, image_data = parse_pdf_and_return_markdown(await pdf_file.read(), extract_images=extract_images)
# return {"markdown": markdown_text, "metadata": metadata, "images": image_data }
# def main():
# # Parse command-line arguments
# parser = argparse.ArgumentParser(description="Run the marker-api server.")
# parser.add_argument("--host", default="127.0.0.1", help="Host IP address")
# parser.add_argument("--port", type=int, default=8000, help="Port number")
# args = parser.parse_args()
# # Load all models before starting the server
# configure_logging() # Assuming this function initializes logging
# # Start the server
# import uvicorn
# uvicorn.run(app, host=args.host, port=args.port)
with gr.Blocks() as server:
gr.Markdown("Upload a PDF file to convert to markdown.")
gr.Interface(
parse_pdf_and_return_markdown,
inputs=[gr.File(label="Upload PDF", type="filepath"), gr.Checkbox(label="Extract Images")],
outputs=[gr.Textbox(label="Markdown"), gr.JSON(label="Metadata"), gr.JSON(label="Images")]
)
if __name__ == "__main__":
server.launch() |