Update app.py

app.py

@@ -1,4 +1,3 @@
-# streamlit_app.py
 import streamlit as st
 from fastai.vision.all import *
 import shutil
@@ -40,14 +39,23 @@ def main():
 
     if choice == "Train Model":
         st.subheader("Training the Model")
-        food_path = untar_data(URLs.FOOD)
-        label_a = st.text_input("Enter label A:", "samosa")
-        label_b = st.text_input("Enter label B:", "hot_and_sour_soup")
+        food_path = Path("~/.fastai/data/food-101/food-101").expanduser()
+        if not food_path.exists():
+            food_path = untar_data(URLs.FOOD)
+            label_a = st.text_input("Enter label A:", "samosa")
+            label_b = st.text_input("Enter label B:", "hot_and_sour_soup")
 
-        prepare_data(food_path, label_a, label_b)
-        learn = train_model(food_path, get_label)
+            prepare_data(food_path, label_a, label_b)
+            learn = train_model(food_path, get_label)
 
-        st.success("Model trained successfully!")
+            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")