import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image
import json

# Load model 
@st.cache_resource
def load_model():
    model = tf.keras.models.load_model("./src/intel_classifier_v2.h5")
    return model

model = load_model()

# Load class indices 
with open("./src/class_indices.json", "r") as f:
    class_indices = json.load(f)

class_names = list(class_indices.keys())


# Main
def run():

    st.title("Environment Image Classifier")

    # Form input 
    with st.form(key='form_image_classifier'):
        uploaded_file = st.file_uploader(
            "Upload an image | Classes = (Buildings, Forest, Glacier, Mountain, Sea, Street)",
            type=["jpg", "jpeg", "png"],
            help="Upload an environment image to classify"
        )

        submitted = st.form_submit_button('Predict')

    if submitted:

        if uploaded_file is None:
            st.warning("Please upload an image first.")
            return

        # Display image
        image = Image.open(uploaded_file).convert('RGB')
        st.image(image, caption="Uploaded Image", use_container_width=True)

        # Preprocess
        img = image.resize((150, 150))
        img_array = np.array(img) / 255.0
        img_array = np.expand_dims(img_array, axis=0)

        # Predict
        prediction = model.predict(img_array)[0]
        predicted_index = np.argmax(prediction)
        predicted_class = class_names[predicted_index]
        confidence = np.max(prediction)

        # Output
        st.subheader("Prediction Result")

        st.success(f"Predicted Class: {predicted_class}")
        st.write(f"Confidence: {confidence:.2%}")

        st.write("### Class Probabilities")
        for i, prob in enumerate(prediction):
            st.write(f"{class_names[i]}: {prob:.2%}")

if __name__ == '__main__':
    run()