Update app.py

app.py

@@ -8,38 +8,41 @@ import gradio as gr
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = models.densenet121(pretrained=True)
 
-# Modify the classifier layer to output probabilities for 14 classes (14 pathologies)
+# Modify the classifier layer to output probabilities for 14 classes (pathologies)
 num_classes = 14
-model.classifier = nn.Sequential(
-    nn.Linear(model.classifier.in_features, num_classes),
-    nn.Sigmoid()  # Use Sigmoid for multi-label classification
+model.fc = nn.Sequential(
+    nn.Linear(model.fc.in_features, num_classes),
+    nn.Sigmoid(),  # Use Sigmoid for multi-label classification
 )
-model.load_state_dict(torch.load('chexnet.pth', map_location=device))  # Load your pre-trained weights
-model = model.to(device)
+
+try:
+    model.load_state_dict(torch.load('chexnet.pth', map_location=device))
+except Exception as e:
+    print(f"Error loading pre-trained weights: {e}")
+model.to(device)
 model.eval()
 
 # Define image transformations (resize, normalize)
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    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 = [
-    'Atelectasis', 'Cardiomegaly', 'Effusion', 'Infiltration', 'Mass', 
-    'Nodule', 'Pneumonia', 'Pneumothorax', 'Consolidation', 'Edema', 
+    'Atelectasis', 'Cardiomegaly', 'Effusion', 'Infiltration', 'Mass',
+    'Nodule', 'Pneumonia', 'Pneumothorax', 'Consolidation', 'Edema',
     'Emphysema', 'Fibrosis', 'Pleural Thickening', 'Hernia'
 ]
 
 # Prediction function
 def predict_disease(image):
     image = transform(image).unsqueeze(0).to(device)  # Transform and add batch dimension
+
     with torch.no_grad():
         outputs = model(image)
         outputs = outputs.cpu().numpy().flatten()
-    
-    # Create a dictionary of disease probabilities
     result = {class_name: float(prob) for class_name, prob in zip(class_names, outputs)}
     return result
 
@@ -49,9 +52,9 @@ interface = gr.Interface(
     inputs=gr.inputs.Image(type='pil'),  # Input is an image
     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."
+    description="Upload a chest X-ray to detect the probability of 14 different diseases.",
 )
 
 # Launch the Gradio app
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch()