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

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

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

def predict_pokemon(image):
    # Preprocess image
    image = Image.fromarray(np.array(image).astype('uint8'))  # Convert to PIL image
    image = image.resize((224, 224))  # Resize image to 224x224
    image = np.array(image)
    image = np.expand_dims(image, axis=0)  # Add batch dimension

    # 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:.2f}%"
    return result

st.title("Pokémon Classifier")

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

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)