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

# Lade dein Modell (hier als Beispiel die Keras .h5 Datei)
model = tf.keras.models.load_model('pokemon_model.keras')

# Klassennamen, sollten deinem Dataset entsprechen
class_names = ['Jolteon', 'Kakuna', 'Mr. Mime']

def classify_image(image):
    image = Image.fromarray(image.astype('uint8'), 'RGB')
    img = image.resize((150, 150))
    img_array = tf.keras.preprocessing.image.img_to_array(img)
    img_array = tf.expand_dims(img_array, 0)  # Erstelle einen Batch
    predictions = model.predict(img_array)
    predicted_class = class_names[np.argmax(predictions[0])]
    confidence = np.max(predictions[0])
    return {predicted_class: float(confidence)}


image_input = gr.Image()  # Entferne den `shape` Parameter
label = gr.Label(num_top_classes=3)

iface = gr.Interface(
    fn=classify_image, 
    inputs=image_input, 
    outputs=label,
    title='Pokémon Classifier',
             description='Upload an image of Jolteon, Kakuna, Mr. Mime and the classifier will tell you which one it is and the confidence level of the prediction.').launch()

iface.launch()