Update app.py

app.py

@@ -32,7 +32,7 @@ from datasets import Features, Sequence, ClassLabel, Value, Array2D, Array3D
 from transformers import LayoutLMv3ForTokenClassification,LayoutLMv3FeatureExtractor
 
 # define id2label
-id2label={0: 'container id', 1: 'seal number', 2: 'container quantity', 3: 'container type', 4: 'tare', 5: 'package quantity', 6: 'weight', 7: 'others'}
+id2label = {0: 'song name', 1: 'artist', 2: 'year', 3: 'album', 4: 'genres', 5: 'song writer', 6: 'lyrics', 7: 'others'}
 custom_config = r'--oem 3 --psm 6'
 # lang='eng+deu+ita+chi_sim'
 lang='spa'
@@ -62,19 +62,19 @@ def unnormalize_box(bbox, width, height):
 
 def iob_to_label(label):
   if label == 0:
-    return 'container id'
+    return 'song name'
   if label == 1:
-    return 'seal number'
+    return 'artist'
   if label == 2:
-    return 'container quantity'
+    return 'year'
   if label == 3:
-    return 'container type'
+    return 'album'
   if label == 4:
-    return 'tare'
+    return 'genres'
   if label == 5:
-    return 'package quantity'
+    return 'song writer'
   if label == 6:
-    return 'weight'
+    return 'lyrics'
   if label == 7:
     return 'others'
 
@@ -176,9 +176,10 @@ def visualize_image(final_bbox, final_preds, l_words, image):
       draw = ImageDraw.Draw(image)
       font = ImageFont.load_default()
 
-      label2color = {'container id':'red', 'seal number':'blue', 'container quantity':'black', 'container type':'green', 'tare':'brown', 'package quantity':'purple', 'weight':'orange', 'others': 'white'}
-      l2l = {'container id':'red', 'seal number':'blue', 'container quantity':'black', 'container type':'green', 'tare':'brown', 'package quantity':'purple', 'weight':'orange', 'others': 'white'}
-      f_labels = {'container id':'red', 'seal number':'blue', 'container quantity':'black', 'container type':'green', 'tare':'brown', 'package quantity':'purple', 'weight':'orange', 'others': 'white'}
+      label2color = {'song name':'red', 'artist':'blue', 'year':'black', 'album':'green', 'genres':'brown', 'song writer':'blue', 'lyrics':'purple', 'others': 'white'}
+      l2l = {'song name':'red', 'artist':'blue', 'year':'black', 'album':'green', 'genres':'brown', 'song writer':'blue','lyrics':'purple', 'others':'white'}
+      f_labels = {'song name':'red', 'artist':'blue', 'year':'black', 'album':'green', 'genres':'brown', 'song writer':'blue','lyrics':'purple', 'others':'white'}
+
 
       json_df = []