Update app.py

app.py

@@ -1,31 +1,32 @@
-import streamlit as st
-from fastai.vision.all import *
-
-# Load the trained model
-learn = load_learner('export.pkl')
-
-# Define a function to make predictions on an image
-def predict_image(image):
-    img = PILImage.create(image)
-    pred, _, probs = learn.predict(img)
-    return pred, probs[pred]
-
-# Main function for Streamlit app
-def main():
-    st.title("Image Classifier")
-    st.write("Upload an image to classify")
-
-    # Upload image file
-    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
-
-    # Display prediction when an image is uploaded
-    if uploaded_file is not None:
-        image = Image.open(uploaded_file)
-        st.image(image, caption='Uploaded Image.', use_column_width=True)
-        st.write("Classifying...")
-        prediction, probability = predict_image(uploaded_file)
-        st.write(f"Prediction: {prediction}; Probability: {probability:.4f}")
-
-# Run the Streamlit app
-if __name__ == '__main__':
-    main()
+import streamlit as st
+from fastai.vision.all import *
+from PIL import Image
+
+# Load the trained model
+learn = load_learner('export.pkl')
+
+# Define a function to make predictions on an image
+def predict_image(image):
+    img = PILImage.create(image)
+    pred, pred_idx, probs = learn.predict(img)
+    return pred, probs[pred_idx].item()
+
+# Main function for Streamlit app
+def main():
+    st.title("Image Classifier")
+    st.write("Upload an image to classify")
+
+    # Upload image file
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
+
+    # Display prediction when an image is uploaded
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)
+        st.image(image, caption='Uploaded Image.', use_column_width=True)
+        st.write("Classifying...")
+        prediction, probability = predict_image(uploaded_file)
+        st.write(f"Prediction: {prediction}; Probability: {probability:.4f}")
+
+# Run the Streamlit app
+if __name__ == '__main__':
+    main()