import time
import streamlit as st
import numpy as np
import PIL
import urllib.request
from utils import *
#from fastai.data.external import *

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

    labels = gen_labels()
    predicted_class = np.argmax(prediction, axis=1)[0]
    classification_result = 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

# 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
            
            # Ensure image shape is correct and add batch dimension
            img_array = preprocess(image)  # This should return an array of shape (1, 256, 256, 3)
            
            predicted_class, confidence = classify_garbage(img_array, model)
            
            st.info('The uploaded image has been classified as "{}" with {:.2f}% confidence.'.format(predicted_class, confidence))
            
        except Exception as e:
            st.error(f"An error occurred: {e}")