Update app.py

app.py

@@ -29,6 +29,13 @@ processor = LiltForTokenClassification.from_pretrained("SCUT-DLVCLab/lilt-robert
 model = LayoutLMv3Processor.from_pretrained(
     "jinhybr/LiLt-funsd-en"
 )
+####
+
+
+
+
+
+####
 
 # load image example
 dataset = load_dataset("nielsr/funsd-layoutlmv3", split="test")
@@ -38,14 +45,11 @@ image.save("document.png")
 
 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",
-}
 
 
+
+
+# helper function to unnormalize bboxes for drawing onto the image
 def unnormalize_box(bbox, width, height):
     return [
         width * (bbox[0] / 1000),
@@ -55,6 +59,18 @@ def unnormalize_box(bbox, width, height):
     ]
 
 
+label2color = {
+    "B-HEADER": "blue",
+    "B-QUESTION": "red",
+    "B-ANSWER": "green",
+    "I-HEADER": "blue",
+    "I-QUESTION": "red",
+    "I-ANSWER": "green",
+}
+
+
+
+
 def iob_to_label(label):
     label = label[2:]
     if not label:
@@ -62,6 +78,27 @@ def iob_to_label(label):
     return label
 
 
+
+
+# draw results onto the image
+def draw_boxes(image, boxes, predictions):
+    width, height = image.size
+    normalizes_boxes = [unnormalize_box(box, width, height) for box in boxes]
+
+    # draw predictions over the image
+    draw = ImageDraw.Draw(image)
+    font = ImageFont.load_default()
+    for prediction, box in zip(predictions, normalizes_boxes):
+        if prediction == "O":
+            continue
+        draw.rectangle(box, outline="black")
+        draw.rectangle(box, outline=label2color[prediction])
+        draw.text((box[0] + 10, box[1] - 10), text=prediction, fill=label2color[prediction], font=font)
+    return image
+
+
+
+
 def process_image(image):
     width, height = image.size
 
@@ -89,7 +126,9 @@ def process_image(image):
         if not is_subword[idx]
     ]
 
-    # draw predictions over the image
+    draw_boxes(image, true_boxes, true_predictions)
+
+''''    # draw predictions over the image
     draw = ImageDraw.Draw(image)
     font = ImageFont.load_default()
     for prediction, box in zip(true_predictions, true_boxes):
@@ -101,6 +140,9 @@ def process_image(image):
             fill=label2color[predicted_label],
             font=font,
         )
+''''
+
+
 
     return image