oussama commited on
Commit
a1f54f4
1 Parent(s): a45b18d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -12
app.py CHANGED
@@ -9,27 +9,21 @@ from transformers import AutoProcessor
9
  from datasets import Features, Sequence, ClassLabel, Value, Array2D, Array3D
10
  import torch
11
  from datasets import load_metric
12
- from transformers import LayoutLMv3ForTokenClassification
13
  from transformers.data.data_collator import default_data_collator
14
 
15
 
16
- from transformers import AutoModelForTokenClassification
17
  from datasets import load_dataset
18
  from PIL import Image, ImageDraw, ImageFont
19
 
20
 
21
- processor = AutoProcessor.from_pretrained("microsoft/layoutlmv3-base", apply_ocr=True)
22
- model = AutoModelForTokenClassification.from_pretrained("oussama/Layoutlm_Form_information_extraction")
 
23
 
24
 
25
 
26
- # load image example
27
- dataset = load_dataset("darentang/generated", split="test")
28
- Image.open(dataset[2]["image_path"]).convert("RGB").save("example1.png")
29
- Image.open(dataset[1]["image_path"]).convert("RGB").save("example2.png")
30
- Image.open(dataset[0]["image_path"]).convert("RGB").save("example3.png")
31
- # define id2label, label2color
32
- labels = dataset.features['ner_tags'].feature.names
33
  id2label = {0: 'O', 1: 'B-HEADER', 2: 'I-HEADER', 3: 'B-QUESTION', 4: 'I-QUESTION', 5: 'B-ANSWER', 6: 'I-ANSWER'}
34
  label2color = {'question':'blue', 'answer':'green', 'header':'orange', 'other':'violet'}
35
 
@@ -54,7 +48,7 @@ def process_image(image):
54
  width, height = image.size
55
 
56
  # encode
57
- encoding = processor(image, truncation=True, return_offsets_mapping=True, return_tensors="pt")
58
  offset_mapping = encoding.pop('offset_mapping')
59
 
60
  # forward pass
 
9
  from datasets import Features, Sequence, ClassLabel, Value, Array2D, Array3D
10
  import torch
11
  from datasets import load_metric
12
+ from transformers import LayoutLMTokenizer
13
  from transformers.data.data_collator import default_data_collator
14
 
15
 
16
+ from transformers import LayoutLMForTokenClassification
17
  from datasets import load_dataset
18
  from PIL import Image, ImageDraw, ImageFont
19
 
20
 
21
+ tokenizer = LayoutLMTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
22
+ model = LayoutLMForTokenClassification.from_pretrained("microsoft/layoutlm-base-uncased", num_labels=13)
23
+
24
 
25
 
26
 
 
 
 
 
 
 
 
27
  id2label = {0: 'O', 1: 'B-HEADER', 2: 'I-HEADER', 3: 'B-QUESTION', 4: 'I-QUESTION', 5: 'B-ANSWER', 6: 'I-ANSWER'}
28
  label2color = {'question':'blue', 'answer':'green', 'header':'orange', 'other':'violet'}
29
 
 
48
  width, height = image.size
49
 
50
  # encode
51
+ encoding = tokenizer(image, truncation=True, return_offsets_mapping=True, return_tensors="pt")
52
  offset_mapping = encoding.pop('offset_mapping')
53
 
54
  # forward pass