Spaces:
Runtime error
Runtime error
File size: 1,127 Bytes
d7c377f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
import os
import torch
import torchvision.transforms as transforms
import torch.nn.functional as F
import gradio as gr
from model import Net
# loads demo model
if torch.cuda.is_available():
dev = "cuda:0"
else:
dev = "cpu"
device = torch.device(dev)
model = torch.load(f"./demo_model.pt", map_location=device)
model.eval()
# inference function
def inference(img):
transform = transforms.Compose([transforms.ToTensor(), transforms.Resize((28, 28))])
img = transform(img).unsqueeze(0) # transforms ndarray and adds batch dimension
with torch.no_grad():
output_probabilities = F.softmax(model(img), dim=1)[0] # probability prediction for each label
return {labels[i]: float(output_probabilities[i]) for i in range(len(labels))}
# Creates and launches gradio interface
labels = range(10) # 1-9 labels
outputs = gr.outputs.Label(num_top_classes=5)
gr.Interface(fn=inference, inputs='sketchpad', outputs=outputs, title="MNIST Interface",
description="Draw a number from 0-9 in the box and click submit to see the model's predictions.").launch()
|