|
|
import gradio as gr |
|
|
import numpy as np |
|
|
import torch |
|
|
import torchvision |
|
|
from torch import nn |
|
|
from huggingface_hub import snapshot_download |
|
|
|
|
|
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) |
|
|
|
|
|
@torch.no_grad() |
|
|
def predict(self, input): |
|
|
input = input.reshape(1, 1, 28, 28) |
|
|
out = self(input) |
|
|
return nn.functional.softmax(out[0], dim = 0) |
|
|
|
|
|
lenet = LeNet() |
|
|
|
|
|
lenet_pt = snapshot_download('stanimirovb/ibob-lenet-v1') + '/lenet-v1.pth' |
|
|
lenet.load_state_dict(torch.load(lenet_pt, map_location='cpu')) |
|
|
|
|
|
resize = torchvision.transforms.Resize((28, 28), antialias=True) |
|
|
def on_submit(img): |
|
|
with torch.no_grad(): |
|
|
img = img['composite'].astype(np.float32) |
|
|
img = torch.from_numpy(img) |
|
|
img = resize(img.unsqueeze(0)) |
|
|
|
|
|
result = lenet.predict(img) |
|
|
|
|
|
sorted = [[i, e] for i, e in enumerate(result.numpy())] |
|
|
sorted.sort(key = lambda a : -a[1]) |
|
|
|
|
|
return "\n".join(map(str, sorted)) |
|
|
|
|
|
iface = gr.Interface( |
|
|
title = "LeNet", |
|
|
fn = on_submit, |
|
|
inputs=gr.Sketchpad(image_mode='P'), |
|
|
outputs=gr.Text(), |
|
|
) |
|
|
iface.launch() |
|
|
|
