from ctypes import resize
import torch
from torchvision import transforms
from PIL import Image
import gradio as gr
from pathlib import Path

convert_tensor = transforms.ToTensor()
simple_net = torch.load("model.pkl", map_location=torch.device("cpu"))
simple_net.eval()

examples = list(map(str, list(Path("examples").iterdir())))
title = "MNIST classification with a simple neural net"
article = """The model has been trained on handwritten images similar to the ones shown in examples.
             The training images were of 28x28 resolution with black background. If one wants to test
             the model with images apart from the ones provided in the examples, it is required that the image
             be of black background. Other thing to note is that the more the test image is similar to the examples
             images in style, the better the model will work."""


def predict(im):
    im_tensor = convert_tensor(im).squeeze().flatten()
    model_outs = simple_net(im_tensor).sigmoid().tolist()
    preds = {i: o for i, o in enumerate(model_outs)}
    return preds


iface = gr.Interface(
    fn=predict,
    title=title,
    article=article,
    inputs=gr.Image(type="pil", image_mode="L", shape=(28, 28)),
    outputs=gr.Label(num_top_classes=10),
    examples=examples,
)

iface.launch()