Spaces:
Running
Running
| import numpy as np | |
| import streamlit as st | |
| 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 | |
| # (Your existing layout code remains unchanged) | |
| # 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 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.") | |
| 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 | |
| # Inside the "Train Mode" section | |
| # Train button | |
| if st.button('Train Model'): | |
| # Update the model with the user-provided image and label | |
| target = np.zeros((1, len(labels))) # Adjust dimensions for one-hot encoding | |
| target[0, labels.index(user_label)] = 1 | |
| # Reshape the image to match the model's input shape | |
| image = tf.image.resize(image, (224, 224)) # adjust the size as needed | |
| image = tf.expand_dims(image, axis=0) | |
| # Create a dataset with the new image and label | |
| dataset = tf.data.Dataset.from_tensor_slices((image, target)) | |
| dataset = dataset.batch(1) # Batch size 1 as we are using a single image | |
| # Train the model | |
| model.fit(dataset, epochs=1) # You might want to adjust the number of epochs | |
| st.success(f'Model has been trained with the new image and label.') | |