Update app.py

app.py

@@ -1,80 +1,169 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
 import os
-import json
-from PIL import Image
 import torch
-import torchvision.transforms as T
+import torch.nn.functional as F
 import gradio as gr
+import numpy as np
+import torchvision.transforms as T
+from PIL import Image
+
+from lib.framework import create_model
+from lib.options import ParamSet, _retrieve_parameter, _dispatch_by_group
+from lib.dataloader import ImageMixin
 
-# Import your model creation function
-from model import create_model  # ここで create_model を定義しているファイルを指定してください
+# ===========================================
+# 1) パス設定
+# ===========================================
+WEIGHT_PATH   = "./cxp_projection_rotation.pt"
+PARAMETER_JSON = "./parameters.json"
 
-# Load parameters and model
-with open("parameters.json", "r") as f:
-    parameters = json.load(f)
+# ===========================================
+# 2) クラスラベル定義
+# ===========================================
+LABEL_APorPA = ["AP", "PA", "Lateral"]
+LABEL_ROUND  = ["Upright", "Inverted", "Left rotation", "Right rotation"]
 
-model = create_model(parameters)
-weights = torch.load("cxp_projection_rotation.pt", map_location="cpu")
-model.load_state_dict(weights)
+# ===========================================
+# 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()
 
-# Transformation for grayscale images
-transform = T.Compose([
-    T.ToTensor(),  # converts [H,W] to [1,H,W]
-])
+# ===========================================
+# 6) 推論＋HTML生成
+# ===========================================
+def predict_html(image_path: str) -> str:
+    # 画像読み込み
+    img = Image.open(image_path).convert("L")
+    handler = ImageHandler(args_dataloader)
+    batch = handler.set_image(img)
 
-# Prediction and HTML rendering
-def predict_html(image_path):
-    # Preprocess and infer
-    img = Image.open(image_path)
-    x = transform(img).unsqueeze(0)
     with torch.no_grad():
-        proj_logits, rot_logits = model(x)
-    proj_idx = proj_logits.argmax(dim=1).item()
-    rot_idx = rot_logits.argmax(dim=1).item()
-    proj_pred = parameters["projection_labels"][proj_idx]
-    rot_pred = parameters["rotation_labels"][rot_idx]
-
-    # Parse file name: ID_Projection_Rotation.png
-    filename = os.path.basename(image_path)
-    name, _ = os.path.splitext(filename)
-    parts = name.split("_")
+        outputs = model(batch)
+        # raw logits
+        logits_proj = outputs.get("label_APorPA")
+        logits_rot  = outputs.get("label_round")
+
+    # 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]
+
+    # ファイル名から元ラベル取得（例: "1_AP_Upright.png" → orig_proj="AP", orig_rot="Upright"）
+    base = os.path.splitext(os.path.basename(image_path))[0]
+    parts = base.split("_")
     if len(parts) >= 3:
         orig_proj = parts[1]
-        orig_rot = parts[2]
-        orig_label = f"{orig_proj}_{orig_rot}"
+        orig_rot  = parts[2]
     else:
-        orig_label = None
-
-    # Build HTML output
-    html = f"<h3>Prediction: {proj_pred} / {rot_pred}</h3>"
-    if orig_label:
-        if orig_label != f"{proj_pred}_{rot_pred}":
-            html += f"<p style='color:red;'>Warning: original label '<strong>{orig_label}</strong>' does not match prediction.</p>"
-        else:
-            html += f"<p>Original label '<strong>{orig_label}</strong>' matches prediction.</p>"
+        orig_proj = orig_rot = None
+
+    # 警告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>"
+
+    # 結果表示用HTML
+    html = (
+        f"<p><strong>Projection :</strong> {pred_proj}</p>"
+        f"<p><strong>Rotation   :</strong> {pred_rot}</p>"
+        f"{warn_html}"
+    )
     return html
 
-# Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("## Chest X-ray Projection and Rotation Classifier")
+# ===========================================
+# 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)
+
     with gr.Row():
-        img_input = gr.Image(type="filepath", image_mode="L", label="Upload PNG Image (256×256)")
-        html_output = gr.HTML(label="Result")
+        input_image = gr.Image(
+            label="Upload PNG (256×256)", 
+            type="filepath", 
+            image_mode="L"
+        )
+        output_html = gr.HTML()
 
-    classify_btn = gr.Button("Classify Image")
-    classify_btn.click(fn=predict_html, inputs=img_input, outputs=html_output)
+    send_btn = gr.Button("Run Inference")
+    send_btn.click(
+        fn=predict_html,
+        inputs=input_image,
+        outputs=output_html
+    )
 
-    # Sample images with filenames shown
-    sample_dir = "samples"
-    sample_files = sorted([f for f in os.listdir(sample_dir) if f.endswith('.png')])
-    sample_paths = [os.path.join(sample_dir, f) for f in sample_files]
+    # サンプル例
     gr.Examples(
-        examples=sample_paths,
-        inputs=img_input,
-        outputs=html_output,
-        fn=predict_html,
-        label="Sample Images"
+        examples=[
+            "./sample/1_AP_Upright.png",
+            "./sample/1_PA_Inverted.png",
+            "./sample/2_AP_Right90.png",
+            "./sample/2_Lateral_Left90.png.png",
+        ],
+        inputs=input_image
+    )
+
+    # サンプルのファイル名を一覧で表示
+    gr.Markdown(
+        "**Sample filenames:**  \n"
+        "- 1_AP_Upright.png  \n"
+        "- 1_PA_Inverted.png  \n"
+        "- 2_AP_Right90.png  \n"
+        "- 2_Lateral_Left90.png"
     )
 
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    demo.launch(debug=True)