Update app.py

app.py

@@ -1,48 +1,43 @@
 import os
 import numpy as np
-import onnxruntime as ort
+import torch
+import torch.nn.functional as F
 import gradio as gr
 from PIL import Image
 
-# モデルをONNX Runtimeで読み込む
-session = ort.InferenceSession("cxp_projection_rotation.pt", None)
-input_name = session.get_inputs()[0].name
-# 出力ヘッド名（順序に注意）
-projection_name = session.get_outputs()[0].name
-rotation_name = session.get_outputs()[1].name
+# 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"]
 
-def softmax(x: np.ndarray) -> np.ndarray:
-    """数値安定版ソフトマックス"""
-    e = np.exp(x - np.max(x))
-    return e / np.sum(e)
-
 def predict(image_path: str) -> str:
     # 画像をグレースケールで読み込み（Lモード）
     img = Image.open(image_path).convert("L")
-    arr = np.array(img, dtype=np.float32) / 255.0        # 正規化
-    arr = arr[np.newaxis, np.newaxis, :, :]              # [1,1,H,W] バッチ＆チャネル次元追加
+    arr = np.array(img, dtype=np.float32) / 255.0        # 0-1 正規化
+    # バッチ＆チャンネル次元を追加して Tensor 化
+    tensor = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0)  # [1,1,256,256]
 
     # 推論
-    proj_logits, rot_logits = session.run(
-        [projection_name, rotation_name],
-        {input_name: arr}
-    )
-    proj_probs = softmax(proj_logits[0])
-    rot_probs  = softmax(rot_logits[0])
+    with torch.no_grad():
+        proj_logits, rot_logits = model(tensor)  # 2 ヘッド出力
 
-    # 予測結果と確率
-    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]
+    # ソフトマックスで確率化
+    proj_probs = F.softmax(proj_logits, dim=1)[0].cpu().numpy()
+    rot_probs  = F.softmax(rot_logits,  dim=1)[0].cpu().numpy()
 
-    # ファイル名から元ラベルを推定（例: "AP_Upright.png" → ["AP","Upright"]）
+    # 最も確率の高いラベルと確率
+    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]
+
+    # ファイル名から元ラベルを取得（例: "AP_Upright.png" → ["AP","Upright"]）
     base = os.path.splitext(os.path.basename(image_path))[0]
     try:
         orig_proj, orig_rot = base.split("_", 1)
@@ -56,7 +51,7 @@ def predict(image_path: str) -> str:
     if orig_rot and orig_rot != rot_lbl:
         warnings_html += "<p style='color:red'>⚠ Potential mislabeled rotation</p>"
 
-    # HTML形式で結果を返す
+    # 結果を 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>"
@@ -64,35 +59,33 @@ def predict(image_path: str) -> str:
     )
     return html
 
-# Gradio UI構築
+# Gradio UI 定義
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
-            # 画像アップロード（PNG/L, 256×256前提）
             image_input = gr.Image(
                 label="Upload PNG (256×256)", 
                 type="filepath", 
                 tool=None
             )
-            # sample_imagesフォルダから4枚まで例示
-            sample_list = sorted(
-                [os.path.join("sample_images", f) 
-                 for f in os.listdir("sample_images") 
-                 if f.lower().endswith(".png")]
-            )
+            # sample_images フォルダからサンプル画像を読み込む
+            sample_list = 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, 
+                examples=sample_list,
+                inputs=image_input,
                 label="Sample Images"
             )
         with gr.Column():
-            # 推論結果用HTML表示エリア
             result = gr.HTML()
 
-    # 画像選択・アップロード時に自動でpredictを実行
+    # 画像を選択／アップロードしたら自動で predict を実行
     image_input.change(
-        fn=predict, 
-        inputs=image_input, 
+        fn=predict,
+        inputs=image_input,
         outputs=result
     )