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

# Lade dein Modell
model_path = "pokemon-model.keras"
model = tf.keras.models.load_model(model_path)
model.summary()  # Check if the model architecture loaded matches the expected one

# Klassen Labels für deine vier Pokémon
labels = ['Squirtle', 'Pikachu', 'Charizard', 'Butterfree']


def predict_pokemon(image):
    # Bildvorverarbeitung
    image = Image.fromarray(image.astype('uint8'), 'RGB')
    image = image.resize((150, 150))  # Anpassen der Bildgröße an das Modell
    image = np.array(image)  # Normalisieren der Pixelwerte
    print(image.shape)
    
    # Bild in das Modell einspeisen und Vorhersage treffen
    prediction = model.predict(image[None, ...])
    confidences = {labels[i]: np.round(float(prediction[0][i]), 2) for i in range(len(labels))}
    return confidences

# Gradio Interface definieren
input_image = gr.Image()
output_text = gr.Textbox(label="Predicted Pokemon")
iface = gr.Interface(
    fn=predict_pokemon,
    inputs=input_image, 
    outputs=gr.Label(),
    title="Pokémon Classifier",
    description="Upload an image of a Pokémon and see the model classify it!"
)

# Starte die Gradio-Schnittstelle
iface.launch()