Update app.py

app.py

@@ -4,23 +4,25 @@ import torch.nn.functional as F
 from PIL import Image
 import gradio as gr
 
-# -----------------------------
-# 1. Load the pretrained model
-# -----------------------------
-model_name = "microsoft/resnet-50"  # fine-tuned for chest x-ray multi-disease
+# Load the properly fine-tuned chest X-ray model
+model_name = "Lucario-K17/biomedclip_radiology_diagnosis"
 model = AutoModelForImageClassification.from_pretrained(model_name)
 processor = AutoImageProcessor.from_pretrained(model_name)
 model.eval()
 
-# Example disease list (adjust depending on model config)
-diseases = ["Pneumonia", "Effusion", "Atelectasis"]
+# All 14 disease labels
+all_diseases = [
+    "Atelectasis", "Cardiomegaly", "Effusion", "Infiltration", "Mass",
+    "Nodule", "Pneumonia", "Pneumothorax", "Consolidation", "Edema",
+    "Emphysema", "Fibrosis", "Pleural_Thickening", "Hernia"
+]
+
+# Lock to desired 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():
@@ -28,25 +30,18 @@ def predict(image):
     
     probs = F.softmax(logits, dim=1).squeeze()
     
-    # Get top-3 predictions
-    top_probs, top_idxs = torch.topk(probs, k=3)
-    
     results = []
-    for idx, prob in zip(top_idxs, top_probs):
-        disease_name = diseases[idx] if idx < len(diseases) else f"Class {idx.item()}"
-        results.append(f"{disease_name}: {prob.item():.2f}")
+    for i, d in zip(target_idxs, target_diseases):
+        results.append(f"{d}: {probs[i].item():.2f}")
     
     return "\n".join(results)
 
-# -----------------------------
-# 3. Gradio interface
-# -----------------------------
 iface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title="Chest X-ray Detector",
-    description="Upload a chest X-ray. The model predicts Pneumonia, Effusion, or Atelectasis with probability."
+    title="Chest X-ray: Pneumonia / Effusion / Atelectasis",
+    description="Upload a chest X-ray. Model predicts probability for Pneumonia, Effusion, and Atelectasis."
 )
 
 iface.launch()