Update app.py

app.py

@@ -1,3 +1,4 @@
+# streamlit_app.py
 import streamlit as st
 from fastai.vision.all import *
 import shutil
@@ -51,12 +52,6 @@ def main():
             st.session_state.model = learn  # Save the model to session state
             st.success("Model trained successfully!")
 
-    # ... (rest of the code remains unchanged)
-
-# Run the Streamlit app
-if __name__ == "__main__":
-    main()
-
     elif choice == "Upload Image":
         st.subheader("Upload Your Own Images")
         if "model" not in st.session_state:
@@ -73,6 +68,21 @@ if __name__ == "__main__":
                     st.write(f"{label_a}: {probs[1].item():.6f}")
                     st.write(f"{label_b}: {probs[0].item():.6f}")
 
+    # ... (rest of the code remains unchanged)
+
+# Run the Streamlit app
+if __name__ == "__main__":
+    main()
+
+            if uploaded_files:
+                for img in uploaded_files:
+                    img = PILImage.create(img)
+                    label, _, probs = st.session_state.model.predict(img)
+
+                    st.image(img, caption=f"This is a {label}.")
+                    st.write(f"{label_a}: {probs[1].item():.6f}")
+                    st.write(f"{label_b}: {probs[0].item():.6f}")
+
     elif choice == "Test Random Images":
         st.subheader("Test Using Images in Dataset")
         if "model" not in st.session_state: