Update app.py

app.py

@@ -1,19 +1,16 @@
 import streamlit as st
 from PIL import Image
-from fastai.vision.all import load_learner, PILImage
+import torch
+from torchvision import transforms
+from torchvision.models import resnet50
+import torch.nn.functional as F
 
 # Set up a title for the app
 st.title("Image Recognition App")
 
-# Load the pre-trained model from Hugging Face's repository or use a publicly available model
-model_path = 'model.pkl'  # Replace this with a model path hosted on Hugging Face, if needed
-
-try:
-    learn = load_learner(model_path)
-except Exception as e:
-    st.error("Error loading the model. Make sure 'model.pkl' is available in your Hugging Face Space.")
-    st.error(str(e))
-    st.stop()
+# Load the pre-trained PyTorch model
+model = resnet50(pretrained=True)
+model.eval()  # Set the model to evaluation mode
 
 # Upload an image
 uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
@@ -23,11 +20,20 @@ if uploaded_file is not None:
     image = Image.open(uploaded_file).convert("RGB")
     st.image(image, caption='Uploaded Image.', use_column_width=True)
 
+    # Preprocess the image
+    preprocess = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    image_tensor = preprocess(image).unsqueeze(0)  # Add a batch dimension
+
     # Run the model to make predictions
     st.write("Classifying...")
-    try:
-        pred, pred_idx, probs = learn.predict(PILImage.create(uploaded_file))
-        st.write(f"Prediction: {pred} with probability {probs[pred_idx]:.2f}")
-    except Exception as e:
-        st.error("Error making a prediction.")
-        st.error(str(e))
+    with torch.no_grad():
+        outputs = model(image_tensor)
+        probabilities = F.softmax(outputs[0], dim=0)
+        top3_prob, top3_classes = torch.topk(probabilities, 3)
+        for i in range(3):
+            st.write(f"Label: {top3_classes[i].item()}, Confidence: {top3_prob[i].item():.2f}")