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

st.set_page_config(page_title="Image Classification", layout="wide")
st.title("Image Classification")

# ---------------------------------------------------
# CLASS LABELS (MUST MATCH TRAINING ORDER EXACTLY)
# ---------------------------------------------------
CLASS_NAMES = [
    "person",          # 0
    "bicycle",         # 1
    "car",             # 2
    "motorcycle",      # 3
    "airplane",        # 4
    "bus",             # 5
    "train",           # 6
    "truck",           # 7
    "traffic light",   # 8
    "stop sign",       # 9
    "bench",           # 10
    "bird",            # 11
    "cat",             # 12
    "dog",             # 13
    "horse",           # 14
    "cow",             # 15
    "elephant",        # 16
    "bottle",          # 17
    "cup",             # 18
    "bowl",            # 19
    "pizza",           # 20
    "cake",            # 21
    "chair",           # 22
    "couch",           # 23
    "potted plant",    # 24
    "bed"              # 25
]

# ---------------------------------------------------
# MODEL LOADING (CACHED)
# ---------------------------------------------------
@st.cache_resource
def load_model(path):
    return tf.keras.models.load_model(path)

MODEL_PATHS = {
    "VGG16": "models/vgg16_smartvision.h5",
    "ResNet50": "models/resnet50_smartvision.h5",
    "MobileNetV2": "models/mobilenetv2_smartvision.h5",
    "EfficientNetB0": "models/efficientnetb0_smartvision.h5"
}

# ---------------------------------------------------
# FILE UPLOAD
# ---------------------------------------------------
uploaded_file = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])

if uploaded_file:

    try:
        col1, col2 = st.columns(2)

        # -----------------------
        # Image Preprocessing
        # -----------------------
        image = Image.open(uploaded_file).convert("RGB")
        resized_image = image.resize((224, 224))

        with col1:
            st.subheader("Uploaded Image")
            st.image(image, width="stretch")

        img_array = np.array(resized_image) / 255.0
        img_array = np.expand_dims(img_array, axis=0)

        # -----------------------
        # Predictions
        # -----------------------
        with col2:
            st.subheader("Model Predictions")

            for model_name, path in MODEL_PATHS.items():

                model = load_model(path)
                predictions = model.predict(img_array, verbose=0)[0]

                top10_indices = np.argsort(predictions)[-10:][::-1]

                st.markdown(f"### {model_name} - Top 10 Predictions")

                for idx in top10_indices:
                    class_name = CLASS_NAMES[idx]
                    confidence = predictions[idx]
                    st.write(f"{class_name} — {confidence:.3f}")

                st.divider()

    except Exception as e:
        st.error(f"Error during classification: {e}")