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

# Chargement des modèles avec message
with st.spinner("Chargement des modèles..."):
    try:
        model_infect = tf.keras.models.load_model("src/exo1.keras")
        model_animals = tf.keras.models.load_model("src/exo2.keras")
        st.success("Modèles chargés avec succès ✅")
    except Exception as e:
        st.error(f"Erreur lors du chargement des modèles : {e}")
        st.stop()

# Fonction de prédiction
def predict(model, img_array, classes):
    prediction = model.predict(img_array)
    index = np.argmax(prediction)
    return classes[index], float(prediction[0][index])

# Interface utilisateur
st.title("🧠 Classification d’images")

# Choix du modèle
option = st.selectbox("Choisissez le modèle :", ("Infecté / Non Infecté", "Chat / Chien"))

# Upload d'image
uploaded_file = st.file_uploader("Uploader une image", type=["jpg", "jpeg", "png"])

if uploaded_file:
    try:
        image = Image.open(uploaded_file).convert("RGB")
        st.image(image, caption="Image chargée", use_column_width=True)

        # Bouton pour déclencher la prédiction
        if st.button("Prédire"):
            with st.spinner("Prétraitement de l’image..."):
                image = image.resize((224, 224))
                img_array = np.array(image) / 255.0
                img_array = np.expand_dims(img_array, axis=0)

            with st.spinner("Prédiction en cours..."):
                if option == "Infecté / Non Infecté":
                    label, confidence = predict(model_infect, img_array, ["Non Infecté", "Infecté"])
                else:
                    label, confidence = predict(model_animals, img_array, ["Chat", "Chien"])

            st.success(f"✅ Classe prédite : **{label}** avec une confiance de **{confidence:.2f}**")
    except Exception as e:
        st.error(f"❌ Erreur lors du traitement : {e}")