Update app.py

app.py

@@ -29,13 +29,30 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
-# Class names for the 14 diseases (labels from ChestX-ray14 dataset)
+# Class names and their interpretations for the 14 diseases
 class_names = [
     'Atelectasis', 'Cardiomegaly', 'Effusion', 'Infiltration', 'Mass',
     'Nodule', 'Pneumonia', 'Pneumothorax', 'Consolidation', 'Edema',
     'Emphysema', 'Fibrosis', 'Pleural Thickening', 'Hernia'
 ]
 
+interpretations = {
+    'Atelectasis': "Partial or complete collapse of the lung.",
+    'Cardiomegaly': "Enlargement of the heart.",
+    'Effusion': "Fluid accumulation in the chest cavity.",
+    'Infiltration': "Substances such as fluid in the lungs.",
+    'Mass': "An abnormal growth in the lung.",
+    'Nodule': "Small round or oval-shaped growth in the lung.",
+    'Pneumonia': "Infection causing inflammation in the air sacs.",
+    'Pneumothorax': "Air in the pleural space causing lung collapse.",
+    'Consolidation': "Lung tissue that has filled with liquid.",
+    'Edema': "Excess fluid in the lungs.",
+    'Emphysema': "Damage to air sacs causing difficulty breathing.",
+    'Fibrosis': "Thickening or scarring of lung tissue.",
+    'Pleural Thickening': "Thickening of the pleura (lining of the lungs).",
+    'Hernia': "Displacement of an organ through a structure."
+}
+
 # Prediction function
 def predict_disease(image):
     image = transform(image).unsqueeze(0).to(device)  # Transform and add batch dimension
@@ -43,7 +60,12 @@ def predict_disease(image):
     with torch.no_grad():
         outputs = model(image)
         outputs = outputs.cpu().numpy().flatten()
-    result = {class_name: float(prob) for class_name, prob in zip(class_names, outputs)}
+    
+    # Result with interpretations
+    result = {
+        f"{class_name} ({interpretations[class_name]})": float(prob)
+        for class_name, prob in zip(class_names, outputs)
+    }
     return result
 
 # Gradio Interface