nielsr HF staff commited on
Commit
b35c806
1 Parent(s): da42f5e

First commit

Browse files
Files changed (2) hide show
  1. app.py +87 -0
  2. requirements.txt +4 -0
app.py ADDED
@@ -0,0 +1,87 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ os.system('pip install git+https://github.com/huggingface/transformers.git --upgrade')
3
+ os.system('pip install pyyaml==5.1')
4
+ # workaround: install old version of pytorch since detectron2 hasn't released packages for pytorch 1.9 (issue: https://github.com/facebookresearch/detectron2/issues/3158)
5
+ os.system('pip install torch==1.8.0+cu101 torchvision==0.9.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html')
6
+
7
+ # install detectron2 that matches pytorch 1.8
8
+ # See https://detectron2.readthedocs.io/tutorials/install.html for instructions
9
+ os.system('pip install -q detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.8/index.html')
10
+
11
+ import gradio as gr
12
+ import numpy as np
13
+ from transformers import LayoutLMv2FeatureExtractor, LayoutLMv2Tokenizer, LayoutLMV2ForTokenClassification
14
+ from datasets import load_dataset
15
+ from PIL import Image, ImageDraw, ImageFont
16
+
17
+ ds = load_dataset("hf-internal-testing/fixtures_docvqa", split="test")
18
+
19
+ image = Image.open(ds[0]["file"]).convert("RGB")
20
+ image.save("document.png")
21
+
22
+ feature_extractor = LayoutLMv2FeatureExtractor.from_pretrained("microsoft/layoutlmv2-base-uncased")
23
+ tokenizer = LayoutLMv2Tokenizer.from_pretrained("microsoft/layoutlmv2-base-uncased")
24
+ model = LayoutLMv2ForTokenClassification.from_pretrained("nielsr/layoutlmv2-finetuned-funsd")
25
+
26
+ def unnormalize_box(bbox, width, height):
27
+ return [
28
+ width * (bbox[0] / 1000),
29
+ height * (bbox[1] / 1000),
30
+ width * (bbox[2] / 1000),
31
+ height * (bbox[3] / 1000),
32
+ ]
33
+
34
+ def iob_to_label(label):
35
+ label = label[2:]
36
+ if not label:
37
+ return 'other'
38
+ return label
39
+
40
+ def process_image(image):
41
+ width, height = image.size
42
+
43
+ # get words, boxes
44
+ encoding_feature_extractor = feature_extractor(image, return_tensors="pt")
45
+ words, boxes = encoding_feature_extractor.words, encoding_feature_extractor.boxes
46
+
47
+ # encode
48
+ encoding = tokenizer(words, boxes=boxes, return_offsets_mapping=True, return_tensors="pt")
49
+ offset_mapping = encoding.pop('offset_mapping')
50
+ encoding["image"] = encoding_feature_extractor.pixel_values
51
+
52
+ # forward pass
53
+ outputs = model(**encoding)
54
+
55
+ # get predictions
56
+ predictions = outputs.logits.argmax(-1).squeeze().tolist()
57
+ token_boxes = encoding.bbox.squeeze().tolist()
58
+
59
+ # only keep non-subword predictions
60
+ is_subword = np.array(offset_mapping.squeeze().tolist())[:,0] != 0
61
+ true_predictions = [id2label[pred] for idx, pred in enumerate(predictions) if not is_subword[idx]]
62
+ true_boxes = [unnormalize_box(box, width, height) for idx, box in enumerate(token_boxes) if not is_subword[idx]]
63
+
64
+ # draw predictions over the image
65
+ draw = ImageDraw.Draw(image)
66
+ font = ImageFont.load_default()
67
+ for prediction, box in zip(true_predictions, true_boxes):
68
+ predicted_label = iob_to_label(prediction).lower()
69
+ draw.rectangle(box, outline=label2color[predicted_label])
70
+ draw.text((box[0]+10, box[1]-10), text=predicted_label, fill=label2color[predicted_label], font=font)
71
+
72
+ return image
73
+
74
+
75
+ title = "Interactive demo: LayoutLMv2"
76
+ description = "Demo for Microsoft's LayoutLMv2, a Transformer for state-of-the-art document image understanding tasks. To use it, simply upload an image or use the example image below. Results will show up in a few seconds."
77
+ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2012.14740'>LayoutLMv2: Multi-modal Pre-training for Visually-Rich Document Understanding</a> | <a href='https://github.com/microsoft/unilm'>Github Repo</a></p>"
78
+ examples =[['document.png']]
79
+
80
+ iface = gr.Interface(fn=process_image,
81
+ inputs=gr.inputs.Image(shape=(480, 480), type="pil"),
82
+ outputs=gr.outputs.Image(type='pil', label=f'annotated image'),
83
+ title=title,
84
+ description=description,
85
+ article=article,
86
+ examples=examples)
87
+ iface.launch()
requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ gradio
2
+ Pillow
3
+ numpy
4
+ datasets