Update app.py

app.py

@@ -4,49 +4,36 @@ import torch.nn.functional as F
 from PIL import Image
 import gradio as gr
 
-# -----------------------------
-# 1. Load pretrained model
-# -----------------------------
-model_name = "microsoft/resnet-50-finetuned-chestxray14"
+# 1️⃣ Load fine-tuned vit-chest-xray model
+model_name = "codewithdark/vit-chest-xray"
 model = AutoModelForImageClassification.from_pretrained(model_name)
 processor = AutoImageProcessor.from_pretrained(model_name)
 model.eval()
 
-# Get labels from config
-id2label = model.config.id2label
+# 2️⃣ Define disease indices based on CheXpert labels
+# The model expects: ['Cardiomegaly', 'Edema', 'Consolidation', 'No Finding', 'Pneumonia']
+label_list = ['Cardiomegaly', 'Edema', 'Consolidation', 'No Finding', 'Pneumonia']
+# We only care about Pneumonia, Consolidation, Edema
+target_labels = ['Pneumonia', 'Consolidation', 'Edema']
+target_idxs = [label_list.index(lbl) for lbl in target_labels]
 
-# Focus only on 3 diseases
-target_diseases = ["Pneumonia", "Effusion", "Atelectasis"]
-
-# -----------------------------
-# 2. Prediction function
-# -----------------------------
 def predict(image):
     img = image.convert("RGB").resize((224, 224))
     inputs = processor(images=img, return_tensors="pt")
-
     with torch.no_grad():
         logits = model(**inputs).logits
-
-    probs = F.softmax(logits, dim=1).squeeze()
-
+    probs = torch.sigmoid(logits).squeeze()  # multi-label => sigmoid
     results = []
-    for idx, label in id2label.items():
-        if label in target_diseases:
-            prob = probs[idx].item()
-            results.append(f"{label}: {'YES' if prob > 0.5 else 'NO'} ({prob:.2f})")
-
+    for idx, lbl in zip(target_idxs, target_labels):
+        prob = probs[idx].item()
+        results.append(f"{lbl}: {'YES' if prob > 0.5 else 'NO'} ({prob:.2f})")
     return "\n".join(results)
 
-# -----------------------------
-# 3. Gradio interface
-# -----------------------------
 iface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title="Chest X-ray: Pneumonia / Effusion / Atelectasis",
-    description="Upload a chest X-ray. Model predicts YES/NO with probabilities for Pneumonia, Effusion, and Atelectasis."
+    title="Chest X-ray Multi-Disease Detector",
+    description="Upload a chest X-ray. Predicts Pneumonia, Consolidation, and Edema."
 )
-
 iface.launch()