src/streamlit_app.py

import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image
import os
import platform
from tensorflow.keras.preprocessing.image import img_to_array

hf_token = os.getenv("HF_TOKEN")

# Constants
IMG_SIZE = 128
CLASS_NAMES = ['cardboard', 'glass', 'metal', 'paper', 'plastic']

# Eco tips per class
ECO_TIPS = {
    'cardboard': "📦 *Tip:* Flatten cardboard boxes before recycling to save space and ensure proper processing.",
    'glass': "🍾 *Fact:* Glass can be recycled endlessly without losing quality. Rinse before placing in the bin.",
    'metal': "🛠️ *Tip:* Aluminum and tin cans are highly recyclable. Crushing them can save space in the recycling bin.",
    'paper': "📄 *Fact:* Paper fibers can only be recycled about 5-7 times. Avoid contamination with food or oil.",
    'plastic': "🧴 *Tip:* Not all plastics are recyclable. Check the plastic code and always clean before disposal."
}

try:
    current_dir = os.path.abspath(os.path.dirname(__file__))
except NameError:
    current_dir = os.getcwd()

if "HF_SPACE_ID" in os.environ:
    # Running on Hugging Face Spaces
    MODEL_PATH = os.path.join("/app", "src", "models", "trashsort_cnn.h5")
else:
    # Local development
    MODEL_PATH = os.path.join(current_dir, "models", "trashsort_cnn.h5") 

# Load the trained model
@st.cache_resource(show_spinner=False)
def load_model():
    return tf.keras.models.load_model(MODEL_PATH, compile=False)

# Preprocess uploaded image
def preprocess_image(img):
    img = img.resize((IMG_SIZE, IMG_SIZE))
    img = img_to_array(img)
    img = img / 255.0
    img = np.expand_dims(img, axis=0)
    return img

# Predict class and confidence
def predict_image(model, image):
    processed = preprocess_image(image)
    prediction = model.predict(processed)
    class_idx = np.argmax(prediction)
    confidence = float(np.max(prediction))
    label = CLASS_NAMES[class_idx]
    return label, confidence

# Streamlit UI
def main():
    st.set_page_config(page_title="TrashSort", page_icon="♻️")
    st.title("TrashSort: Smart Waste Classifier ♻️")

    tab1, tab2, tab3 = st.tabs(["📝 About", "📷 Classify Image", "📊 Model Info"])

    with tab1:
        st.header("About TrashSort")
        st.write("""
            TrashSort is a smart waste classifier app that can identify types of trash:
            **cardboard, glass, metal, paper,** or **plastic** from an image you upload.
            This helps promote proper waste segregation and recycling.
            
            **How to use:**
            1. Go to the 'Classify Image' tab.
            2. Upload a photo of the waste item.
            3. See the classification result along with a confidence score.
            4. Learn eco tips for proper disposal!
        """)

    with tab2:
        st.header("Upload Image for Classification")
        model = load_model()
        uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
        if uploaded_file is not None:
            image = Image.open(uploaded_file).convert("RGB")
            st.image(image, caption="Uploaded Image", use_container_width=True)

            with st.spinner("Classifying..."):
                label, confidence = predict_image(model, image)

            st.markdown(f"### Prediction: **{label.capitalize()}**")
            st.markdown(f"### Confidence: **{confidence * 100:.2f}%**")

            st.markdown("#### ♻️ Eco Tip:")
            st.info(ECO_TIPS[label])

    with tab3:
        st.header("Model Information")
        st.write("""
        This model is a Convolutional Neural Network trained on images of common trash categories.
        It classifies images into the following classes:
        - Cardboard
        - Glass
        - Metal
        - Paper
        - Plastic
        
        The model accuracy is around 76% on the validation set, with room for improvement on some classes.
        """)

    # Footer
    st.markdown("---")
    st.markdown(
        "<div style='text-align: center; color: gray;'>"
        "© 2025 Trash Sort App. Developed by <b>Cherilyn</b>."
        "</div>",
        unsafe_allow_html=True
    )

if __name__ == "__main__":
    main()