Update app.py

app.py

@@ -1,29 +1,43 @@
 import os
-os.system('pip install git+https://github.com/huggingface/transformers.git --upgrade')
-# workaround: install old version of pytorch since detectron2 hasn't released packages for pytorch 1.9 (issue: https://github.com/facebookresearch/detectron2/issues/3158)
-os.system('pip install torch==1.8.0+cu101 torchvision==0.9.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html')
-
-## install PyTesseract
-os.system('pip install -q pytesseract')
+os.system('pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu')
 
 import gradio as gr
 import numpy as np
+from transformers import AutoModelForTokenClassification
+from datasets.features import ClassLabel
+from transformers import AutoProcessor
+from datasets import Features, Sequence, ClassLabel, Value, Array2D, Array3D
+import torch
+from datasets import load_metric
+from transformers import LayoutLMv3ForTokenClassification
+from transformers.data.data_collator import default_data_collator
+
 
+from transformers import AutoModelForTokenClassification
 from datasets import load_dataset
 from PIL import Image, ImageDraw, ImageFont
 
-from transformers import AutoProcessor, AutoModelForTokenClassification
-processor = AutoProcessor.from_pretrained("Ammar-alhaj-ali/LayoutLMv3-Fine-Tuning-FUNSD")
+
+processor = AutoProcessor.from_pretrained("Ammar-alhaj-ali/LayoutLMv3-Fine-Tuning-FUNSD", apply_ocr=True)
 model = AutoModelForTokenClassification.from_pretrained("Ammar-alhaj-ali/LayoutLMv3-Fine-Tuning-FUNSD")
 
 # load image example
-dataset = load_dataset("nielsr/funsd", split="test")
-image = Image.open(dataset[0]["image_path"]).convert("RGB")
-image.save("document.png")
+dataset = load_dataset("darentang/generated", split="test")
+Image.open(dataset[2]["image_path"]).convert("RGB").save("img1.png")
+Image.open(dataset[1]["image_path"]).convert("RGB").save("img2.png")
+Image.open(dataset[0]["image_path"]).convert("RGB").save("img3.png")
 # define id2label, label2color
-labels = dataset.features['ner_tags'].feature.names
+labels = ['O', 'B-HEADER', 'I-HEADER', 'B-QUESTION', 'I-QUESTION', 'B-ANSWER', 'I-ANSWER']
 id2label = {v: k for v, k in enumerate(labels)}
-label2color = {'question':'blue', 'answer':'green', 'header':'orange', 'other':'violet'}
+label2color = {
+    "B-HEADER": 'red',
+    "I-HEADER": 'green',
+    "B-QUESTION": 'orange',
+    "I-QUESTION": "blue",
+    "B-ANSWER": 'gray',
+    "I-ANSWERE": 'violet',
+    "O": 'orange'
+  } 
 
 def unnormalize_box(bbox, width, height):
      return [
@@ -33,13 +47,15 @@ def unnormalize_box(bbox, width, height):
          height * (bbox[3] / 1000),
      ]
 
+
 def iob_to_label(label):
-    label = label[2:]
-    if not label:
-      return 'other'
     return label
 
+
+
 def process_image(image):
+
+    print(type(image))
     width, height = image.size
 
     # encode
@@ -62,16 +78,19 @@ def process_image(image):
     draw = ImageDraw.Draw(image)
     font = ImageFont.load_default()
     for prediction, box in zip(true_predictions, true_boxes):
-        predicted_label = iob_to_label(prediction).lower()
+        predicted_label = iob_to_label(prediction)
         draw.rectangle(box, outline=label2color[predicted_label])
         draw.text((box[0]+10, box[1]-10), text=predicted_label, fill=label2color[predicted_label], font=font)
     
     return image
 
-title = "Interactive demo: LayoutLMv2"
-description = "Demo for Microsoft's LayoutLMv2, a Transformer for state-of-the-art document image understanding tasks. This particular model is fine-tuned on FUNSD, a dataset of manually annotated forms. It annotates the words into QUESTION/ANSWER/HEADER/OTHER. To use it, simply upload an image or use the example image below. Results will show up in a few seconds."
-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>"
-examples =[['img1.png']]
+
+title = "Extraction d'informations de factures en utilisant le modèle LayoutLMv3"
+description = "J'utilise LayoutLMv3 de Microsoft formé sur un ensemble de données de factures pour prédire le nom de l'émetteur de factures, l'adresse de l'émetteur de factures, le code postal de l'émetteur de factures, la date d'échéance, la TPS, la date de facturation, le numéro de facture, le sous-total et le total. Pour l'utiliser, il suffit de télécharger une image ou d'utiliser l'exemple d'image ci-dessous. Les résultats seront affichés en quelques secondes."
+
+article="<b>References</b><br>[1] Y. Xu et al., “LayoutLMv3: Pre-training for Document AI with Unified Text and Image Masking.” 2022. <a href='https://arxiv.org/abs/2204.08387'>Paper Link</a><br>[2]  <a href='https://github.com/NielsRogge/Transformers-Tutorials/tree/master/LayoutLMv3'>LayoutLMv3 training and inference</a>" 
+
+examples =[['img1.png'],['img2.png'],['img3.png']]
 
 css = """.output_image, .input_image {height: 600px !important}"""
 
@@ -82,5 +101,7 @@ iface = gr.Interface(fn=process_image,
                      description=description,
                      article=article,
                      examples=examples,
-                     css=css)
-iface.launch(debug=True)
+                     css=css,
+                     analytics_enabled = True, enable_queue=True)
+
+iface.launch(inline=False, share=False, debug=False)