Update app.py

app.py

@@ -51,7 +51,14 @@ processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base",appl
 
 model = LayoutLMv3ForTokenClassification.from_pretrained("nielsr/layoutlmv3-finetuned-funsd")
 
-dataset = load_dataset("nielsr/funsd-layoutlmv3")
+dataset = load_dataset("nielsr/funsd", split="test")
+image = Image.open(dataset[0]["image_path"]).convert("RGB")
+image = Image.open("./invoice.png")
+image.save("document.png")
+
+
+
+#dataset = load_dataset("nielsr/funsd-layoutlmv3")
 
 example = dataset["test"][0]
 example["image"].save("example1.png")
@@ -106,7 +113,7 @@ def get_label_list(labels):
     label_list.sort()
     return label_list
 
-label2color = {'question':'blue', 'answer':'green', 'header':'orange', 'other':'violet'}
+#label2color = {'question':'blue', 'answer':'green', 'header':'orange', 'other':'violet'}
 
 def unnormalize_box(bbox, width, height):
      return [
@@ -116,14 +123,11 @@ def unnormalize_box(bbox, width, height):
          height * (bbox[3] / 1000),
      ]
 
-# we need to define custom features for `set_format` (used later on) to work properly
-features = Features({
-    'pixel_values': Array3D(dtype="float32", shape=(3, 224, 224)),
-    'input_ids': Sequence(feature=Value(dtype='int64')),
-    'attention_mask': Sequence(Value(dtype='int64')),
-    'bbox': Array2D(dtype="int64", shape=(512, 4)),
-    'labels': Sequence(feature=Value(dtype='int64')),
-})
+def iob_to_label(label):
+        label= label[2:]
+        if not label:
+            return 'other'
+    return label
 
 def process_image(image):
     width, height = image.size
@@ -148,12 +152,6 @@ def process_image(image):
     draw = ImageDraw.Draw(image)
     font = ImageFont.load_default()
 
-    def iob_to_label(labels):
-        labels= labels[2:]
-        if not labels:
-            return 'other'
-    return labels
-
     label2color = {'question':'blue', 'answer':'green', 'header':'orange', 'other':'violet'}
     for prediction, box in zip(true_predictions, true_boxes):
         predicted_label = iob_to_label(prediction) #.lower()