import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

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

# Modellpfad
model_path = "pokemon_classifier_model.keras"

# Modell laden
model = tf.keras.models.load_model(model_path)

# Klassenlabels (Passe diese entsprechend deinem Modell an)
labels = ['Abra', 'Cloyster', 'Dodrio']

# Vorhersagefunktion
def predict(image):
    # Bildvorverarbeitung
    image = image.resize((64, 64))
    image = np.array(image) / 255.0
    image = np.expand_dims(image, axis=0)
    predictions = model.predict(image)
    confidences = {labels[i]: float(predictions[0][i]) for i in range(len(labels))}
    return confidences

# Gradio-Interface erstellen
iface = gr.Interface(
    fn=predict, 
    inputs=gr.Image(),  # Keine shape-Parameter hier
    outputs=gr.Label(),
    description="Pokémon Classifier"
)

if __name__ == "__main__":
    iface.launch()