from model import load_model

import gradio as gr
import os
import torch



model,transforming,classes = load_model()

def predict(img):
    img = transforming(img)
    img = img.unsqueeze(0)


    model.eval()
    with torch.inference_mode():
        pred_probs = torch.softmax(model(img), dim=1)
    return {str(i): float(pred_probs[0][i]) for i in range(len(pred_probs[0]))}

title = 'MNIST Digit Prediction'
description = 'Predict handwritten digits (0-9) using a trained model.'
inputs = gr.Sketchpad(label='Draw a digit')
outputs = gr.Label(num_top_classes=3, label='Predictions')
demo = gr.Interface(fn=predict, inputs=inputs, outputs=outputs, title=title, description=description)

# Launch the interface
demo.launch(share=True)