Update app.py

app.py

@@ -1,41 +1,45 @@
-from transformers import AutoModelForImageClassification, AutoImageProcessor
+from transformers import AutoImageProcessor, AutoModelForImageClassification
 import torch
 import torch.nn.functional as F
 from PIL import Image
 import gradio as gr
 
-# Load the properly fine-tuned chest X-ray model
-model_name = "Lucario-K17/biomedclip_radiology_diagnosis"
+# -----------------------------
+# 1. Load pretrained chest X-ray model
+# -----------------------------
+model_name = "yikuan8/resnet50_chestxray14"  # real NIH ChestX-ray14 model
 model = AutoModelForImageClassification.from_pretrained(model_name)
 processor = AutoImageProcessor.from_pretrained(model_name)
 model.eval()
 
-# All 14 disease labels
-all_diseases = [
-    "Atelectasis", "Cardiomegaly", "Effusion", "Infiltration", "Mass",
-    "Nodule", "Pneumonia", "Pneumothorax", "Consolidation", "Edema",
-    "Emphysema", "Fibrosis", "Pleural_Thickening", "Hernia"
-]
+# Get labels directly from the model config
+id2label = model.config.id2label
 
-# Lock to desired diseases
+# Pick only these 3 diseases
 target_diseases = ["Pneumonia", "Effusion", "Atelectasis"]
-target_idxs = [all_diseases.index(d) for d in target_diseases]
 
+# -----------------------------
+# 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()
-    
+
     results = []
-    for i, d in zip(target_idxs, target_diseases):
-        results.append(f"{d}: {probs[i].item():.2f}")
-    
+    for idx, label in id2label.items():
+        if label in target_diseases:
+            results.append(f"{label}: {probs[idx].item():.2f}")
+
     return "\n".join(results)
 
+# -----------------------------
+# 3. Gradio interface
+# -----------------------------
 iface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),