Update sroie_inference.py

sroie_inference.py

@@ -5,8 +5,8 @@ from PIL import Image, ImageDraw, ImageFont
 from transformers import LayoutLMv3TokenizerFast, LayoutLMv3Processor, LayoutLMv3ForTokenClassification
 from utils import OCR, unnormalize_box
 
-
-labels = ["O", "B-COMPANY", "I-COMPANY", "B-DATE", "I-DATE", "B-ADDRESS", "I-ADDRESS", "B-TOTAL", "I-TOTAL"]
+#  ["O", "B-COMPANY", "I-COMPANY", "B-DATE", "I-DATE", "B-ADDRESS", "I-ADDRESS", "B-TOTAL", "I-TOTAL"]
+labels = ["O", "COMPANY", "COMPANY", "DATE", "DATE", "ADDRESS", "ADDRESS", "TOTAL", "TOTAL"]
 id2label = {v: k for v, k in enumerate(labels)}
 label2id = {k: v for v, k in enumerate(labels)}
 
@@ -83,19 +83,9 @@ def prediction(image):
             d[i] = d[i] + ", " + true_words[id]
     d = {k: v.strip() for (k, v) in d.items()}
 
-    keys_to_pop = []
-    for k, v in d.items():
-        if k[:2] == "I-":
-            d["B-" + k[2:]] = d["B-" + k[2:]] + ", " + v
-            keys_to_pop.append(k)
-
     if "O" in d: d.pop("O")
-    if "B-TOTAL" in d: d.pop("B-TOTAL")
-    for k in keys_to_pop: d.pop(k)
-
-    for k in d.keys():
-        k = k[2:]
-
+    if "TOTAL" in d: d.pop("TOTAL")
+        
     blur_boxes = []
     for prediction, box in zip(true_predictions, true_boxes):
         if prediction != 'O' and prediction[2:] != 'TOTAL':