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

# Load the pre-trained Pokémon model
model_path = "kia_pokemon_keras_model.keras"
model = tf.keras.models.load_model(model_path)

# Pokémon classifier labels
labels = ['Bulbasaur', 'Charmander', 'Squirtle']

def preprocess_image(image):
    # Convert image to RGB if it has an alpha channel
    if image.mode != 'RGB':
        image = image.convert('RGB')
    image = image.resize((224, 224))  # Resize image to 224x224
    image = np.array(image)
    image = np.expand_dims(image, axis=0)  # Add batch dimension
    return image

def predict_pokemon(image):
    # Preprocess image
    image = preprocess_image(image)

    # Predict
    predictions = model.predict(image)
    prediction = np.argmax(predictions, axis=1)[0]
    confidence = np.max(predictions)

    # Prepare output
    result = f"Predicted Pokémon: {labels[prediction]} with confidence: {confidence * 100:.2f}"
    return result

st.title("Pokémon Classifier")

file_uploader = st.file_uploader("Upload an image of a Pokémon", type=['png', 'jpg', 'jpeg'])

# Beispielbilder anzeigen
st.subheader("Or select an example image:")
example_images = ["images/bulbasaur.png", "images/charmander.png", "images/squirtle.png"]
example_labels = ["Bulbasaur", "Charmander", "Squirtle"]

for label, img_path in zip(example_labels, example_images):
    if st.button(label):
        file_uploader = open(img_path, "rb")

if file_uploader is not None:
    # Display the image
    image = Image.open(file_uploader)
    st.image(image, caption='Uploaded Image', use_column_width=True)

    # Make prediction
    result = predict_pokemon(image)
    st.subheader(result)