Update app.py

app.py

@@ -18,31 +18,31 @@ parameter = './parameters.json'
 
 # ===========================================
 # 2) クラスラベルの定義
-#    - 今回は撮像方向 (3クラス) と回転方向 (4クラス)
+#    - label_APorPA (3クラス): 0=AP, 1=PA, 2=Lateral
+#    - label_round (4クラス): 0=Upright, 1=Inverted, 2=Left rotation, 3=Right rotation
 # ===========================================
 LABEL_APorPA = [
-    "AP",        # class 0
-    "PA",        # class 1
-    "Lateral",   # class 2
-    # ↑ 実際のクラス名に合わせて変更してください
+    "AP",       # class 0
+    "PA",       # class 1
+    "Lateral",  # class 2
 ]
 
 LABEL_ROUND = [
-    "0° Rotation",    # class 0
-    "90° Rotation",   # class 1
-    "180° Rotation",  # class 2
-    "270° Rotation"   # class 3
+    "Upright",         # class 0
+    "Inverted",        # class 1
+    "Left rotation",   # class 2
+    "Right rotation"   # class 3
 ]
 
 # ===========================================
 # 3) 前処理用の ImageHandlerクラス
-#    - 画像が既に256×256前提。Resizeはコメントアウトで残す
+#    - 画像が既に256×256前提
 # ===========================================
 class ImageHandler(ImageMixin):
     def __init__(self, params):
         self.params = params
         self.transform = T.Compose([
-            # T.Resize((256, 256)),  # 必要ならコメントアウトを外す
+            # T.Resize((256, 256)),  # 必要であればコメントアウトを外す
             T.ToTensor(),
         ])
 
@@ -88,8 +88,8 @@ model.eval()  # 推論モード
 def classify_APorPA_and_round(image):
     """
     モデルが以下を出力する想定:
-      outputs["label_APorPA"]  -> shape=[1, 3] (3クラス)
-      outputs["label_round"]   -> shape=[1, 4] (4クラス)
+      outputs["label_APorPA"]  -> shape=[1, 3] (3クラス: AP, PA, Lateral)
+      outputs["label_round"]   -> shape=[1, 4] (4クラス: Upright, Inverted, Left rotation, Right rotation)
     """
     image_handler = ImageHandler(args_dataloader)
     image_tensor = image_handler.set_image(image)
@@ -109,10 +109,8 @@ def classify_APorPA_and_round(image):
             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,3]想定
+        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 ---
@@ -120,7 +118,7 @@ def classify_APorPA_and_round(image):
             print(f"[ERROR] 'label_round' not found in outputs. Actual keys: {list(outputs.keys())}")
             return predicted_APorPA, "ERROR: Missing 'label_round'"
 
-        scores_round = outputs["label_round"]  # shape=[1,4]想定
+        scores_round = outputs["label_round"]  # shape=[1,4]
         pred_round_idx = torch.argmax(scores_round, dim=1).item()
         predicted_round = LABEL_ROUND[pred_round_idx]
 
@@ -133,7 +131,8 @@ 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>胸部レントゲン画像に対して、撮像方向(3クラス)と回転方向(4クラス)を同時に推定します。</p>
+    <p>胸部レントゲン画像に対して、撮像方向(3クラス: AP, PA, Lateral) と
+    回転方向(4クラス: Upright, Inverted, Left rotation, Right rotation)を同時に推定します。</p>
 </div>
 """
 
@@ -143,7 +142,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/PA/Lateral")  # 3クラス想定に合わせた名前
+        output_APorPA = gr.Label(label="Predicted AP/PA/Lateral")  
         output_round = gr.Label(label="Predicted Rotation")
 
     send_btn = gr.Button("Inference")
@@ -157,9 +156,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'
+                './sample/sample_AP_upright.png',
+                './sample/sample_PA_inverted.png',
+                './sample/sample_lateral_left.png'
             ],
             inputs=input_image
         )