Update app.py

app.py

@@ -8,90 +8,47 @@ import matplotlib.image as mpimg
 def get_label(file_name):
     return file_name.split('-')[0]
 
-# Function to prepare data (similar to your code)
-def prepare_data(food_path, label_a, label_b):
-    for img in get_image_files(food_path):
-        if label_a in str(img):
-            img.rename(f"{img.parent}/{label_a}-{img.name}")
-        elif label_b in str(img):
-            img.rename(f"{img.parent}/{label_b}-{img.name}")
-        else:
-            os.remove(img)
-
-# Function to train the model
-def train_model(food_path, label_func):
-    dls = ImageDataLoaders.from_name_func(
-        food_path, get_image_files(food_path), valid_pct=0.2, seed=420,
-        label_func=label_func, item_tfms=Resize(230)
-    )
-
-    learn = cnn_learner(dls, resnet34, metrics=error_rate, pretrained=True)
-    learn.fine_tune(epochs=1)
-
-    return learn
-
 # Streamlit app
 def main():
     st.title("Food Classifier Streamlit App")
 
+    # Load the pre-trained model
+    model_path = "export.pkl"  # Update with the correct path to your export.pkl
+    model = load_learner(model_path)
+
     # Sidebar options
-    options = ["Train Model", "Upload Image", "Test Random Images", "Confusion Matrix"]
+    options = ["Upload Image", "Test Random Images", "Confusion Matrix"]
     choice = st.sidebar.selectbox("Choose an option", options)
 
-    if choice == "Train Model":
-        st.subheader("Training the Model")
-        food_path = Path("~/.fastai/data/food-101/food-101").expanduser()
-        if not food_path.exists():
-            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")
-
-            if st.button("Train Model"):
-                prepare_data(food_path, label_a, label_b)
-                learn = train_model(food_path, get_label)
-                st.session_state.model = learn  # Save the model to session state
-                st.success("Model trained successfully!")
-
-    elif choice == "Upload Image":
+    if 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)
+        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)
+        if uploaded_files:
+            for img in uploaded_files:
+                img = PILImage.create(img)
+                label, _, probs = model.predict(img)
 
-                    st.image(img, caption=f"This is a {label}.")
-                    st.write(f"{label}: {probs[1].item():.6f}")
-                    st.write(f"{label}: {probs[0].item():.6f}")
+                st.image(img, caption=f"This is a {label}.")
+                st.write(f"{label}: {probs[1].item():.6f}")
+                st.write(f"{label}: {probs[0].item():.6f}")
 
     elif choice == "Test Random Images":
         st.subheader("Test Using Images in Dataset")
-        if "model" not in st.session_state:
-            st.warning("Please train the model first.")
-        else:
-            for i in range(0, 5):  # Change 5 to the number of images you want to display
-                random_index = random.randint(0, len(get_image_files(food_path)) - 1)
-                img_path = get_image_files(food_path)[random_index]
-                img = mpimg.imread(img_path)
-                label, _, probs = st.session_state.model.predict(img)
+        for i in range(0, 5):  # Change 5 to the number of images you want to display
+            random_index = random.randint(0, len(get_image_files(food_path)) - 1)
+            img_path = get_image_files(food_path)[random_index]
+            img = mpimg.imread(img_path)
+            label, _, probs = model.predict(img)
 
-                st.image(img, caption=f"Predicted label: {label}")
+            st.image(img, caption=f"Predicted label: {label}")
 
     elif choice == "Confusion Matrix":
         st.subheader("Confusion Matrix")
-        if "model" not in st.session_state:
-            st.warning("Please train the model first.")
-        else:
-            interp = ClassificationInterpretation.from_learner(st.session_state.model)
-            st.pyplot(interp.plot_confusion_matrix())
+        interp = ClassificationInterpretation.from_learner(model)
+        st.pyplot(interp.plot_confusion_matrix())
 
 # Run the Streamlit app
 if __name__ == "__main__":
     main()
+