import streamlit as st
import numpy as np
from tensorflow.keras.preprocessing import image
import tensorflow as tf

# Lade das gespeicherte Modell
model = tf.keras.models.load_model("pokemon_classification_model_xception_v3.keras")

# Setze die Bildabmessungen
img_height, img_width = 299, 299  # Eingabegröße für Xception

# Definiere eine Funktion zur Vorhersage und Rückgabe des Labels und der Wahrscheinlichkeit
def predict_label_and_probability(image_path):
    # Lade das Bild und passe es an die Eingabegröße des Modells an
    img = image.load_img(image_path, target_size=(img_height, img_width))
    x = image.img_to_array(img)
    x = np.expand_dims(x, axis=0)
    x /= 255.  # Skalieren der Bildpixel

    # Vorhersage mit dem Modell
    preds = model.predict(x)
    class_idx = np.argmax(preds[0])
    
    # Mappe Klassenindizes auf Klassennamen
    class_labels = {0: 'Abra', 1: 'Butterfree', 2: 'Eevee'}
    predicted_class = class_labels[class_idx]
    
    # Gib das vorhergesagte Label und die Wahrscheinlichkeit zurück
    probability = preds[0][class_idx]
    return predicted_class, probability

# Streamlit App
st.title("Pokémon Classification")

uploaded_file = st.file_uploader("Choose a Pokémon image...", type="jpg")

if uploaded_file is not None:
    # Zeige das hochgeladene Bild
    st.image(uploaded_file, caption='Uploaded Pokémon Image.', use_column_width=True)

    # Führe die Vorhersage durch und zeige das Ergebnis
    label, probability = predict_label_and_probability(uploaded_file)
    st.write("Prediction:", label)
    st.write("Probability:", probability)