import gradio as gr
import pickle
import numpy as np

# Load your model
with open('your_mnist_model.pkl', 'rb') as f:
    model = pickle.load(f)

def predict(image):
    # Preprocess the image
    image = image.reshape(1, 28, 28) / 255.0  # Normalize to 0-1
    
    # Make prediction
    probabilities = model.predict_proba(image.reshape(1, -1))[0]
    
    # Format output
    return {str(i): float(prob) for i, prob in enumerate(probabilities)}

iface = gr.Interface(
    fn=predict,
    inputs=gr.Sketchpad(shape=(28, 28)),
    outputs=gr.Label(num_top_classes=10),
    title="MNIST Digit Classifier",
    description="Draw a digit (0-9) in the canvas and the model will predict it."
)

iface.launch()