Update app.py

app.py

@@ -35,11 +35,11 @@ labels = ['O', 'B-HEADER', 'I-HEADER', 'B-QUESTION', 'I-QUESTION', 'B-ANSWER', '
 id2label = {v: k for v, k in enumerate(labels)}
 label2color = {
     "B-HEADER": 'red',
-    "I-HEADER": 'green',
-    "B-QUESTION": 'orange',
-    "I-QUESTION": "blue",
-    "B-ANSWER": 'gray',
-    "I-ANSWERE": 'violet',
+    "I-HEADER": 'red',
+    "B-QUESTION": 'red',
+    "I-QUESTION": "red",
+    "B-ANSWER": 'blue',
+    "I-ANSWERE": 'blue',
     "O": 'orange'
   } 
 
@@ -83,9 +83,9 @@ def process_image(image):
     font = ImageFont.load_default()
     for prediction, box in zip(true_predictions, true_boxes):
         predicted_label = iob_to_label(prediction)
-        draw.rectangle(box, outline='red') #label2color[predicted_label]
+        draw.rectangle(box, outline=label2color[predicted_label]) #label2color[predicted_label]
         print(predicted_label)
-        draw.text((box[0]+10, box[1]-10), text=predicted_label, fill='red', font=font) #label2color[predicted_label]
+        draw.text((box[0]+10, box[1]-10), text=predicted_label, fill=label2color[predicted_label], font=font) #label2color[predicted_label]
     
     return image