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

# Lade das trainierte Modell
model_path = "Fruits_fruits.keras"
model = tf.keras.models.load_model(model_path)

# Definiere die Vorhersagefunktion
def predict_fruit(image):
    # Bildvorverarbeitung
    image = Image.fromarray(image.astype('uint8'))  # Konvertiere numpy-Array zu PIL-Bild
    image = image.resize((150, 150))  # Ändere die Bildgröße auf 150x150
    image = np.array(image)
    image = np.expand_dims(image, axis=0)  # Füge Batch-Dimension hinzu
    image = image / 255.0  # Normiere die Pixelwerte auf den Bereich [0, 1]
    
    # Vorhersage
    prediction = model.predict(image)
    
    # Softmax anwenden, um Wahrscheinlichkeiten für jede Klasse zu erhalten
    probabilities = tf.nn.softmax(prediction)
    
    # Klassenbezeichnungen für die Früchte
    fruit_classes = ['Apple', 'Lemon', 'Strawberry']
    probabilities_dict = {fruit_class: round(float(probability), 2) for fruit_class, probability in zip(fruit_classes, probabilities[0])}
    
    return probabilities_dict

# Erstelle die Gradio-Schnittstelle
input_image = gr.inputs.Image(shape=(150, 150), label="Upload a Fruit Image")
output_label = gr.outputs.Label(num_top_classes=3, label="Prediction")

iface = gr.Interface(
    fn=predict_fruit,
    inputs=input_image, 
    outputs=output_label,
    live=True,
    title="Fruit Classification",
    description="Upload an image of an Apple, Lemon, or Strawberry to get the predicted class. The model uses a CNN trained on the Fruits-360 dataset.",
    examples=["images/01.jpg", "images/02.jpg", "images/03.jpg"],
    theme="dark"
)

iface.launch()