File size: 1,727 Bytes
c348f38 897bb56 c348f38 c595fbe c348f38 2a0cd88 |
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 |
from typing import Dict, Any, List
import os
import base64
current_dir = os.getcwd()
os.environ['HF_HOME'] = os.path.join(current_dir)
os.environ['PAGINATE_OUTPUT']='True'
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=""):
# Initialize the OCR model
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:
inputs (:obj: dict): A dictionary containing the inputs.
max_pages (:obj: int): The maximum number of pages to process.
file (:obj: str): The base64-encoded PDF file content.
Return:
A list of dictionaries containing the extracted text.
"""
inputs = data.get("inputs", {})
file_content = inputs.get("file")
max_pages = inputs.get("max_pages", None)
# Decode the base64-encoded file content
file_bytes = base64.b64decode(file_content)
self.upload_file(BytesIO(file_bytes))
pdf_path = self.file_location
# Perform OCR on the input PDF
extracted_text, _, _ = convert_single_pdf(pdf_path, self.models, max_pages=max_pages, langs=["vi"])
# Return the extracted text
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 |