Update README.md

README.md

@@ -25,17 +25,73 @@ It achieves the following results on the evaluation set:
 - Overall F1: 0.8900
 - Overall Accuracy: 0.8204
 
-## Model description
+## Model Usage
+
+```python
+from transformers import LiltForTokenClassification, LayoutLMv3Processor
+from PIL import Image, ImageDraw, ImageFont
+import torch
+
+# load model and processor from huggingface hub
+model = LiltForTokenClassification.from_pretrained("philschmid/lilt-en-funsd")
+processor = LayoutLMv3Processor.from_pretrained("philschmid/lilt-en-funsd")
+
+
+# helper function to unnormalize bboxes for drawing onto the image
+def unnormalize_box(bbox, width, height):
+    return [
+        width * (bbox[0] / 1000),
+        height * (bbox[1] / 1000),
+        width * (bbox[2] / 1000),
+        height * (bbox[3] / 1000),
+    ]
+
+
+label2color = {
+    "B-HEADER": "blue",
+    "B-QUESTION": "red",
+    "B-ANSWER": "green",
+    "I-HEADER": "blue",
+    "I-QUESTION": "red",
+    "I-ANSWER": "green",
+}
+# 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
+
+
+# run inference
+def run_inference(image, model=model, processor=processor, output_image=True):
+    # create model input
+    encoding = processor(image, return_tensors="pt")
+    del encoding["pixel_values"]
+    # run inference
+    outputs = model(**encoding)
+    predictions = outputs.logits.argmax(-1).squeeze().tolist()
+    # get labels
+    labels = [model.config.id2label[prediction] for prediction in predictions]
+    if output_image:
+        return draw_boxes(image, encoding["bbox"][0], labels)
+    else:
+        return labels
+
+
+run_inference(dataset["test"][34]["image"])
+
+```
 
-More information needed
-
-## Intended uses & limitations
-
-More information needed
-
-## Training and evaluation data
-
-More information needed
 
 ## Training procedure