Update app.py

app.py

@@ -1,39 +1,48 @@
 from transformers import AutoModelForImageClassification, AutoImageProcessor
 import torch
-import torch.nn.functional as F
 from PIL import Image
 import gradio as gr
 
-# 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()
 
-# 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', 'Cardiomegaly', 'No Finding', 'Edema']
-target_idxs = [label_list.index(lbl) for lbl in target_labels]
+labels = ['Cardiomegaly', 'Edema', 'Consolidation', 'No Finding', 'Pneumonia']
+target_labels = ['Pneumonia', 'Consolidation', 'Edema']
+target_idxs = [labels.index(lbl) for lbl in target_labels]
 
 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 = torch.sigmoid(logits).squeeze()  # multi-label => sigmoid
-    results = []
-    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)
+image = image.convert("RGB").resize((224, 224))
+inputs = processor(images=image, return_tensors="pt")
+with torch.no_grad():
+logits = model(**inputs).logits
+probs = torch.sigmoid(logits).squeeze()
+
+```
+detected = []
+results = []
+for idx, lbl in zip(target_idxs, target_labels):
+    prob = probs[idx].item()
+    status = "YES" if prob > 0.5 else "NO"
+    results.append(f"{lbl}: {status} ({prob:.2f})")
+    if status == "YES":
+        detected.append(lbl)
+
+if detected:
+    summary = f"⚠️ Patient shows signs of: {', '.join(detected)}."
+else:
+    summary = "✅ Patient appears healthy — no major lung issues detected."
+
+return "\n".join(results + ["\n" + summary])
+```
 
 iface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="Chest X-ray Multi-Disease Detector",
-    description="Upload a chest X-ray. Predicts Pneumonia, Consolidation, and Edema."
+fn=predict,
+inputs=gr.Image(type="pil"),
+outputs="text",
+title="Chest X-ray Disease Detector",
+description="Upload a chest X-ray to detect Pneumonia, Consolidation, and Edema. Gives clear patient health summary."
 )
-iface.launch()
+
+iface.launch()