Update app.py

app.py

@@ -58,7 +58,6 @@ tfm = transforms.Compose([
     ),
 ])
 
-
 @torch.no_grad()
 def predict(image: Image.Image):
     if image is None:
@@ -68,11 +67,20 @@ def predict(image: Image.Image):
     x = tfm(image).unsqueeze(0).to(DEVICE)
 
     log_sp, log_st = model(x)
-    sp_id = int(log_sp.argmax(dim=1))
-    st_id = int(log_st.argmax(dim=1))
 
-    return SPECIES[sp_id], STATE[st_id]
+    prob_sp = torch.softmax(log_sp, dim=1)[0]  # [num_species]
+    prob_st = torch.softmax(log_st, dim=1)[0]  # [num_states]
+
+    sp_id = int(prob_sp.argmax().item())
+    st_id = int(prob_st.argmax().item())
+
+    sp_conf = float(prob_sp[sp_id].item())
+    st_conf = float(prob_st[st_id].item())
+
+    sp_text = f"{SPECIES[sp_id]}  (id={sp_id}, conf={sp_conf:.3f})"
+    st_text = f"{STATE[st_id]}  (id={st_id}, conf={st_conf:.3f})"
 
+    return sp_text, st_text
 
 demo = gr.Interface(
     fn=predict,