Update app.py

app.py

@@ -1,3 +1,4 @@
+%%writefile /content/lung_disease_detection/app.py
 import gradio as gr
 import torch
 import torch.nn.functional as F
@@ -15,22 +16,20 @@ transform = transforms.Compose([
 
 # Load the trained ResNet50 model
 model = FineTunedResNet(num_classes=3)
-model.load_state_dict(torch.load('/content/lung_disease_detection/models/final_fine_tuned_resnet50.pth',
-                                 map_location=torch.device('cpu')))
+model.load_state_dict(torch.load('models/final_fine_tuned_resnet50.pth', map_location=torch.device('cpu')))
 model.eval()
 
-
 # Define a function to make predictions
 def predict(image):
     start_time = time.time()  # Start the timer
     image = transform(image).unsqueeze(0)  # Transform and add batch dimension
-
+    
     with torch.no_grad():
         output = model(image)
         probabilities = F.softmax(output, dim=1)[0]
         top_prob, top_class = torch.topk(probabilities, 3)
         classes = ['🦠 COVID', '🫁 Normal', '🦠 Pneumonia']  # Adjust based on the classes in your model
-
+        
     end_time = time.time()  # End the timer
     prediction_time = end_time - start_time  # Calculate the prediction time
 
@@ -39,10 +38,9 @@ def predict(image):
     for i in range(top_prob.size(0)):
         result += f"{classes[top_class[i]]}: {top_prob[i].item() * 100:.2f}%\\n"
     result += f"Prediction Time: {prediction_time:.2f} seconds"
-
+    
     return result
 
-
 # Example images with labels
 examples = [
     ['examples/Pneumonia/02009view1_frontal.jpg', '🦠 Pneumonia'],
@@ -67,4 +65,4 @@ interface = gr.Interface(
 )
 
 # Launch the interface
-interface.launch()
+interface.launch()