Update app.py

app.py

@@ -1,93 +1,95 @@
 import os
 import numpy as np
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
+import timm
 import gradio as gr
 from PIL import Image
 
-# PyTorch モデルをロード
-# map_location="cpu" で CPU 上にロード
-model = torch.load("cxp_projection_rotation.pt", map_location="cpu")
-model.eval()  # 評価モードに切り替え
-
-# クラスラベル定義
+# --- クラスラベル定義 ---
 PROJ_LABELS = ["AP", "PA", "Lateral"]
 ROT_LABELS  = ["Upright", "Inverted", "Left90", "Right90"]
 
+# --- モデル定義: EfficientNet backbone + 2-head ---
+class CXRModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        # 1ch 入力, 分類器を持たない backbone
+        self.backbone = timm.create_model(
+            "efficientnet_b0", pretrained=False, in_chans=1, num_classes=0
+        )
+        nf = self.backbone.num_features
+        # 2つの出力ヘッド
+        self.projection_head = nn.Linear(nf, len(PROJ_LABELS))
+        self.rotation_head   = nn.Linear(nf, len(ROT_LABELS))
+
+    def forward(self, x):
+        # 特徴量抽出
+        feats = self.backbone.forward_features(x)
+        # global pool → flatten
+        pooled = self.backbone.global_pool(feats).flatten(1)
+        # 2ヘッド出力
+        return self.projection_head(pooled), self.rotation_head(pooled)
+
+# --- モデル・重みロード ---
+model = CXRModel()
+state_dict = torch.load("cxp_projection_rotation.pt", map_location="cpu")
+model.load_state_dict(state_dict)   # OrderedDict → モデルに流し込む
+model.eval()                        # 評価モード
+
+# --- 推論 & HTML 結果生成関数 ---
 def predict(image_path: str) -> str:
-    # 画像をグレースケールで読み込み（Lモード）
     img = Image.open(image_path).convert("L")
-    arr = np.array(img, dtype=np.float32) / 255.0        # 0-1 正規化
-    # バッチ＆チャンネル次元を追加して Tensor 化
+    arr = np.array(img, dtype=np.float32) / 255.0
     tensor = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0)  # [1,1,256,256]
 
-    # 推論
     with torch.no_grad():
-        proj_logits, rot_logits = model(tensor)  # 2 ヘッド出力
+        proj_logits, rot_logits = model(tensor)
 
-    # ソフトマックスで確率化
+    # softmax で確率化
     proj_probs = F.softmax(proj_logits, dim=1)[0].cpu().numpy()
     rot_probs  = F.softmax(rot_logits,  dim=1)[0].cpu().numpy()
 
-    # 最も確率の高いラベルと確率
-    proj_idx = int(np.argmax(proj_probs))
-    rot_idx  = int(np.argmax(rot_probs))
-    proj_lbl = PROJ_LABELS[proj_idx]
-    rot_lbl  = ROT_LABELS[rot_idx]
-    proj_p   = proj_probs[proj_idx]
-    rot_p    = rot_probs[rot_idx]
+    # ベストラベル＋確率
+    pi = int(np.argmax(proj_probs));  ri = int(np.argmax(rot_probs))
+    pl, rl = PROJ_LABELS[pi], ROT_LABELS[ri]
+    pp, rp = proj_probs[pi],  rot_probs[ri]
 
-    # ファイル名から元ラベルを取得（例: "AP_Upright.png" → ["AP","Upright"]）
+    # 元ラベル推定（ファイル名から）
     base = os.path.splitext(os.path.basename(image_path))[0]
     try:
-        orig_proj, orig_rot = base.split("_", 1)
+        orig_p, orig_r = base.split("_", 1)
     except ValueError:
-        orig_proj = orig_rot = None
+        orig_p = orig_r = None
 
-    # 警告メッセージ（必要なら赤字で表示）
-    warnings_html = ""
-    if orig_proj and orig_proj != proj_lbl:
-        warnings_html += "<p style='color:red'>⚠ Potential mislabeled projection</p>"
-    if orig_rot and orig_rot != rot_lbl:
-        warnings_html += "<p style='color:red'>⚠ Potential mislabeled rotation</p>"
+    # 警告生成
+    warn_html = ""
+    if orig_p and orig_p != pl:
+        warn_html += "<p style='color:red'>⚠ Potential mislabeled projection</p>"
+    if orig_r and orig_r != rl:
+        warn_html += "<p style='color:red'>⚠ Potential mislabeled rotation</p>"
 
-    # 結果を HTML 形式で返す
-    html = (
-        f"<p><strong>Projection :</strong> {proj_lbl} (p={proj_p:.3f})</p>"
-        f"<p><strong>Rotation   :</strong> {rot_lbl} (p={rot_p:.3f})</p>"
-        f"{warnings_html}"
+    return (
+        f"<p><strong>Projection :</strong> {pl} (p={pp:.3f})</p>"
+        f"<p><strong>Rotation   :</strong> {rl} (p={rp:.3f})</p>"
+        f"{warn_html}"
     )
-    return html
 
-# Gradio UI 定義
+# --- Gradio UI ---
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(
-                label="Upload PNG (256×256)", 
-                type="filepath", 
-                tool=None
-            )
-            # sample_images フォルダからサンプル画像を読み込む
-            sample_list = sorted([
+            img_in = gr.Image(label="Upload PNG (256×256)", type="filepath", tool=None)
+            samples = sorted(
                 os.path.join("sample_images", f)
-                for f in os.listdir("sample_images")
-                if f.lower().endswith(".png")
-            ])
-            gr.Examples(
-                examples=sample_list,
-                inputs=image_input,
-                label="Sample Images"
+                for f in os.listdir("sample_images") if f.lower().endswith(".png")
             )
+            gr.Examples(examples=samples, inputs=img_in, label="Sample Images")
         with gr.Column():
-            result = gr.HTML()
+            output = gr.HTML()
 
-    # 画像を選択／アップロードしたら自動で predict を実行
-    image_input.change(
-        fn=predict,
-        inputs=image_input,
-        outputs=result
-    )
+    img_in.change(fn=predict, inputs=img_in, outputs=output)
 
 if __name__ == "__main__":
     demo.launch()