Update app.py

app.py

@@ -14,18 +14,19 @@ mtcnn = MTCNN(margin=40, keep_all=True, post_process=False, device=device)
 learn = load_learner('export.pkl')
 
 labels = learn.dls.vocab
-
+e_colors = {'confused':'orange', 'attentive':'green','bored':'red','interested':'blue','frustrated':'purple','thoughtful':'pink'}
 emotions = {"x": [], "y": [], "State": []}
 
 def predict(img):
     img = PILImage.create(img)
     boxes, _ = mtcnn.detect(img)
     o_img = img.copy()
-    draw = ImageDraw.Draw(o_img)
+    #draw = ImageDraw.Draw(o_img)
+    draw = Image.new("RGBA", img.size, color = "white")
     for box in boxes:
         coords = tuple(box.tolist())
         pred,pred_idx,probs = learn.predict(img.crop(coords))
-        draw.rectangle(coords, outline=(0, 0, 0), width=1)
+        draw.rectangle(coords, fill=e_colors[pred], outline=(0, 0, 0), width=1)
         draw.text((coords[0]-10, coords[1]-10), pred, font=ImageFont.truetype("arial.ttf"))
     return o_img
     '''