Update sroie_inference.py

sroie_inference.py

@@ -6,9 +6,9 @@ from transformers import LayoutLMv3TokenizerFast, LayoutLMv3Processor, LayoutLMv
 from utils import OCR, unnormalize_box
 
 
-tokenizer = LayoutLMv3TokenizerFast.from_pretrained("mp-02/layoutlmv3-finetuned-cord-sroie", apply_ocr=False)
-processor = LayoutLMv3Processor.from_pretrained("mp-02/layoutlmv3-finetuned-cord-sroie", apply_ocr=False)
-model = LayoutLMv3ForTokenClassification.from_pretrained("mp-02/layoutlmv3-finetuned-cord-sroie")
+tokenizer = LayoutLMv3TokenizerFast.from_pretrained("mp-02/layoutlmv3-finetuned-sroie", apply_ocr=False)
+processor = LayoutLMv3Processor.from_pretrained("mp-02/layoutlmv3-finetuned-sroie", apply_ocr=False)
+model = LayoutLMv3ForTokenClassification.from_pretrained("mp-02/layoutlmv3-finetuned-sroie")
 
 id2label = model.config.id2label
 label2id = model.config.label2id
@@ -85,8 +85,8 @@ def prediction(image):
             d[i] = d[i] + ", " + true_words[id]
     d = {k: v.strip() for (k, v) in d.items()}
 
-    #if "O" in d: d.pop("O")
-    #if "TOTAL" in d: d.pop("TOTAL")
+    if "O" in d: d.pop("O")
+    if "TOTAL" in d: d.pop("TOTAL")
         
     blur_boxes = []
     for prediction, box in zip(true_predictions, true_boxes):
@@ -95,12 +95,12 @@ def prediction(image):
 
     image = (blur(image, blur_boxes))
 
-    draw = ImageDraw.Draw(image, "RGBA")
-    font = ImageFont.load_default()
+    #draw = ImageDraw.Draw(image, "RGBA")
+    #font = ImageFont.load_default()
 
-    for prediction, box in zip(true_predictions, true_boxes):
-       draw.rectangle(box)
-       draw.text((box[0]+10, box[1]-10), text=prediction, font=font, fill="black", font_size="8")
+    #for prediction, box in zip(true_predictions, true_boxes):
+    #   draw.rectangle(box)
+    #   draw.text((box[0]+10, box[1]-10), text=prediction, font=font, fill="black", font_size="8")
 
     return d, image