Update app.py

app.py

@@ -4,77 +4,12 @@ from PIL import Image
 import tensorflow as tf
 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
-st.set_page_config(
-    page_title="EcoIdentify by EcoClim Solutions",
-    page_icon="https://ecoclimsolutions.files.wordpress.com/2024/01/rmcai-removebg.png?resize=48%2C48",
-    layout="wide",
-    initial_sidebar_state="expanded",
-)
-
-# Customized Streamlit styles
-st.markdown(
-    """
-    <style>
-        body {
-            color: #333333;
-            background-color: #f9f9f9;
-            font-family: 'Helvetica', sans-serif;
-        }
-        .st-bb {
-            padding: 0rem;
-        }
-        .st-ec {
-            color: #666666;
-        }
-        .st-ef {
-            color: #666666;
-        }
-        .st-ei {
-            color: #333333;
-        }
-        .st-dh {
-            font-size: 36px;
-            font-weight: bold;
-            color: #4CAF50;
-            text-align: center;
-            margin-bottom: 20px;
-        }
-        .st-gf {
-            background-color: #4CAF50;
-            color: white;
-            padding: 15px 30px;
-            font-size: 18px;
-            border: none;
-            border-radius: 8px;
-            cursor: pointer;
-            transition: background-color 0.3s;
-        }
-        .st-gf:hover {
-            background-color: #45a049;
-        }
-        .st-gh {
-            text-align: center;
-            font-size: 24px;
-            font-weight: bold;
-            margin-bottom: 20px;
-        }
-        .st-logo {
-            max-width: 100%;
-            height: auto;
-            margin: 20px auto;
-            display: block;
-        }
-    </style>
-    """,
-    unsafe_allow_html=True,
-)
+# (Your existing layout code remains unchanged)
 
 # Logo
 st.image("https://ecoclimsolutions.files.wordpress.com/2024/01/rmcai-removebg.png?resize=48%2C48")
@@ -82,30 +17,71 @@ st.image("https://ecoclimsolutions.files.wordpress.com/2024/01/rmcai-removebg.pn
 # Page title
 st.title("EcoIdentify by EcoClim Solutions")
 
-# Subheader
-st.header("Upload a waste image to find its category")
+# 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 white background would produce maximum accuracy *")
 
-# Note
-st.markdown("* Please note that our dataset is trained primarily with images that contain a white background.  Therefore, images with 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 upload section
-opt = st.selectbox("How do you want to upload the image for classification?", ("Please Select", "Upload image from device"))
+    image = None
 
-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}')
 
-if opt == 'Upload image from device':
+                # 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:
+                        # Update the model with the user-provided image and label
+                        image = preprocess_image(file)  # preprocess the image again
+                        target = np.zeros(len(labels))
+                        target[labels.index(user_label)] = 1
+                        model.train_on_batch(image[np.newaxis, ...], target)
+                        st.success(f'Thank you for providing feedback. Model has been updated with the new label.')
+    except Exception as e:
+        st.error(f"An error occurred: {e}.  Please contact us EcoClim Solutions at EcoClimSolutions.wordpress.com.")
+
+elif mode == "Train Mode":
+    # Train the model with a new image and label
+    st.header("Train the model with a new image and label")
+
+    # Image upload section
     file = st.file_uploader('Select', type=['jpg', 'png', 'jpeg'])
     if file:
         try:
             image = preprocess_image(file)
+            st.image(image, width=256, caption='Uploaded Image')
+
+            # Label input
+            user_label = st.selectbox("Select the correct label", labels)
+
+            # Train button
+            if st.button('Train Model'):
+                # Update the model with the user-provided image and label
+                target = np.zeros(len(labels))
+                target[labels.index(user_label)] = 1
+                model.train_on_batch(image[np.newaxis, ...], target)
+                st.success(f'Model has been trained with the new image and label.')
         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, ...])
-            st.success(f'Prediction: {labels[np.argmax(prediction[0], axis=-1)]}')
-except Exception as e:
-    st.error(f"An error occurred: {e}.  Please contact us EcoClim Solutions at EcoClimSolutions.wordpress.com.")
-