Spaces:
Runtime error
Runtime error
File size: 1,598 Bytes
cdfbbe4 3d16b2e cbc3a01 3d16b2e cdfbbe4 cbc3a01 cdfbbe4 |
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 37 38 39 40 41 42 43 44 |
import gradio as gr
import torch
from torch import nn
from pathlib import Path
LABELS = Path('class_names.txt').read_text().splitlines()
title = "Sketch Classifier"
description = "This machine has vision. To test the machine, if you sketch something below, the machine will attempt to classify your sketch. It can guess from 100 objects and it will present to you its top three guesses. Beside each guess, the length of the bar indicates the confidence with which the machine has identified your sketch. The longer the bar, the more confident the machine is."
model = nn.Sequential(
nn.Conv2d(1, 32, 3, padding='same'),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding='same'),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding='same'),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(1152, 256),
nn.ReLU(),
nn.Linear(256, len(LABELS)),
)
state_dict = torch.load('pytorch_model.bin', map_location='cpu')
model.load_state_dict(state_dict, strict=False)
model.eval()
def predict(img):
x = torch.tensor(img, dtype=torch.float32).unsqueeze(0).unsqueeze(0) / 255.
with torch.no_grad():
out = model(x)
probabilities = torch.nn.functional.softmax(out[0], dim=0)
values, indices = torch.topk(probabilities, 3)
confidences = {LABELS[i]: v.item() for i, v in zip(indices, values)}
return confidences
gr.Interface(fn=predict,
inputs="sketchpad",
outputs="label",
title=title,
description=description,
live=True).launch(debug=True)
|