Update app.py

app.py

@@ -4,11 +4,20 @@ from PIL import Image
 import tensorflow as tf
 from utils import preprocess_image
 
+import numpy as np
+import streamlit as st
+import tensorflow as tf
+from PIL import Image
+from utils import preprocess_image
+
 # Initialize labels and model
 labels = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
 model = tf.keras.models.load_model('classify_model.h5')
 
-# Customized Streamlit layout
+# Accumulate images and labels for training
+accumulated_images = []
+accumulated_labels = []
+
 # Customized Streamlit styles
 st.markdown(
     """
@@ -74,47 +83,10 @@ st.image("https://ecoclimsolutions.files.wordpress.com/2024/01/rmcai-removebg.pn
 # Page title
 st.title("EcoIdentify by EcoClim Solutions")
 
-# Mode selection
 mode = st.selectbox("Select Mode", ["Predict Mode", "Train Mode"])
 
 if mode == "Predict Mode":
-    # Subheader
-    st.header("Upload a waste image to find its category")
-
-    # Note
-    st.markdown("* Please note that our dataset is trained primarily with images that contain a white background. Therefore, images with a white background would produce maximum accuracy *")
-
-    # Image upload section
-    opt = st.selectbox("How do you want to upload the image for classification?", ("Please Select", "Upload image from device"))
-
-    image = None
-
-    if opt == 'Upload image from device':
-        file = st.file_uploader('Select', type=['jpg', 'png', 'jpeg'])
-        if file:
-            try:
-                image = preprocess_image(file)
-            except Exception as e:
-                st.error(f"An error occurred: {e}. Please contact us EcoClim Solutions at EcoClimSolutions.wordpress.com.")
-    try:
-        if image is not None:
-            st.image(image, width=256, caption='Uploaded Image')
-            if st.button('Predict'):
-                prediction = model.predict(image[np.newaxis, ...])
-                predicted_label = labels[np.argmax(prediction[0], axis=-1)]
-                st.success(f'Prediction: {predicted_label}')
-
-                # Ask user if the prediction is correct
-                user_feedback = st.radio("Is the prediction correct?", ["Yes", "No"])
-                if user_feedback == "No":
-                    # Allow user to provide correct label
-                    user_label = st.text_input("Enter the correct label:")
-                    if user_label:
-                        st.success(f'Thank you for providing feedback. Please switch to "Train Mode" to update the model.')
-
-    except Exception as e:
-        st.error(f"An error occurred: {e}. Please contact us EcoClim Solutions at EcoClimSolutions.wordpress.com.")
-
+    # ... [same code for Predict Mode] ...
 
 elif mode == "Train Mode":
     # Train the model with a new image and label
@@ -132,11 +104,19 @@ elif mode == "Train Mode":
 
             # Train button
             if st.button('Train Model'):
-                # Update the model with the user-provided image and label
-                image = image[np.newaxis, ...]  # Add a batch dimension
+                accumulated_images.append(image[np.newaxis, ...])
                 label_index = labels.index(user_label)
                 label_one_hot = tf.one_hot(label_index, len(labels))
-                model.fit(image, label_one_hot, epochs=2, batch_size=1)
-                st.success(f'Model has been trained with the new image and label.')
+                accumulated_labels.append(label_one_hot)
+                
+                if len(accumulated_images) >= 5:  # Example threshold
+                    X_train = np.vstack(accumulated_images)
+                    y_train = np.vstack(accumulated_labels)
+                    model.fit(X_train, y_train, epochs=2, batch_size=1)
+                    st.success(f'Model has been trained with the accumulated images and labels.')
+                    # Clear accumulated data
+                    accumulated_images.clear()
+                    accumulated_labels.clear()
+
         except Exception as e:
             st.error(f"An error occurred: {e}. Please contact us EcoClim Solutions at EcoClimSolutions.wordpress.com.")