Narsil HF staff commited on
Commit
d8f642d
1 Parent(s): 1cebade

Upload handler.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. handler.py +64 -0
handler.py ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ from transformers import LayoutLMForTokenClassification, LayoutLMv2Processor
3
+ import torch
4
+ from subprocess import run
5
+
6
+ # install tesseract-ocr and pytesseract
7
+ run("apt install -y tesseract-ocr", shell=True, check=True)
8
+ run("pip install pytesseract", shell=True, check=True)
9
+
10
+ # helper function to unnormalize bboxes for drawing onto the image
11
+ def 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 device
21
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
22
+
23
+
24
+ class EndpointHandler:
25
+ def __init__(self, path=""):
26
+ # load model and processor from path
27
+ 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.Image
35
+ """
36
+ # process input
37
+ image = data.pop("inputs", data)
38
+
39
+ # process image
40
+ encoding = self.processor(image, return_tensors="pt")
41
+
42
+ # run prediction
43
+ 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 output
53
+ 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
+ continue
60
+ 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}