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


model_path = "iris_mlp.weights.h5"
model = tf.keras.Sequential([
    tf.keras.layers.InputLayer(input_shape=[4]),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Dense(32, activation="relu"),
    tf.keras.layers.Dense(16, activation="relu"),
    tf.keras.layers.Dense(3, activation="softmax")
])
model.load_weights(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()