Views
No views yet
handler.py1from typing import Dict, List, Any
2from transformers import LayoutLMForTokenClassification, LayoutLMv2Processor
3import torch
4from subprocess import run
5
6# install tesseract-ocr and pytesseract
7run("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 image
11def 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# set device
20device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
21
22class EndpointHandler:
23 def __init__(self, path=""):
24 # load model and processor from path
25 self.model = LayoutLMForTokenClassification.from_pretrained(path).to(device)
26 self.processor = LayoutLMv2Processor.from_pretrained(path)
27
28 def __call__(self, data: Dict[str, bytes]) -> Dict[str, List[Any]]:
29 """
30 Args:
31 data (:obj:):
32 includes the deserialized image file as PIL.Image
33 """
34 # process input
35 image = data.pop("inputs", data)
36
37 # process image
38 encoding = self.processor(image, return_tensors="pt")
39
40 # run prediction
41 with torch.inference_mode():
42 outputs = self.model(
43 input_ids=encoding.input_ids.to(device),
44 bbox=encoding.bbox.to(device),
45 attention_mask=encoding.attention_mask.to(device),
46 token_type_ids=encoding.token_type_ids.to(device),
47 )
48 predictions = outputs.logits.softmax(-1)
49
50 # post process output
51 result = []
52 for item, inp_ids, bbox in zip(
53 predictions.squeeze(0).cpu(), encoding.input_ids.squeeze(0).cpu(), encoding.bbox.squeeze(0).cpu()
54 ):
55 label = self.model.config.id2label[int(item.argmax().cpu())]
56 if label == "O":
57 continue
58 score = item.max().item()
59 text = self.processor.tokenizer.decode(inp_ids)
60 bbox = unnormalize_box(bbox.tolist(), image.width, image.height)
61 result.append({"label": label, "score": score, "text": text, "bbox": bbox})
62 return {"predictions": result}requests to send our requests. (make your you have it installed pip install requests)1import json
2import requests as r
3import mimetypes
4
5ENDPOINT_URL="" # url of your endpoint
6HF_TOKEN="" # organization token where you deployed your endpoint
7
8def predict(path_to_image:str=None):
9 with open(path_to_image, "rb") as i:
10 b = i.read()
11 headers= {
12 "Authorization": f"Bearer {HF_TOKEN}",
13 "Content-Type": mimetypes.guess_type(path_to_image)[0]
14 }
15 response = r.post(ENDPOINT_URL, headers=headers, data=b)
16 return response.json()
17
18prediction = predict(path_to_image="path_to_your_image.png")
19
20print(prediction)
21# {'predictions': [{'label': 'I-ANSWER', 'score': 0.4823932945728302, 'text': '[CLS]', 'bbox': [0.0, 0.0, 0.0, 0.0]}, {'label': 'B-HEADER', 'score': 0.992474377155304, 'text': 'your', 'bbox': [1712.529, 181.203, 1859.949, 228.88799999999998]},1from PIL import Image, ImageDraw, ImageFont
2
3# draw results on image
4def draw_result(path_to_image,result):
5 image = Image.open(path_to_image)
6 label2color = {
7 "B-HEADER": "blue",
8 "B-QUESTION": "red",
9 "B-ANSWER": "green",
10 "I-HEADER": "blue",
11 "I-QUESTION": "red",
12 "I-ANSWER": "green",
13 }
14
15 # draw predictions over the image
16 draw = ImageDraw.Draw(image)
17 font = ImageFont.load_default()
18 for res in result:
19 draw.rectangle(res["bbox"], outline="black")
20 draw.rectangle(res["bbox"], outline=label2color[res["label"]])
21 draw.text((res["bbox"][0] + 10, res["bbox"][1] - 10), text=res["label"], fill=label2color[res["label"]], font=font)
22 return image
23
24draw_result("path_to_your_image.png", prediction["predictions"])