Update cord_inference.py

cord_inference.py

@@ -4,90 +4,50 @@ from transformers import LayoutLMv3TokenizerFast, LayoutLMv3Processor, LayoutLMv
 from PIL import Image, ImageDraw, ImageFont
 from utils import OCR, unnormalize_box
 
-
-id2label = {
-    "0": "O",
-    "1": "B-MENU.CNT",
-    "2": "B-MENU.DISCOUNTPRICE",
-    "3": "B-MENU.NM",
-    "4": "B-MENU.NUM",
-    "5": "B-MENU.PRICE",
-    "6": "B-MENU.SUB.CNT",
-    "7": "B-MENU.SUB.NM",
-    "8": "B-MENU.SUB.PRICE",
-    "9": "B-MENU.UNITPRICE",
-    "10": "B-SUB_TOTAL.DISCOUNT_PRICE",
-    "11": "B-SUB_TOTAL.ETC",
-    "12": "B-SUB_TOTAL.SERVICE_PRICE",
-    "13": "B-SUB_TOTAL.SUBTOTAL_PRICE",
-    "14": "B-SUB_TOTAL.TAX_PRICE",
-    "15": "B-TOTAL.CASHPRICE",
-    "16": "B-TOTAL.CHANGEPRICE",
-    "17": "B-TOTAL.CREDITCARDPRICE",
-    "18": "B-TOTAL.MENUQTY_CNT",
-    "19": "B-TOTAL.TOTAL_PRICE",
-    "20": "I-MENU.CNT",
-    "21": "I-MENU.DISCOUNTPRICE",
-    "22": "I-MENU.NM",
-    "23": "I-MENU.NUM",
-    "24": "I-MENU.PRICE",
-    "25": "I-MENU.SUB.CNT",
-    "26": "I-MENU.SUB.NM",
-    "27": "I-MENU.SUB.PRICE",
-    "28": "I-MENU.UNITPRICE",
-    "29": "I-SUB_TOTAL.DISCOUNT_PRICE",
-    "30": "I-SUB_TOTAL.ETC",
-    "31": "I-SUB_TOTAL.SERVICE_PRICE",
-    "32": "I-SUB_TOTAL.SUBTOTAL_PRICE",
-    "33": "I-SUB_TOTAL.TAX_PRICE",
-    "34": "I-TOTAL.CASHPRICE",
-    "35": "I-TOTAL.CHANGEPRICE",
-    "36": "I-TOTAL.CREDITCARDPRICE",
-    "37": "I-TOTAL.MENUQTY_CNT",
-    "38": "I-TOTAL.TOTAL_PRICE"
-  }
-
-label2id = {
-    "B-MENU.CNT": 1,
-    "B-MENU.DISCOUNTPRICE": 2,
-    "B-MENU.NM": 3,
-    "B-MENU.NUM": 4,
-    "B-MENU.PRICE": 5,
-    "B-MENU.SUB.CNT": 6,
-    "B-MENU.SUB.NM": 7,
-    "B-MENU.SUB.PRICE": 8,
-    "B-MENU.UNITPRICE": 9,
-    "B-SUB_TOTAL.DISCOUNT_PRICE": 10,
-    "B-SUB_TOTAL.ETC": 11,
-    "B-SUB_TOTAL.SERVICE_PRICE": 12,
-    "B-SUB_TOTAL.SUBTOTAL_PRICE": 13,
-    "B-SUB_TOTAL.TAX_PRICE": 14,
-    "B-TOTAL.CASHPRICE": 15,
-    "B-TOTAL.CHANGEPRICE": 16,
-    "B-TOTAL.CREDITCARDPRICE": 17,
-    "B-TOTAL.MENUQTY_CNT": 18,
-    "B-TOTAL.TOTAL_PRICE": 19,
-    "I-MENU.CNT": 20,
-    "I-MENU.DISCOUNTPRICE": 21,
-    "I-MENU.NM": 22,
-    "I-MENU.NUM": 23,
-    "I-MENU.PRICE": 24,
-    "I-MENU.SUB.CNT": 25,
-    "I-MENU.SUB.NM": 26,
-    "I-MENU.SUB.PRICE": 27,
-    "I-MENU.UNITPRICE": 28,
-    "I-SUB_TOTAL.DISCOUNT_PRICE": 29,
-    "I-SUB_TOTAL.ETC": 30,
-    "I-SUB_TOTAL.SERVICE_PRICE": 31,
-    "I-SUB_TOTAL.SUBTOTAL_PRICE": 32,
-    "I-SUB_TOTAL.TAX_PRICE": 33,
-    "I-TOTAL.CASHPRICE": 34,
-    "I-TOTAL.CHANGEPRICE": 35,
-    "I-TOTAL.CREDITCARDPRICE": 36,
-    "I-TOTAL.MENUQTY_CNT": 37,
-    "I-TOTAL.TOTAL_PRICE": 38,
-    "O": 0
-  }
+labels = [
+    "O",
+    "B-MENU.CNT",
+    "B-MENU.DISCOUNTPRICE",
+    "B-MENU.NM",
+    "B-MENU.NUM",
+    "B-MENU.PRICE",
+    "B-MENU.SUB.CNT",
+    "B-MENU.SUB.NM",
+    "B-MENU.SUB.PRICE",
+    "B-MENU.UNITPRICE",
+    "B-SUB_TOTAL.DISCOUNT_PRICE",
+    "B-SUB_TOTAL.ETC",
+    "B-SUB_TOTAL.SERVICE_PRICE",
+    "B-SUB_TOTAL.SUBTOTAL_PRICE",
+    "B-SUB_TOTAL.TAX_PRICE",
+    "B-TOTAL.CASHPRICE",
+    "B-TOTAL.CHANGEPRICE",
+    "B-TOTAL.CREDITCARDPRICE",
+    "B-TOTAL.MENUQTY_CNT",
+    "B-TOTAL.TOTAL_PRICE",
+    "I-MENU.CNT",
+    "I-MENU.DISCOUNTPRICE",
+    "I-MENU.NM",
+    "I-MENU.NUM",
+    "I-MENU.PRICE",
+    "I-MENU.SUB.CNT",
+    "I-MENU.SUB.NM",
+    "I-MENU.SUB.PRICE",
+    "I-MENU.UNITPRICE",
+    "I-SUB_TOTAL.DISCOUNT_PRICE",
+    "I-SUB_TOTAL.ETC",
+    "I-SUB_TOTAL.SERVICE_PRICE",
+    "I-SUB_TOTAL.SUBTOTAL_PRICE",
+    "I-SUB_TOTAL.TAX_PRICE",
+    "I-TOTAL.CASHPRICE",
+    "I-TOTAL.CHANGEPRICE",
+    "I-TOTAL.CREDITCARDPRICE",
+    "I-TOTAL.MENUQTY_CNT",
+    "I-TOTAL.TOTAL_PRICE"
+]
+
+id2label = dict(enumerate(label_list))
+label2id = {v: k for k, v in enumerate(label_list)}
 
 # nielsr/layoutlmv3-finetuned-cord
 tokenizer = LayoutLMv3TokenizerFast.from_pretrained("mp-02/layoutlmv3-finetuned-cord", apply_ocr=False)