aliabd HF staff commited on
Commit
53a1b09
β€’
1 Parent(s): f674a9a

Upload with huggingface_hub

Browse files
Files changed (2) hide show
  1. README.md +1 -2
  2. run.py +51 -0
README.md CHANGED
@@ -6,7 +6,6 @@ colorFrom: indigo
6
  colorTo: indigo
7
  sdk: gradio
8
  sdk_version: 3.4.1
9
-
10
- app_file: app.py
11
  pinned: false
12
  ---
6
  colorTo: indigo
7
  sdk: gradio
8
  sdk_version: 3.4.1
9
+ app_file: run.py
 
10
  pinned: false
11
  ---
run.py ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pathlib import Path
2
+
3
+ import torch
4
+ import gradio as gr
5
+ from torch import nn
6
+ import gdown
7
+
8
+ url = 'https://drive.google.com/uc?id=1dsk2JNZLRDjC-0J4wIQX_FcVurPaXaAZ'
9
+ output = 'pytorch_model.bin'
10
+ gdown.download(url, output, quiet=False)
11
+
12
+ LABELS = Path('class_names.txt').read_text().splitlines()
13
+
14
+ model = nn.Sequential(
15
+ nn.Conv2d(1, 32, 3, padding='same'),
16
+ nn.ReLU(),
17
+ nn.MaxPool2d(2),
18
+ nn.Conv2d(32, 64, 3, padding='same'),
19
+ nn.ReLU(),
20
+ nn.MaxPool2d(2),
21
+ nn.Conv2d(64, 128, 3, padding='same'),
22
+ nn.ReLU(),
23
+ nn.MaxPool2d(2),
24
+ nn.Flatten(),
25
+ nn.Linear(1152, 256),
26
+ nn.ReLU(),
27
+ nn.Linear(256, len(LABELS)),
28
+ )
29
+ state_dict = torch.load('pytorch_model.bin', map_location='cpu')
30
+ model.load_state_dict(state_dict, strict=False)
31
+ model.eval()
32
+
33
+ def predict(input):
34
+ im = input
35
+ if im is None:
36
+ return None
37
+
38
+ x = torch.tensor(im, dtype=torch.float32).unsqueeze(0).unsqueeze(0) / 255.
39
+
40
+ with torch.no_grad():
41
+ out = model(x)
42
+
43
+ probabilities = torch.nn.functional.softmax(out[0], dim=0)
44
+
45
+ values, indices = torch.topk(probabilities, 5)
46
+
47
+ return {LABELS[i]: v.item() for i, v in zip(indices, values)}
48
+
49
+
50
+ interface = gr.Interface(predict, inputs=gr.templates.Sketchpad(label="Draw Here"), outputs=gr.Label(label="Guess"), theme="default", css=".footer{display:none !important}", live=True)
51
+ interface.launch(enable_queue=False)