Update app.py

app.py

@@ -1,16 +1,12 @@
 import torch
 import torch.nn as nn
 from torchvision import transforms, models
-from torchvision.models import DenseNet121_Weights
 from PIL import Image
 import gradio as gr
 
 # Load the pre-trained DenseNet-121 model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Use DenseNet121 with updated weights parameter
-weights = DenseNet121_Weights.IMAGENET1K_V1
-model = models.densenet121(weights=weights)
+model = models.densenet121(pretrained=True)
 
 # Modify the classifier layer to output probabilities for 14 classes (pathologies)
 num_classes = 14
@@ -21,8 +17,8 @@ model.classifier = nn.Sequential(
 
 try:
     model.load_state_dict(torch.load('chexnet.pth', map_location=device))
-except FileNotFoundError:
-    print("Error loading pre-trained weights: 'chexnet.pth' file not found.")
+except Exception as e:
+    print(f"Error loading pre-trained weights: {e}")
 model.to(device)
 model.eval()
 
@@ -72,6 +68,7 @@ def predict_disease(image):
     }
     return result
 
+
 # References to display
 references = """
 1. Huang, G., et al. (2017). Densely Connected Convolutional Networks. Proceedings of the IEEE conference on computer vision and pattern recognition.
@@ -89,6 +86,15 @@ interface = gr.Interface(
     description="Upload a chest X-ray to detect the probability of 14 different diseases.",
 )
 
+# Gradio Interface
+interface = gr.Interface(
+    fn=predict_disease,
+    inputs=gr.components.Image(type='pil'),  # Updated input component
+    outputs="label",  # Output is a dictionary of labels with probabilities
+    title="CheXNet Pneumonia Detection",
+    description="Upload a chest X-ray to detect the probability of 14 different diseases.",
+)
+
 # Launch the Gradio app
 if __name__ == "__main__":
     interface.launch()