Update app.py

app.py

@@ -32,6 +32,8 @@ def train_model(food_path, label_func):
 
 # ... (previous code)
 
+# ... (previous code)
+
 # Streamlit app
 def main():
     st.title("Food Classifier Streamlit App")
@@ -44,7 +46,10 @@ def main():
         st.subheader("Training the Model")
         food_path = Path("~/.fastai/data/food-101/food-101").expanduser()
         if not food_path.exists():
-            food_path = untar_data(URLs.FOOD)
+            try:
+                food_path = untar_data(URLs.FOOD)
+            except FileExistsError:
+                st.warning("Data directory already exists. Skipping download.")
             label_a = st.text_input("Enter label A:", "samosa")
             label_b = st.text_input("Enter label B:", "hot_and_sour_soup")
 
@@ -54,23 +59,12 @@ def main():
             st.session_state.model = learn  # Save the model to session state
             st.success("Model trained successfully!")
 
-    elif choice == "Upload Image":
-        st.subheader("Upload Your Own Images")
-        if "model" not in st.session_state:
-            st.warning("Please train the model first.")
-        else:
-            uploaded_files = st.file_uploader("Choose images", type=["jpg", "jpeg", "png"], accept_multiple_files=True)
-
-            if uploaded_files:
-                for img in uploaded_files:
-                    img = PILImage.create(img)
-                    label, _, probs = st.session_state.model.predict(img)
+    # ... (rest of the code remains unchanged)
 
-                    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}")
+# Run the Streamlit app
+if __name__ == "__main__":
+    main()
 
-    # ... (rest of the code remains unchanged)