import gradio as gr
import torch

labels = ['Zero','Um','Dois','Três','Quatro','Cinco','Seis','Sete','Oito', 'Nove']

# LOADING MODEL
model.load_state_dict(torch.load("model_mnist.pth", map_location=torch.device('cuda')))


def predict(input):
  input = torch.from_numpy(input.reshape(1, 1, 28, 28)).to(dtype=torch.float32, device=device)

  with torch.no_grad():
    outputs = model(input)  
    prediction = torch.nn.functional.softmax(outputs[0], dim=0)
    confidences = {labels[i]: float(prediction[i]) for i in range(10)}    
  return confidences

gr.Interface(title='Classificador de dígitos', fn=predict, 
             inputs="sketchpad",
             outputs=gr.Label(num_top_classes=3)).launch(share=True, debug=True)