Update app.py

app.py

@@ -11,6 +11,24 @@ model = ViTForImageClassification.from_pretrained(model_name)
 image_processor = ViTImageProcessor.from_pretrained(model_name)
 model.eval()
 
+# NIH Chest X-ray predefined conditions
+labels = [
+    "Atelectasis",
+    "Cardiomegaly",
+    "Effusion",
+    "Infiltration",
+    "Mass",
+    "Nodule",
+    "Pneumonia",
+    "Pneumothorax",
+    "Consolidation",
+    "Edema",
+    "Emphysema",
+    "Fibrosis",
+    "Pleural Thickening",
+    "Hernia"
+]
+
 # Function to apply Grad-CAM visualization
 def generate_grad_cam(image, target_layer):
     # Preprocess the image
@@ -53,8 +71,7 @@ def predict_and_explain(image):
     logits = outputs.logits
     predicted_class_idx = logits.argmax(-1).item()
     
-    # Predefined medical conditions (adjust based on your dataset)
-    labels = ["Class 1 - Normal", "Class 2 - Condition A", "Class 3 - Condition B"]
+    # Get the predicted label based on NIH Chest X-ray conditions
     predicted_label = labels[predicted_class_idx]
     
     # Generate Grad-CAM heatmap
@@ -82,8 +99,8 @@ interface = gr.Interface(
         "text",
         gr.Image(type="file", label="Grad-CAM Visualization")
     ],
-    title="Medical Image Analysis Tool with Explainability",
-    description="Upload an X-ray or MRI image to get a prediction for a medical condition with explainability through Grad-CAM.",
+    title="Medical Image Analysis Tool with NIH Chest X-ray",
+    description="Upload a Chest X-ray image to get a prediction for common thoracic conditions based on the NIH dataset, with explainability through Grad-CAM.",
     live=True
 )