File size: 968 Bytes
e7604fc 6dcd7d1 e7604fc bdfbfc5 e7604fc 6dcd7d1 bdfbfc5 |
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 |
from typing import Dict, List, Any
from transformers import LayoutLMForTokenClassification, LayoutLMv2Processor
import torch
from subprocess import run
import pytesseract
from pytesseract import Output
# install tesseract-ocr and pytesseract
run("apt install -y tesseract-ocr", shell=True, check=True)
# set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class EndpointHandler:
def __init__(self, path=""):
self.pytesseract_installed = False
try:
import pytesseract
self.pytesseract_installed = True
except ImportError:
print("Pytesseract not installed, will not use OCR")
def __call__(self, data: Dict[str, bytes]) -> Dict[str, List[Any]]:
"""
Args:
data (:obj:):
includes the deserialized image file as PIL.Image
"""
# process input
image = data.pop("inputs", data)
result = pytesseract.image_to_string(image, config='--psm 3', output_type=Output.STRING)
return {"predictions": result}
|