import gradio as gr
import tensorflow as tf
from PIL import Image
import numpy as np

# Laden des vortrainierten Pokémon-Modells
model_path = "kia_pokemon_keras_model.h5"
model = tf.keras.models.load_model(model_path)

# Labels für den Pokémon Classifier 
labels = [
    'Bulbasaur','Charmander','Squirtle'
]

def predict_pokemon(image):
    # Preprocess image
    image = Image.fromarray(image.astype('uint8'))  # Convert numpy array to PIL image
    image = image.resize((224, 224))  # Resize the image to 224x224
    image = np.array(image)
    image = np.expand_dims(image, axis=0)  # same as image[None, ...]
    
    # Predict
    predictions = model.predict(image)
    prediction = np.argmax(predictions, axis=1)[0]
    confidence = np.max(predictions)
    
    # Vorbereiten der Ausgabe
    result = f"Predicted Pokémon: {labels[prediction]} with confidence: {confidence:.2f}"
    return result

# Erstellen der Gradio-Oberfläche
input_image = gr.Image() 
output_label = gr.Label()
interface = gr.Interface(fn=predict_pokemon,
                         inputs=input_image,
                         outputs=output_label,
                         examples=["images/bulbasaur.png", "images/charmander.png", "images/squirtle.png"], 
                         title="Pokémon Classifier",
                         description="Drag and drop an image or select an example below to predict the Pokémon.")

# Interface starten
interface.launch()