Update app.py

app.py

@@ -1,15 +1,15 @@
 import streamlit as st
 from PIL import Image
 import torch
-from torchvision import transforms
-from torchvision.models import resnet50
+import torchvision.transforms as transforms
+from torchvision import models
 import torch.nn.functional as F
 
 # Set up a title for the app
-st.title("Image Recognition App")
+st.title("Simple Image Recognition App")
 
-# Load the pre-trained PyTorch model
-model = resnet50(pretrained=True)
+# Load a pre-trained model from torchvision (e.g., ResNet50)
+model = models.resnet50(pretrained=True)
 model.eval()  # Set the model to evaluation mode
 
 # Upload an image
@@ -35,5 +35,8 @@ if uploaded_file is not None:
         outputs = model(image_tensor)
         probabilities = F.softmax(outputs[0], dim=0)
         top3_prob, top3_classes = torch.topk(probabilities, 3)
+
+        # Display the top 3 predictions
+        st.write("Predictions:")
         for i in range(3):
             st.write(f"Label: {top3_classes[i].item()}, Confidence: {top3_prob[i].item():.2f}")