Update app.py

app.py

@@ -4,19 +4,19 @@ 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')
 
-# Accumulate images and labels for training
-accumulated_images = []
-accumulated_labels = []
+# 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(
@@ -76,48 +76,34 @@ st.markdown(
     unsafe_allow_html=True,
 )
 
-
 # Logo
 st.image("https://ecoclimsolutions.files.wordpress.com/2024/01/rmcai-removebg.png?resize=48%2C48")
 
 # Page title
 st.title("EcoIdentify by EcoClim Solutions")
 
-mode = st.selectbox("Select Mode", ["Predict Mode", "Train 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 *")
 
-if mode == "Predict Mode":
-    # ... [same code for Predict Mode] ...
+# Image upload section
+opt = st.selectbox("How do you want to upload the image for classification?", ("Please Select", "Upload image from device"))
 
-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 = None
 
-    # Image upload section
+if opt == 'Upload image from device':
     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'):
-                accumulated_images.append(image[np.newaxis, ...])
-                label_index = labels.index(user_label)
-                label_one_hot = tf.one_hot(label_index, len(labels))
-                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()
+        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, ...])
+            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.")
+