philschmid/layoutlm-funsd

1from typing import Dict, List, Any2from transformers import LayoutLMForTokenClassification, LayoutLMv2Processor3import torch4from subprocess import run5 6# install tesseract-ocr and pytesseract7run("apt install -y tesseract-ocr", shell=True, check=True)8run("pip install pytesseract", shell=True, check=True)9 10# helper function to unnormalize bboxes for drawing onto the image11def unnormalize_box(bbox, width, height):12    return [13        width * (bbox[0] / 1000),14        height * (bbox[1] / 1000),15        width * (bbox[2] / 1000),16        height * (bbox[3] / 1000),17    ]18 19 20# set device21device = torch.device("cuda" if torch.cuda.is_available() else "cpu")22 23 24class EndpointHandler:25    def __init__(self, path=""):26        # load model and processor from path27        self.model = LayoutLMForTokenClassification.from_pretrained(path).to(device)28        self.processor = LayoutLMv2Processor.from_pretrained(path)29 30    def __call__(self, data: Dict[str, bytes]) -> Dict[str, List[Any]]:31        """32        Args:33            data (:obj:):34                includes the deserialized image file as PIL.Image35        """36        # process input37        image = data.pop("inputs", data)38 39        # process image40        encoding = self.processor(image, return_tensors="pt")41 42        # run prediction43        with torch.inference_mode():44            outputs = self.model(45                input_ids=encoding.input_ids.to(device),46                bbox=encoding.bbox.to(device),47                attention_mask=encoding.attention_mask.to(device),48                token_type_ids=encoding.token_type_ids.to(device),49            )50            predictions = outputs.logits.softmax(-1)51 52        # post process output53        result = []54        for item, inp_ids, bbox in zip(55            predictions.squeeze(0).cpu(), encoding.input_ids.squeeze(0).cpu(), encoding.bbox.squeeze(0).cpu()56        ):57            label = self.model.config.id2label[int(item.argmax().cpu())]58            if label == "O":59                continue60            score = item.max().item()61            text = self.processor.tokenizer.decode(inp_ids)62            bbox = unnormalize_box(bbox.tolist(), image.width, image.height)63            result.append({"label": label, "score": score, "text": text, "bbox": bbox})64        return {"predictions": result}65