Update app.py

app.py

@@ -38,7 +38,7 @@ def convert_l2n_n2l(dataset):
     return label_list, id2label, label2id, len(label_list)
 
 def label_colour(label):
-    label2color = {'MENU.PRICE':'blue', 'MENU.NM':'green', 'other':'green','MENU.TOTAL_PRICE':'red'}
+    label2color ={'MENU.PRICE':'blue','MENU.NM':'red','other':None,'MENU.NUM':'orange','TOTAL.TOTAL_PRICE':'green'}
     if label in label2color:
         colour = label2color.get(label)
     else:
@@ -109,7 +109,15 @@ def predict(image):
     words = text_extraction(image)
     extracted_words = convert_results(words,true_predictions)
     
-    return image,extracted_words
+    menu_list = []
+    price_list = []
+    for idx,item in enumerate(final_result):
+        if item[1] == 'MENU.NM':
+            menu_list.append(f"item {idx}.{item[0]}")
+        if item[1] == 'MENU.PRICE':
+            price_list.append(f"item {idx}. ${item[0]}")
+    
+    return image,menu_list,price_list
     
 def text_extraction(image):
     feature_extractor = LayoutLMv3FeatureExtractor()
@@ -120,7 +128,7 @@ css = """.output_image, .input_image {height: 600px !important}"""
 
 demo = gr.Interface(fn = predict,
                     inputs = gr.inputs.Image(type="pil"),
-                    outputs = [gr.outputs.Image(type="pil", label="annotated image"),'text'],
+                    outputs = [gr.outputs.Image(type="pil", label="annotated image"),'text','text'],
                     css = css,
                     examples = examples,
                     allow_flagging=True,