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

# =========================
# Load trained model
# =========================
# Make sure you've trained and saved it as best_model.h5 in your notebook
model = tf.keras.models.load_model("mnist_model.h5")

# =========================
# Prediction function
# =========================
def predict(image):
    """
    Takes a PIL image, preprocesses it (grayscale + resize),
    runs prediction using trained model, and returns predicted digit.
    """
    # Convert to grayscale + resize
    image = image.convert("L").resize((28, 28))

    # Convert to numpy and normalize
    img_array = np.array(image) / 255.0
    img_array = img_array.reshape(1, 28, 28, 1)  # batch shape

    # Predict
    prediction = model.predict(img_array)
    predicted_class = np.argmax(prediction, axis=1)[0]

    # Also return top-3 predictions with probabilities
    top3_indices = prediction[0].argsort()[-3:][::-1]
    top3_probs = prediction[0][top3_indices]

    result = {str(d): float(p) for d, p in zip(top3_indices, top3_probs)}
    return result

# =========================
# Gradio interface
# =========================
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", image_mode="L"),
    outputs=gr.Label(num_top_classes=3),   # show top 3 predictions
    title="MNIST Digit Classifier",
    description="Upload a handwritten digit (0–9) image. The model will predict the digit."
)

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