Update app.py

app.py

@@ -1,29 +1,14 @@
 import streamlit as st
 from PIL import Image
-import torch
-import torchvision.transforms as transforms
-from torchvision import models
-
-# Check if CUDA (GPU) is available, otherwise use CPU
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Load the pre-trained model (ResNet-50)
-model = models.resnet50(pretrained=True)
-model = model.to(device)
-model.eval()
-
-# Define a function to preprocess the image
-def preprocess_image(image):
-    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]),
-    ])
-    return transform(image).unsqueeze(0).to(device)
+from fastai.vision.all import load_learner, PILImage, Path
 
 # Set up a title for the app
 st.title("Image Recognition App")
 
+# Load the pre-trained model (make sure it's stored in the same directory as 'app.py')
+model_path = Path("model.pkl")  # Adjust the path if needed
+learn = load_learner(model_path)
+
 # Upload an image
 uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
 
@@ -32,15 +17,9 @@ 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 and make predictions
+    # Run the model to make predictions
     st.write("Classifying...")
-    input_tensor = preprocess_image(image)
-    with torch.no_grad():
-        output = model(input_tensor)
-        probabilities = torch.nn.functional.softmax(output[0], dim=0)
-        top5_prob, top5_classes = torch.topk(probabilities, 5)
-        
-        # Display the top 5 predictions
-        st.write("Predictions:")
-        for i in range(5):
-            st.write(f"Label: {top5_classes[i].item()}, Confidence: {top5_prob[i].item():.2f}")
+    pred, pred_idx, probs = learn.predict(PILImage.create(uploaded_file))
+
+    # Display the predictions
+    st.write(f"Prediction: {pred} with probability {probs[pred_idx]:.2f}")