Update app.py

app.py

@@ -16,7 +16,7 @@ from lib.dataloader import ImageMixin
 # ===========================================
 # 1) パス設定
 # ===========================================
-WEIGHT_PATH   = "./cxp_projection_rotation.pt"
+WEIGHT_PATH    = "./cxp_projection_rotation.pt"
 PARAMETER_JSON = "./parameters.json"
 
 # ===========================================
@@ -32,8 +32,6 @@ class ImageHandler(ImageMixin):
     def __init__(self, params):
         self.params = params
         self.transform = T.Compose([
-            # 256×256 前提なら Resize は不要
-            # T.Resize((256, 256)),
             T.ToTensor(),
         ])
 
@@ -82,15 +80,20 @@ def predict_html(image_path: str) -> str:
 
     with torch.no_grad():
         outputs = model(batch)
-        # raw logits
         logits_proj = outputs.get("label_APorPA")
         logits_rot  = outputs.get("label_round")
 
+        # softmax で確率に変換
+        probs_proj = F.softmax(logits_proj, dim=1)[0].cpu().numpy()
+        probs_rot  = F.softmax(logits_rot,  dim=1)[0].cpu().numpy()
+
     # argmax でラベル選択
-    idx_proj = int(torch.argmax(logits_proj, dim=1).item())
-    idx_rot  = int(torch.argmax(logits_rot,  dim=1).item())
+    idx_proj = int(probs_proj.argmax())
+    idx_rot  = int(probs_rot.argmax())
     pred_proj = LABEL_APorPA[idx_proj]
     pred_rot  = LABEL_ROUND[idx_rot]
+    conf_proj = float(probs_proj[idx_proj])
+    conf_rot  = float(probs_rot[idx_rot])
 
     # ファイル名から元ラベル取得（例: "1_AP_Upright.png" → orig_proj="AP", orig_rot="Upright"）
     base = os.path.splitext(os.path.basename(image_path))[0]
@@ -101,17 +104,48 @@ def predict_html(image_path: str) -> str:
     else:
         orig_proj = orig_rot = None
 
+    # 警告HTML作成用ヘルパー
+    def make_warning(kind, orig, pred, conf):
+        # kind: "projection" or "rotation"
+        high_thr = 0.8
+        med_thr  = 0.5
+        if orig and orig != pred:
+            if conf >= high_thr:
+                return (
+                    f"<p style='color:red'>⚠ Potentially mislabeled {kind}: "
+                    f"filename says {orig}, model predicts {pred} (confidence {conf:.2f})</p>"
+                )
+            elif conf >= med_thr:
+                return (
+                    f"<p style='color:orange'>⚠ There is a possibility of mislabeled {kind}: "
+                    f"model predicts {pred} with moderate confidence ({conf:.2f})</p>"
+                )
+        if conf < med_thr:
+            return (
+                f"<p style='color:orange'>⚠ Low confidence for {kind} ({conf:.2f}); "
+                f"please check image quality or framing.</p>"
+            )
+        return ""
+
     # 警告HTML
     warn_html = ""
-    if orig_proj and orig_proj != pred_proj:
-        warn_html += "<p style='color:red'>⚠ Potential mislabeled projection</p>"
-    if orig_rot and orig_rot != pred_rot:
-        warn_html += "<p style='color:red'>⚠ Potential mislabeled rotation</p>"
+    warn_html += make_warning("projection", orig_proj, pred_proj, conf_proj)
+    warn_html += make_warning("rotation",   orig_rot,  pred_rot,  conf_rot)
+
+    # クラスごとのスコア表示用HTML
+    scores_proj = ", ".join(
+        f"{LABEL_APorPA[i]}: {p:.2f}" for i, p in enumerate(probs_proj)
+    )
+    scores_rot = ", ".join(
+        f"{LABEL_ROUND[i]}: {p:.2f}" for i, p in enumerate(probs_rot)
+    )
 
     # 結果表示用HTML
     html = (
-        f"<p><strong>Projection :</strong> {pred_proj}</p>"
-        f"<p><strong>Rotation   :</strong> {pred_rot}</p>"
+        f"<p><strong>Projection :</strong> {pred_proj} "
+        f"<small>({scores_proj})</small></p>"
+        f"<p><strong>Rotation   :</strong> {pred_rot} "
+        f"<small>({scores_rot})</small></p>"
         f"{warn_html}"
     )
     return html
@@ -123,7 +157,8 @@ html_header = """
 <div style="padding:10px;border:1px solid #ddd;border-radius:5px">
   <h2>Chest X‑ray Projection & Rotation Classification</h2>
   <p>Upload a 256×256 grayscale PNG. The model predicts projection (AP/PA/Lateral)
-     and rotation (Upright/Inverted/Left/Right) and warns if filename label differs.</p>
+     and rotation (Upright/Inverted/Left/Right) and shows softmax confidences. 
+     It warns if filename label differs or if confidence is low.</p>
 </div>
 """
 
@@ -158,7 +193,7 @@ with gr.Blocks(title="CXR Projection & Rotation") as demo:
 
     # サンプルのファイル名を一覧で表示
     gr.Markdown(
-        "**Sample filenames:**  \n"
+        "**Sample filenames:**  𝚮\n"
         "- 1_AP_Upright.png  \n"
         "- 1_PA_Inverted.png  \n"
         "- 2_AP_Right90.png  \n"
@@ -166,4 +201,4 @@ with gr.Blocks(title="CXR Projection & Rotation") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    demo.launch(debug=True)