Add certainty metric to UI

app.py

@@ -75,7 +75,7 @@ def process_image_str(groupImageFilePath: str):
     groupImage_d = ImageDraw.ImageDraw(groupImage)
     font = ImageFont.truetype("Arial Bold.ttf", 30)
     
-    labels = ""
+    labels = "| | Name | Certainty | \n | -------- | ------- | ------- |\n"
     n = 1
     
     for image, location in zip(images, locations):
@@ -85,22 +85,22 @@ def process_image_str(groupImageFilePath: str):
         # do AI stuff here
         with torch.no_grad():
             outputs_t = model(intputTensor)
-            _, pred_t = torch.max(outputs_t, dim=1)
+            cert, pred_t = torch.max(torch.softmax(outputs_t, dim=1), dim=1)
             groupImage_d.rectangle(location, outline=(0, 255, 0), width=2)
             groupImage_d.text((location[0] + 4, location[1] + 2), f"{n}", fill=(0, 255, 0), font=font)
-            labels += f"{n}: {celeb_list[pred_t.item()]}\n"
+            labels += f"| {n} | {celeb_list[pred_t.item()]} | {int(cert.item() * 100)}% | \n"
         
         n += 1
     
-    return [gr.Image(groupImage), gr.Text(labels)]
+    return [gr.Image(groupImage), gr.Markdown(labels)]
     # return gr.Image(images[0])
 
 
 # 5. Create the Gradio interface
 interface = gr.Interface(
     fn=process_image_str,  # Function to process the image
-    inputs=gr.Image(type='filepath'),  # Upload input
-    outputs=[gr.Image(), gr.Text()],  # Display output
+    inputs=gr.Image(type='filepath', label="Input Image"),  # Upload input
+    outputs=[gr.Image(label="Output"), gr.Markdown(label="Output Legend")],  # Display output
     allow_flagging='never',
     title="Celebrity Face Detector",
     description="Upload a picture of a celebrity or group of celebrities to identify them"