Update app.py

app.py

@@ -1,157 +1,76 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
 import os
+import json
+from PIL import Image
 import torch
-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
-
-# ===========================================
-# 1) パス設定
-# ===========================================
-WEIGHT_PATH   = "./cxp_projection_rotation.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"),
-    )
+import gradio as gr
 
-args_model, args_dataloader = load_parameter(PARAMETER_JSON)
+# Load parameters and model
+with open("parameters.json", "r") as f:
+    parameters = json.load(f)
 
-# ===========================================
-# 5) モデル作成＆重みロード
-# ===========================================
-model = create_model(args_model)
-print(f"Loading weights from {WEIGHT_PATH}")
-model.load_weight(WEIGHT_PATH)
+model = create_model(parameters)  # your existing create_model function
+weights = torch.load("cxp_projection_rotation.pt", map_location="cpu")
+model.load_state_dict(weights)
 model.eval()
 
-# ===========================================
-# 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)
+# Transformation for grayscale images
+transform = T.Compose([
+    T.ToTensor(),  # converts [H,W] to [1,H,W]
+])
 
+# 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():
-        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]
-
-    # ファイル名から元ラベル取得（例: "AP_Upright.png"）
-    base = os.path.splitext(os.path.basename(image_path))[0]
-    try:
-        orig_proj, orig_rot = base.split("_", 1)
-    except ValueError:
-        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}"
-    )
+        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("_")
+    if len(parts) >= 3:
+        orig_proj = parts[1]
+        orig_rot = parts[2]
+        orig_label = f"{orig_proj}_{orig_rot}"
+    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>"
     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:
+    gr.Markdown("## Chest X-ray Projection and Rotation Classifier")
     with gr.Row():
-        input_image = gr.Image(
-            label="Upload PNG (256×256)", 
-            type="filepath", 
-            image_mode="L"
-        )
-        output_html = gr.HTML()
+        img_input = gr.Image(type="filepath", image_mode="L", label="Upload PNG Image (256×256)")
+        html_output = gr.HTML(label="Result")
 
-    send_btn = gr.Button("Run Inference")
-    send_btn.click(
-        fn=predict_html,
-        inputs=input_image,
-        outputs=output_html
-    )
+    classify_btn = gr.Button("Classify Image")
+    classify_btn.click(fn=predict_html, inputs=img_input, outputs=html_output)
 
-    # サンプル例
+    # 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/1_AP_Upright.png",
-            "./sample/1_PA_Inverted.png",
-            "./sample/2_AP_Right90.png",
-            "./sample/2_Lateral_Left90.png.png",
-        ],
-        inputs=input_image
+        examples=sample_paths,
+        inputs=img_input,
+        outputs=html_output,
+        fn=predict_html,
+        label="Sample Images"
     )
 
 if __name__ == "__main__":