Update app.py

app.py

@@ -1,21 +1,21 @@
-from transformers import AutoImageProcessor, AutoModelForImageClassification
+from transformers import AutoModelForImageClassification, AutoImageProcessor
 import torch
 import torch.nn.functional as F
 from PIL import Image
 import gradio as gr
 
 # -----------------------------
-# 1. Load pretrained chest X-ray model
+# 1. Load pretrained model
 # -----------------------------
-model_name = "yikuan8/resnet50_chestxray14"  # real NIH ChestX-ray14 model
+model_name = "microsoft/resnet-50-finetuned-chestxray14"
 model = AutoModelForImageClassification.from_pretrained(model_name)
 processor = AutoImageProcessor.from_pretrained(model_name)
 model.eval()
 
-# Get labels directly from the model config
+# Get labels from config
 id2label = model.config.id2label
 
-# Pick only these 3 diseases
+# Focus only on 3 diseases
 target_diseases = ["Pneumonia", "Effusion", "Atelectasis"]
 
 # -----------------------------
@@ -33,7 +33,8 @@ def predict(image):
     results = []
     for idx, label in id2label.items():
         if label in target_diseases:
-            results.append(f"{label}: {probs[idx].item():.2f}")
+            prob = probs[idx].item()
+            results.append(f"{label}: {'YES' if prob > 0.5 else 'NO'} ({prob:.2f})")
 
     return "\n".join(results)
 
@@ -45,7 +46,7 @@ iface = gr.Interface(
     inputs=gr.Image(type="pil"),
     outputs="text",
     title="Chest X-ray: Pneumonia / Effusion / Atelectasis",
-    description="Upload a chest X-ray. Model predicts probability for Pneumonia, Effusion, and Atelectasis."
+    description="Upload a chest X-ray. Model predicts YES/NO with probabilities for Pneumonia, Effusion, and Atelectasis."
 )
 
 iface.launch()