frank1206/quickdraw_duplicate

1from pathlib import Path2 3import torch4import gradio as gr5from torch import nn6 7 8LABELS = Path('class_names.txt').read_text().splitlines()9 10model = nn.Sequential(11    nn.Conv2d(1, 32, 3, padding='same'),12    nn.ReLU(),13    nn.MaxPool2d(2),14    nn.Conv2d(32, 64, 3, padding='same'),15    nn.ReLU(),16    nn.MaxPool2d(2),17    nn.Conv2d(64, 128, 3, padding='same'),18    nn.ReLU(),19    nn.MaxPool2d(2),20    nn.Flatten(),21    nn.Linear(1152, 256),22    nn.ReLU(),23    nn.Linear(256, len(LABELS)),24)25state_dict = torch.load('pytorch_model.bin', map_location='cpu')26model.load_state_dict(state_dict, strict=False)27model.eval()28 29def predict(im):30    x = torch.tensor(im, dtype=torch.float32).unsqueeze(0).unsqueeze(0) / 255.31 32    with torch.no_grad():33        out = model(x)34 35    probabilities = torch.nn.functional.softmax(out[0], dim=0)36 37    values, indices = torch.topk(probabilities, 5)38 39    return {LABELS[i]: v.item() for i, v in zip(indices, values)}40 41 42interface = gr.Interface(predict, inputs='sketchpad', outputs='label', live=True)43interface.launch(debug=True)44