Update app.py

app.py

@@ -30,22 +30,31 @@ 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")
 # define id2label, label2color
-labels = dataset.features['ner_tags'].feature.names
-id2label = {v: k for v, k in enumerate(labels)}
-label2color = {
-    "B-ABN": 'blue',
-    "B-BILLER": 'blue',
-    "B-BILLER_ADDRESS": 'green',
-    "B-BILLER_POST_CODE": 'orange',
-    "B-DUE_DATE": "blue",
-    "B-GST": 'green',
-    "B-INVOICE_DATE": 'violet',
-    "B-INVOICE_NUMBER": 'orange',
-    "B-SUBTOTAL": 'green',
-    "B-TOTAL": 'blue',
-    "I-BILLER_ADDRESS": 'blue',
-    "O": 'orange'
-  } 
+#labels = dataset.features['ner_tags'].feature.names
+labels = ['O', 'HEADER', 'SUBHEADER', 'TEXTO', 'NUMERAL', 'RESUMEN']
+#id2label = {v: k for v, k in enumerate(labels)}
+id2label = {0: 'O', 1: 'HEADER', 2: 'SUBHEADER', 3: 'TEXTO', 4: 'NUMERAL', 5: 'RESUMEN'}
+label2color = {'O': 'violet',
+ 'HEADER': 'orange',
+ 'SUBHEADER': 'blue',
+ 'TEXTO': 'green',
+ 'NUMERAL': 'yellow',
+ 'RESUMEN': 'black',}
+
+#label2color = {
+#    "B-ABN": 'blue',
+#    "B-BILLER": 'blue',
+#    "B-BILLER_ADDRESS": 'green',
+#    "B-BILLER_POST_CODE": 'orange',
+#    "B-DUE_DATE": "blue",
+#    "B-GST": 'green',
+#    "B-INVOICE_DATE": 'violet',
+#    "B-INVOICE_NUMBER": 'orange',
+#    "B-SUBTOTAL": 'green',
+#    "B-TOTAL": 'blue',
+#    "I-BILLER_ADDRESS": 'blue',
+#    "O": 'orange'
+#  } 
 
 def unnormalize_box(bbox, width, height):
      return [