nateraw commited on
Commit
bf5ee30
1 Parent(s): e9b7e7a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +43 -0
app.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pathlib import Path
2
+
3
+ import torch
4
+ import gradio as gr
5
+ from torch import nn
6
+
7
+
8
+ LABELS = Path('class_names.txt').read_text().splitlines()
9
+
10
+ model = 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
+ )
25
+ state_dict = torch.load('pytorch_model.bin', map_location='cpu')
26
+ model.load_state_dict(state_dict, strict=False)
27
+ model.eval()
28
+
29
+ def 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
+
42
+ interface = gr.Interface(predict, inputs='sketchpad', outputs='label')
43
+ interface.launch(debug=True)