Update cord_inference.py

cord_inference.py

@@ -4,55 +4,14 @@ from transformers import LayoutLMv3TokenizerFast, LayoutLMv3Processor, LayoutLMv
 from PIL import Image, ImageDraw, ImageFont
 from utils import OCR, unnormalize_box
 
-label_list = [
-    "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)}
 
 tokenizer = LayoutLMv3TokenizerFast.from_pretrained("mp-02/layoutlmv3-finetuned-cord", apply_ocr=False)
 processor = LayoutLMv3Processor.from_pretrained("mp-02/layoutlmv3-finetuned-cord", apply_ocr=False)
 model = LayoutLMv3ForTokenClassification.from_pretrained("mp-02/layoutlmv3-finetuned-cord")
 
+id2label = model.config.id2label
+label2id = model.config.label2id
+
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model.to(device)