commit

app.py

@@ -0,0 +1,43 @@
+from pathlib import Path
+
+import torch
+import gradio as gr
+from torch import nn
+
+
+LABELS = Path('class_names.txt').read_text().splitlines()
+
+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(im):
+    x = torch.tensor(im, 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, 5)
+
+    return {LABELS[i]: v.item() for i, v in zip(indices, values)}
+
+
+interface = gr.Interface(predict, inputs='sketchpad', outputs='label', theme="default", description="Who wants to play Pictionary? Draw a common object like a shovel or a laptop, and the algorithm will guess in real time!", live=True)
+interface.launch(debug=True)

class_names.txt

@@ -0,0 +1,100 @@
+airplane
+alarm_clock
+anvil
+apple
+axe
+baseball
+baseball_bat
+basketball
+beard
+bed
+bench
+bicycle
+bird
+book
+bread
+bridge
+broom
+butterfly
+camera
+candle
+car
+cat
+ceiling_fan
+cell_phone
+chair
+circle
+clock
+cloud
+coffee_cup
+cookie
+cup
+diving_board
+donut
+door
+drums
+dumbbell
+envelope
+eye
+eyeglasses
+face
+fan
+flower
+frying_pan
+grapes
+hammer
+hat
+headphones
+helmet
+hot_dog
+ice_cream
+key
+knife
+ladder
+laptop
+light_bulb
+lightning
+line
+lollipop
+microphone
+moon
+mountain
+moustache
+mushroom
+pants
+paper_clip
+pencil
+pillow
+pizza
+power_outlet
+radio
+rainbow
+rifle
+saw
+scissors
+screwdriver
+shorts
+shovel
+smiley_face
+snake
+sock
+spider
+spoon
+square
+star
+stop_sign
+suitcase
+sun
+sword
+syringe
+t-shirt
+table
+tennis_racquet
+tent
+tooth
+traffic_light
+tree
+triangle
+umbrella
+wheel
+wristwatch

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:effb6ea6f1593c09e8247944028ed9c309b5ff1cef82ba38b822bee2ca4d0f3c
+size 1656903

requirements.txt

@@ -0,0 +1 @@
+torch