|
from typing import Dict, List, Any |
|
from transformers import LayoutLMForTokenClassification, LayoutLMv2Processor |
|
import torch |
|
from subprocess import run |
|
|
|
|
|
run("apt install -y tesseract-ocr", shell=True, check=True) |
|
run("pip install pytesseract", shell=True, check=True) |
|
|
|
|
|
def unnormalize_box(bbox, width, height): |
|
return [ |
|
width * (bbox[0] / 1000), |
|
height * (bbox[1] / 1000), |
|
width * (bbox[2] / 1000), |
|
height * (bbox[3] / 1000), |
|
] |
|
|
|
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
|
|
class EndpointHandler: |
|
def __init__(self, path=""): |
|
|
|
self.model = LayoutLMForTokenClassification.from_pretrained(path).to(device) |
|
self.processor = LayoutLMv2Processor.from_pretrained(path) |
|
|
|
def __call__(self, data: Dict[str, bytes]) -> Dict[str, List[Any]]: |
|
""" |
|
Args: |
|
data (:obj:): |
|
includes the deserialized image file as PIL.Image |
|
""" |
|
|
|
image = data.pop("inputs", data) |
|
|
|
|
|
encoding = self.processor(image, return_tensors="pt") |
|
|
|
|
|
with torch.inference_mode(): |
|
outputs = self.model( |
|
input_ids=encoding.input_ids.to(device), |
|
bbox=encoding.bbox.to(device), |
|
attention_mask=encoding.attention_mask.to(device), |
|
token_type_ids=encoding.token_type_ids.to(device), |
|
) |
|
predictions = outputs.logits.softmax(-1) |
|
|
|
|
|
result = [] |
|
for item, inp_ids, bbox in zip( |
|
predictions.squeeze(0).cpu(), encoding.input_ids.squeeze(0).cpu(), encoding.bbox.squeeze(0).cpu() |
|
): |
|
label = self.model.config.id2label[int(item.argmax().cpu())] |
|
if label == "O": |
|
continue |
|
score = item.max().item() |
|
text = self.processor.tokenizer.decode(inp_ids) |
|
bbox = unnormalize_box(bbox.tolist(), image.width, image.height) |
|
result.append({"label": label, "score": score, "text": text, "bbox": bbox}) |
|
return {"predictions": result} |
|
|
