app.py

import time
import streamlit as st
import numpy as np
from PIL import Image
import urllib.request
from utils import *  # Assuming the gen_labels() and preprocess() functions are in this module

# Function to classify the garbage
def classify_garbage(img_path, model):
    processed_img = preprocess_image(img_path)
    prediction = model.predict(processed_img)

    class_labels = ["cardboard", "glass", "metal", "paper", "plastic", "trash"]
    predicted_class = np.argmax(prediction, axis=1)[0]
    classification_result = class_labels[predicted_class]

    # Get the confidence (probability) of the predicted class
    confidence = prediction[0][predicted_class] * 100  # Convert probability to percentage

    return classification_result, confidence

# Load labels
labels = gen_labels()

# Streamlit app layout
st.markdown('<center><h1>Garbage Segregation</h1></center>', unsafe_allow_html=True)
st.markdown('<center><h3>Please upload Waste Image to find its Category</h3></center>', unsafe_allow_html=True)

opt = st.selectbox("How do you want to upload the image for classification?", ('Please Select', 'Upload image via link', 'Upload image from device'))

if opt == 'Upload image from device':
    file = st.file_uploader('Select', type=['jpg', 'png', 'jpeg'])
    if file is not None:
        image = Image.open(file).resize((256, 256), Image.LANCZOS)

elif opt == 'Upload image via link':
    img_url = st.text_input('Enter the Image Address')
    try:
        image = Image.open(urllib.request.urlopen(img_url)).resize((256, 256), Image.LANCZOS)
    except ValueError:
        st.error("Please Enter a valid Image Address!")

if 'image' in locals():  # Check if image variable exists
    st.image(image, width=300, caption='Uploaded Image')

    if st.button('Predict'):
        try:
            model = model_arc()  # Initialize your model
            predicted_class, confidence = classify_garbage(image, model)
            
            st.info('The uploaded image has been classified as "{}" waste with {:.2f}% confidence.'.format(predicted_class, confidence))
            
        except Exception as e:
            st.error(f"An error occurred: {e}")