Views
No views yet
1import torch
2import numpy as np
3from PIL import Image, ImageDraw, ImageFont
4import pytesseract
5from transformers import LayoutLMForTokenClassification, LayoutLMTokenizer
6
7
8device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
9
10tokenizer = LayoutLMTokenizer.from_pretrained("mrm8488/layoutlm-finetuned-funsd")
11model = LayoutLMForTokenClassification.from_pretrained("mrm8488/layoutlm-finetuned-funsd", num_labels=13)
12model.to(device)
13
14
15image = Image.open("/83443897.png")
16image = image.convert("RGB")
17
18# Display the image
19
20
21# Run Tesseract (OCR) on the image
22
23width, height = image.size
24w_scale = 1000/width
25h_scale = 1000/height
26
27ocr_df = pytesseract.image_to_data(image, output_type='data.frame') \\n
28ocr_df = ocr_df.dropna() \\n .assign(left_scaled = ocr_df.left*w_scale,
29 width_scaled = ocr_df.width*w_scale,
30 top_scaled = ocr_df.top*h_scale,
31 height_scaled = ocr_df.height*h_scale,
32 right_scaled = lambda x: x.left_scaled + x.width_scaled,
33 bottom_scaled = lambda x: x.top_scaled + x.height_scaled)
34
35float_cols = ocr_df.select_dtypes('float').columns
36ocr_df[float_cols] = ocr_df[float_cols].round(0).astype(int)
37ocr_df = ocr_df.replace(r'^\s*$', np.nan, regex=True)
38ocr_df = ocr_df.dropna().reset_index(drop=True)
39ocr_df[:20]
40
41# create a list of words, actual bounding boxes, and normalized boxes
42
43words = list(ocr_df.text)
44coordinates = ocr_df[['left', 'top', 'width', 'height']]
45actual_boxes = []
46for idx, row in coordinates.iterrows():
47 x, y, w, h = tuple(row) # the row comes in (left, top, width, height) format
48 actual_box = [x, y, x+w, y+h] # we turn it into (left, top, left+widght, top+height) to get the actual box
49 actual_boxes.append(actual_box)
50
51def normalize_box(box, width, height):
52 return [
53 int(1000 * (box[0] / width)),
54 int(1000 * (box[1] / height)),
55 int(1000 * (box[2] / width)),
56 int(1000 * (box[3] / height)),
57 ]
58
59boxes = []
60for box in actual_boxes:
61 boxes.append(normalize_box(box, width, height))
62
63# Display boxes
64
65def convert_example_to_features(image, words, boxes, actual_boxes, tokenizer, args, cls_token_box=[0, 0, 0, 0],
66 sep_token_box=[1000, 1000, 1000, 1000],
67 pad_token_box=[0, 0, 0, 0]):
68 width, height = image.size
69
70 tokens = []
71 token_boxes = []
72 actual_bboxes = [] # we use an extra b because actual_boxes is already used
73 token_actual_boxes = []
74 for word, box, actual_bbox in zip(words, boxes, actual_boxes):
75 word_tokens = tokenizer.tokenize(word)
76 tokens.extend(word_tokens)
77 token_boxes.extend([box] * len(word_tokens))
78 actual_bboxes.extend([actual_bbox] * len(word_tokens))
79 token_actual_boxes.extend([actual_bbox] * len(word_tokens))
80
81 # Truncation: account for [CLS] and [SEP] with "- 2".
82 special_tokens_count = 2
83 if len(tokens) > args.max_seq_length - special_tokens_count:
84 tokens = tokens[: (args.max_seq_length - special_tokens_count)]
85 token_boxes = token_boxes[: (args.max_seq_length - special_tokens_count)]
86 actual_bboxes = actual_bboxes[: (args.max_seq_length - special_tokens_count)]
87 token_actual_boxes = token_actual_boxes[: (args.max_seq_length - special_tokens_count)]
88
89 # add [SEP] token, with corresponding token boxes and actual boxes
90 tokens += [tokenizer.sep_token]
91 token_boxes += [sep_token_box]
92 actual_bboxes += [[0, 0, width, height]]
93 token_actual_boxes += [[0, 0, width, height]]
94
95 segment_ids = [0] * len(tokens)
96
97 # next: [CLS] token
98 tokens = [tokenizer.cls_token] + tokens
99 token_boxes = [cls_token_box] + token_boxes
100 actual_bboxes = [[0, 0, width, height]] + actual_bboxes
101 token_actual_boxes = [[0, 0, width, height]] + token_actual_boxes
102 segment_ids = [1] + segment_ids
103
104 input_ids = tokenizer.convert_tokens_to_ids(tokens)
105
106 # The mask has 1 for real tokens and 0 for padding tokens. Only real
107 # tokens are attended to.
108 input_mask = [1] * len(input_ids)
109
110 # Zero-pad up to the sequence length.
111 padding_length = args.max_seq_length - len(input_ids)
112 input_ids += [tokenizer.pad_token_id] * padding_length
113 input_mask += [0] * padding_length
114 segment_ids += [tokenizer.pad_token_id] * padding_length
115 token_boxes += [pad_token_box] * padding_length
116 token_actual_boxes += [pad_token_box] * padding_length
117
118 assert len(input_ids) == args.max_seq_length
119 assert len(input_mask) == args.max_seq_length
120 assert len(segment_ids) == args.max_seq_length
121 assert len(token_boxes) == args.max_seq_length
122 assert len(token_actual_boxes) == args.max_seq_length
123
124 return input_ids, input_mask, segment_ids, token_boxes, token_actual_boxes
125
126input_ids, input_mask, segment_ids, token_boxes, token_actual_boxes = convert_example_to_features(image=image, words=words, boxes=boxes, actual_boxes=actual_boxes, tokenizer=tokenizer, args=args)
127
128input_ids = torch.tensor(input_ids, device=device).unsqueeze(0)
129attention_mask = torch.tensor(input_mask, device=device).unsqueeze(0)
130token_type_ids = torch.tensor(segment_ids, device=device).unsqueeze(0)
131bbox = torch.tensor(token_boxes, device=device).unsqueeze(0)
132
133
134outputs = model(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask, token_type_ids=token_type_ids)
135
136token_predictions = outputs.logits.argmax(-1).squeeze().tolist() # the predictions are at the token level
137
138word_level_predictions = [] # let's turn them into word level predictions
139final_boxes = []
140for id, token_pred, box in zip(input_ids.squeeze().tolist(), token_predictions, token_actual_boxes):
141 if (tokenizer.decode([id]).startswith("##")) or (id in [tokenizer.cls_token_id,
142 tokenizer.sep_token_id,
143 tokenizer.pad_token_id]):
144 # skip prediction + bounding box
145
146 continue
147 else:
148 word_level_predictions.append(token_pred)
149 final_boxes.append(box)
150
151#print(word_level_predictions)
152
153
154draw = ImageDraw.Draw(image)
155
156font = ImageFont.load_default()
157
158def iob_to_label(label):
159 if label != 'O':
160 return label[2:]
161 else:
162 return "other"
163
164label2color = {'question':'blue', 'answer':'green', 'header':'orange', 'other':'violet'}
165
166for prediction, box in zip(word_level_predictions, final_boxes):
167 predicted_label = iob_to_label(label_map[prediction]).lower()
168 draw.rectangle(box, outline=label2color[predicted_label])
169 draw.text((box[0] + 10, box[1] - 10), text=predicted_label, fill=label2color[predicted_label], font=font)
170
171# Display the result (image)
172