|
from typing import Dict, Any, List
|
|
import os
|
|
|
|
current_dir = os.getcwd()
|
|
os.environ['HF_HOME'] = os.path.join(current_dir)
|
|
|
|
from marker.convert import convert_single_pdf
|
|
from marker.logger import configure_logging
|
|
from marker.models import load_all_models
|
|
from marker.output import save_markdown
|
|
from io import BytesIO
|
|
class EndpointHandler:
|
|
def __init__(self, path=""):
|
|
|
|
self.models = load_all_models()
|
|
self.file_location = "input/temp.pdf"
|
|
os.makedirs("input", exist_ok=True)
|
|
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
|
|
"""
|
|
data args:
|
|
max_pages (:obj: int): The maximum number of pages to process.
|
|
file (:obj: UploadFile): The uploaded PDF file.
|
|
Return:
|
|
A list of dictionaries containing the extracted text.
|
|
"""
|
|
|
|
self.upload_file(data['file'])
|
|
pdf_path = self.file_location
|
|
max_pages = data.get("max_pages", None)
|
|
|
|
|
|
extracted_text, _, _ = convert_single_pdf(pdf_path, self.models, max_pages=max_pages, langs=["vi"])
|
|
|
|
return [{"extracted_text": extracted_text}]
|
|
|
|
def upload_file(self, file: BytesIO, max_pages: int = None):
|
|
with open(self.file_location, "wb") as f:
|
|
f.write(file.read())
|
|
return True |
