import gradio as gr
import torch
from torch import nn

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

# Locate device
if torch.cuda.is_available():
    device = torch.device("cuda:0")
    print("GPU")
else:
    device = torch.device("cpu")
    print("CPU")


# Neural Network
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        
        self.convs = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=4, kernel_size=(5, 5)),   
            nn.Tanh(),                                                      
            nn.AvgPool2d(2, 2),                                            

            nn.Conv2d(in_channels=4, out_channels=12, kernel_size=(5, 5)), 
            nn.Tanh(),
            nn.AvgPool2d(2, 2) 
        )

        self.linear = nn.Sequential(
            nn.Linear(4*4*12,10)
        )
    
    def forward(self, x):
        x = self.convs(x)
        x = torch.flatten(x, 1)

        return self.linear(x)

# Loading model
model = LeNet().to(device)
model.load_state_dict(torch.load("model_mnist.pth", map_location=torch.device('cpu')))


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)