import gradio as gr
from fastai.vision.all import *
import pathlib
import os

# Load the exported model
learn = load_learner("model_mnist.pkl")

# Define prediction function
def predict(img):
    pred, pred_idx, probs = learn.predict(img)
    return f"Prediction: {pred} (Confidence: {probs[pred_idx]:.4f})"

# Create Gradio interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload a digit image (3 or 7 only)"),
    outputs="text",
    title="MNIST Digit Classifier (3 vs 7)",
    description="Upload an image of a handwritten digit. Only digits 3 and 7 will work!"
)

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