import gradio as gr
import tensorflow as tf
import numpy as np


model_path = "iris_mlp.keras"
model = tf.keras.models.load_model(model_path)

labels = ['Setosa', 'Versicolour', 'Virginica']

# Define the core prediction function
def predict_iris(sepal_length, sepal_width, petal_length, petal_width):
    features = [sepal_length, sepal_width, petal_length, petal_width]
    features = np.array(features)[None, ...]
    prediction = model.predict(features)
    print(prediction)
    confidences = {labels[i]: np.round(float(prediction[0][i]), 2) for i in range(len(labels))}
    return confidences

# Create the Gradio interface
iface = gr.Interface(
    fn=predict_iris,
    inputs=["number", "number", "number", "number"],
    outputs=gr.Label(),
    examples=[[7.7, 2.6, 6.9, 2.3]]
)

iface.launch()