Update app.py

app.py

@@ -13,16 +13,18 @@ from lib.dataloader import ImageMixin
 # ===========================================
 # 1) パスなど（修正があれば適宜変更）
 # ===========================================
-test_weight = './weight_epoch-011_best.pt'
+test_weight = './weight_epoch-003_best.pt'
 parameter = './parameters.json'
 
 # ===========================================
 # 2) クラスラベルの定義
-#    - 撮像方向(label_APorPA)と回転方向(label_round)のクラス
+#    - 今回は撮像方向 (3クラス) と回転方向 (4クラス)
 # ===========================================
 LABEL_APorPA = [
-    "AP",  # class 0
-    "PA"   # class 1
+    "AP",        # class 0
+    "PA",        # class 1
+    "Lateral",   # class 2
+    # ↑ 実際のクラス名に合わせて変更してください
 ]
 
 LABEL_ROUND = [
@@ -39,15 +41,12 @@ LABEL_ROUND = [
 class ImageHandler(ImageMixin):
     def __init__(self, params):
         self.params = params
-        # ここでリサイズは省略（推論が重いので）
-        # 入力画像は既に256×256であることを想定
         self.transform = T.Compose([
-            # T.Resize((256, 256)),  # コメントアウト: 画像を256×256にリサイズ
-            T.ToTensor(),            # Tensor化 (0~1, shape: C,H,W)
+            # T.Resize((256, 256)),  # 必要ならコメントアウトを外す
+            T.ToTensor(),
         ])
 
     def set_image(self, image):
-        # PIL画像 -> transform -> バッチ次元を付ける
         image = self.transform(image)
         image = {'image': image.unsqueeze(0)}
         return image
@@ -89,8 +88,8 @@ model.eval()  # 推論モード
 def classify_APorPA_and_round(image):
     """
     モデルが以下を出力する想定:
-      outputs["label_APorPA"]  -> shape=[1, 2] (2クラス: AP/PA)
-      outputs["label_round"]   -> shape=[1, 4] (4クラス: 0°, 90°, 180°, 270°)
+      outputs["label_APorPA"]  -> shape=[1, 3] (3クラス)
+      outputs["label_round"]   -> shape=[1, 4] (4クラス)
     """
     image_handler = ImageHandler(args_dataloader)
     image_tensor = image_handler.set_image(image)
@@ -98,13 +97,22 @@ def classify_APorPA_and_round(image):
     with torch.no_grad():
         outputs = model(image_tensor)
 
+        # デバッグ用の出力チェック
+        print("keys in outputs =", outputs.keys())
+        if "label_APorPA" in outputs:
+            print("label_APorPA shape =", outputs["label_APorPA"].shape)
+        if "label_round" in outputs:
+            print("label_round shape =", outputs["label_round"].shape)
+
         # --- label_APorPA ---
         if "label_APorPA" not in outputs:
             print(f"[ERROR] 'label_APorPA' not found in outputs. Actual keys: {list(outputs.keys())}")
             return "ERROR: Missing 'label_APorPA'", "ERROR: Missing 'label_round'"
 
-        scores_APorPA = outputs["label_APorPA"]  # shape=[1,2]想定
+        scores_APorPA = outputs["label_APorPA"]  # shape=[1,3]想定
         pred_APorPA_idx = torch.argmax(scores_APorPA, dim=1).item()
+        
+        # IndexError が発生する場合は、pred_APorPA_idx が 0~2 の範囲外になる
         predicted_APorPA = LABEL_APorPA[pred_APorPA_idx]
 
         # --- label_round ---
@@ -125,7 +133,7 @@ html_content = """
 <div style="padding: 15px; border: 1px solid #e0e0e0; border-radius: 5px;">
     <h3>Chest X-ray: AP/PA & Rotation Classification</h3>
     <p>入力画像は既に256×256（グレースケール）であることを想定し、内部でのリサイズは行いません。</p>
-    <p>胸部レントゲン画像に対して、撮像方向(AP or PA)と回転方向(0°, 90°, 180°, 270°)を同時に推定します。</p>
+    <p>胸部レントゲン画像に対して、撮像方向(3クラス)と回転方向(4クラス)を同時に推定します。</p>
 </div>
 """
 
@@ -135,7 +143,7 @@ with gr.Blocks(title="Chest X-ray: AP/PA & Rotation Classification") as demo:
 
     with gr.Row():
         input_image = gr.Image(type="pil", image_mode="L")
-        output_APorPA = gr.Label(label="Predicted AP or PA")
+        output_APorPA = gr.Label(label="Predicted AP/PA/Lateral")  # 3クラス想定に合わせた名前
         output_round = gr.Label(label="Predicted Rotation")
 
     send_btn = gr.Button("Inference")
@@ -149,9 +157,9 @@ with gr.Blocks(title="Chest X-ray: AP/PA & Rotation Classification") as demo:
         # サンプルファイルは実際のパスに置き換えてください
         gr.Examples(
             examples=[
-                './sample/sample_AP_inverted.png',
-                './sample/sample_PA_right.png',
-                './sample/sample_lateral_upright.png'
+                './samples/sample_chest_AP_0deg.png',
+                './samples/sample_chest_PA_90deg.png',
+                './samples/sample_chest_LAT_180deg.png'
             ],
             inputs=input_image
         )