app.py

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
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,
)

# Logo
st.image("https://ecoclimsolutions.files.wordpress.com/2024/01/rmcai-removebg.png?resize=48%2C48")

# Page title
st.title("EcoIdentify by EcoClim Solutions")

# 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 *")

# 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:
        image = preprocess_image(file)

# Sidebar section
st.sidebar.title("Options")

# Create a radio button widget for training mode
training_mode = st.sidebar.radio("Select training mode", ["None", "Dropout", "Batch Normalization"])

if training_mode == 'None':
    user_choice == 'Predict'
elif training_mode == "Dropout" or "Batch Normalization":
    user_choice == 'Train'


# Display the current training mode and user choice
st.write(f"Training mode: {training_mode}")
st.write(f"User choice: {user_choice}")

try:
    if image is not None:
        st.image(image, width=256, caption='Uploaded Image')
        
        # Execute different code blocks based on user choice
        if user_choice == "Predict":
            # Call the model with the training mode argument
            # Use a dictionary to map the training mode to the corresponding boolean value
            prediction = model.predict_on_batch(image[np.newaxis, ...], training={"None": False, "Dropout": True, "Batch Normalization": True}[training_mode])
            st.success(f'Prediction: {labels[np.argmax(prediction[0], axis=-1)]}')
        elif user_choice == "Train":
            # Generate some dummy target data for demonstration
            # You can replace this with your actual target data
            target = np.random.randint(0, 6, size=(1,))
            # Call the model without the training argument
            loss = model.train_on_batch(image[np.newaxis, ...], target)
            st.success(f'Loss: {loss}')

except Exception as e:
    st.error(f"An error occurred: {e}.  Please contact us EcoClim Solutions at EcoClimSolutions.wordpress.com.")