Update app.py

app.py

@@ -25,7 +25,12 @@ model = AutoModelForTokenClassification.from_pretrained("SickBoy/layoutlm_docume
 
 
 # load image example
-dataset = load_dataset("darentang/generated", split="test")
+dataset = load_dataset("SickBoy/layout_documents", split="train")
+example = dataset[0]
+image1 = example["image"]
+words = example["tokens"]
+boxes = example["bboxes"]
+labels = example["ner_tags"]
 Image.open(dataset[2]["image_path"]).convert("RGB").save("example1.png")
 Image.open(dataset[1]["image_path"]).convert("RGB").save("example2.png")
 Image.open(dataset[0]["image_path"]).convert("RGB").save("example3.png")
@@ -76,8 +81,9 @@ def process_image(image):
     width, height = image.size
 
     # encode
-    encoding = processor(image, truncation=True, return_offsets_mapping=True, return_tensors="pt")
-    offset_mapping = encoding.pop('offset_mapping')
+    #encoding = processor(image, truncation=True, return_offsets_mapping=True, return_tensors="pt")
+    encoding = processor(image1, words, boxes=boxes, word_labels=word_labels, return_tensors="pt")
+    #offset_mapping = encoding.pop('offset_mapping')
 
     # forward pass
     outputs = model(**encoding)