Update app.py

app.py

@@ -1,95 +1,157 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
 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
+import numpy as np
+import torchvision.transforms as T
 from PIL import Image
 
-# --- クラスラベル定義 ---
-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:
+from lib.framework import create_model
+from lib.options import ParamSet, _retrieve_parameter, _dispatch_by_group
+from lib.dataloader import ImageMixin
+
+# ===========================================
+# 1) パス設定
+# ===========================================
+WEIGHT_PATH   = "./weight_epoch-011_best.pt"
+PARAMETER_JSON = "./parameters.json"
+
+# ===========================================
+# 2) クラスラベル定義
+# ===========================================
+LABEL_APorPA = ["AP", "PA", "Lateral"]
+LABEL_ROUND  = ["Upright", "Inverted", "Left rotation", "Right rotation"]
+
+# ===========================================
+# 3) 前処理クラス
+# ===========================================
+class ImageHandler(ImageMixin):
+    def __init__(self, params):
+        self.params = params
+        self.transform = T.Compose([
+            # 256×256 前提なら Resize は不要
+            # T.Resize((256, 256)),
+            T.ToTensor(),
+        ])
+
+    def set_image(self, image: Image.Image):
+        tensor = self.transform(image)            # [C,H,W], float32 in [0,1]
+        return {"image": tensor.unsqueeze(0)}     # バッチ次元追加
+
+# ===========================================
+# 4) パラメータロード
+# ===========================================
+def load_parameter(parameter_path):
+    _args = ParamSet()
+    params = _retrieve_parameter(parameter_path)
+    for k, v in params.items():
+        setattr(_args, k, v)
+    # 推論用に上書き
+    _args.augmentation = "no"
+    _args.sampler      = "no"
+    _args.pretrained   = False
+    _args.mlp          = None
+    _args.net          = _args.model
+    _args.device       = torch.device("cpu")
+    return (
+        _dispatch_by_group(_args, "model"),
+        _dispatch_by_group(_args, "dataloader"),
+    )
+
+args_model, args_dataloader = load_parameter(PARAMETER_JSON)
+
+# ===========================================
+# 5) モデル作成＆重みロード
+# ===========================================
+model = create_model(args_model)
+print(f"Loading weights from {WEIGHT_PATH}")
+model.load_weight(WEIGHT_PATH)
+model.eval()
+
+# ===========================================
+# 6) 推論＋HTML生成
+# ===========================================
+def predict_html(image_path: str) -> str:
+    # 画像読み込み
     img = Image.open(image_path).convert("L")
-    arr = np.array(img, dtype=np.float32) / 255.0
-    tensor = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0)  # [1,1,256,256]
+    handler = ImageHandler(args_dataloader)
+    batch = handler.set_image(img)
 
     with torch.no_grad():
-        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()
+        outputs = model(batch)
+        # raw logits
+        logits_proj = outputs.get("label_APorPA")
+        logits_rot  = outputs.get("label_round")
 
-    # ベストラベル＋確率
-    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]
+    # argmax でラベル選択
+    idx_proj = int(torch.argmax(logits_proj, dim=1).item())
+    idx_rot  = int(torch.argmax(logits_rot,  dim=1).item())
+    pred_proj = LABEL_APorPA[idx_proj]
+    pred_rot  = LABEL_ROUND[idx_rot]
 
-    # 元ラベル推定（ファイル名から）
+    # ファイ���名から元ラベル取得（例: "AP_Upright.png"）
     base = os.path.splitext(os.path.basename(image_path))[0]
     try:
-        orig_p, orig_r = base.split("_", 1)
+        orig_proj, orig_rot = base.split("_", 1)
     except ValueError:
-        orig_p = orig_r = None
+        orig_proj = orig_rot = None
 
-    # 警告生成
+    # 警告HTML
     warn_html = ""
-    if orig_p and orig_p != pl:
+    if orig_proj and orig_proj != pred_proj:
         warn_html += "<p style='color:red'>⚠ Potential mislabeled projection</p>"
-    if orig_r and orig_r != rl:
+    if orig_rot and orig_rot != pred_rot:
         warn_html += "<p style='color:red'>⚠ Potential mislabeled rotation</p>"
 
-    return (
-        f"<p><strong>Projection :</strong> {pl} (p={pp:.3f})</p>"
-        f"<p><strong>Rotation   :</strong> {rl} (p={rp:.3f})</p>"
+    # 結果表示用HTML
+    html = (
+        f"<p><strong>Projection :</strong> {pred_proj}</p>"
+        f"<p><strong>Rotation   :</strong> {pred_rot}</p>"
         f"{warn_html}"
     )
+    return html
+
+# ===========================================
+# 7) Gradio UI
+# ===========================================
+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>
+</div>
+"""
+
+with gr.Blocks(title="CXR Projection & Rotation") as demo:
+    gr.HTML(html_header)
 
-# --- Gradio UI ---
-with gr.Blocks() as demo:
     with gr.Row():
-        with gr.Column():
-            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=samples, inputs=img_in, label="Sample Images")
-        with gr.Column():
-            output = gr.HTML()
-
-    img_in.change(fn=predict, inputs=img_in, outputs=output)
+        input_image = gr.Image(
+            label="Upload PNG (256×256)", 
+            type="filepath", 
+            image_mode="L"
+        )
+        output_html = gr.HTML()
+
+    send_btn = gr.Button("Run Inference")
+    send_btn.click(
+        fn=predict_html,
+        inputs=input_image,
+        outputs=output_html
+    )
+
+    # サンプル例
+    gr.Examples(
+        examples=[
+            "./sample/sample_AP_Upright.png",
+            "./sample/sample_PA_Inverted.png",
+            "./sample/sample_Lateral_Right rotation.png",
+        ],
+        inputs=input_image
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(debug=True)