Fix app.py

app.py

@@ -8,21 +8,29 @@ os.system('pip install torch==1.8.0+cu101 torchvision==0.9.0+cu101 -f https://do
 # See https://detectron2.readthedocs.io/tutorials/install.html for instructions
 os.system('pip install -q detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.8/index.html')
 
+## install PyTesseract
+os.system('sudo apt install tesseract-ocr')
+os.system('pip install -q pytesseract')
+
 import gradio as gr
 import numpy as np
-from transformers import LayoutLMv2FeatureExtractor, LayoutLMv2Tokenizer, LayoutLMV2ForTokenClassification
+from transformers import LayoutLMv2FeatureExtractor, LayoutLMv2TokenizerFast, LayoutLMv2ForTokenClassification
 from datasets import load_dataset
 from PIL import Image, ImageDraw, ImageFont
 
-ds = load_dataset("hf-internal-testing/fixtures_docvqa", split="test")
-
-image = Image.open(ds[0]["file"]).convert("RGB")
-image.save("document.png")
-
 feature_extractor = LayoutLMv2FeatureExtractor.from_pretrained("microsoft/layoutlmv2-base-uncased")
-tokenizer = LayoutLMv2Tokenizer.from_pretrained("microsoft/layoutlmv2-base-uncased")
+tokenizer = LayoutLMv2TokenizerFast.from_pretrained("microsoft/layoutlmv2-base-uncased")
 model = LayoutLMv2ForTokenClassification.from_pretrained("nielsr/layoutlmv2-finetuned-funsd")
 
+# load image example
+dataset = load_dataset("nielsr/funsd", split="test")
+image = Image.open(dataset[0]["image_path"]).convert("RGB")
+image.save("document.png")
+# define id2label, label2color
+labels = dataset.features['ner_tags'].feature.names
+id2label = {v: k for v, k in enumerate(labels)}
+label2color = {'question':'blue', 'answer':'green', 'header':'orange', 'other':'violet'}
+
 def unnormalize_box(bbox, width, height):
      return [
          width * (bbox[0] / 1000),
@@ -78,10 +86,10 @@ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2012.147
 examples =[['document.png']]
 
 iface = gr.Interface(fn=process_image, 
-                     inputs=gr.inputs.Image(shape=(480, 480), type="pil"), 
-                     outputs=gr.outputs.Image(type='pil', label=f'annotated image'),
+                     inputs=gr.inputs.Image(type="pil"), 
+                     outputs=gr.outputs.Image(type="pil", label="annotated image"),
                      title=title,
                      description=description,
                      article=article,
                      examples=examples)
-iface.launch()
+iface.launch(debug=True)