some fixes

.gitignore

@@ -0,0 +1,3 @@
+*.pth
+.venv
+__pycache__

app.py

@@ -4,21 +4,19 @@ import torchvision.transforms as transforms
 import torch.nn.functional as F
 from pathlib import Path
 import gradio as gr
-from PIL import Image
-import numpy as np
+from huggingface_hub import hf_hub_download
 
-
-LABELS = Path('classes.txt').read_text().splitlines()
+LABELS = Path("classes.txt").read_text().splitlines()
 num_classes = len(LABELS)
 
 model = nn.Sequential(
-    nn.Conv2d(1, 64, 3, padding='same'),
+    nn.Conv2d(1, 64, 3, padding="same"),
     nn.ReLU(),
     nn.MaxPool2d(2),
-    nn.Conv2d(64, 128, 3, padding='same'),
+    nn.Conv2d(64, 128, 3, padding="same"),
     nn.ReLU(),
     nn.MaxPool2d(2),
-    nn.Conv2d(128, 256, 3, padding='same'),
+    nn.Conv2d(128, 256, 3, padding="same"),
     nn.ReLU(),
     nn.MaxPool2d(2),
     nn.Flatten(),
@@ -27,32 +25,40 @@ model = nn.Sequential(
     nn.Linear(512, num_classes),
 )
 
-state_dict = torch.load('model.pth', map_location='cpu')
-model.load_state_dict(state_dict, strict=False)
+model_path = hf_hub_download(repo_id="jerilseb/quickdraw-small", filename="model.pth")
+state_dict = torch.load(model_path, map_location="cpu")
+model.load_state_dict(state_dict)
 model.eval()
 
-transform = transforms.Compose([
-    transforms.Resize((28, 28)),
-    transforms.ToTensor(),
-    transforms.Normalize((0.5,), (0.5,))
-])
 
+transform = transforms.Compose(
+    [
+        transforms.Resize((28, 28)),
+        transforms.ToTensor(),
+        transforms.Normalize((0.5,), (0.5,)),
+    ]
+)
 
 def predict(image):
     image = image['composite']
-    image = Image.fromarray(image).convert('L')
-    input = transform(image).unsqueeze(0)
-
+    tensor = transform(image).unsqueeze(0)
     with torch.no_grad():
-        out = model(input)
+        out = model(tensor)
 
-    print(out.shape)
     probabilities = F.softmax(out[0], dim=0)
     values, indices = torch.topk(probabilities, 5)
-    print(values, indices)
-
     return {LABELS[i]: v.item() for i, v in zip(indices, values)}
 
 
-interface = gr.Interface(predict, inputs='sketchpad', outputs='label', live=True)
-interface.launch(debug=True)
+inputs = gr.ImageEditor(
+    type="pil",
+    height=720,
+    width=720,
+    layers=False,
+    image_mode="L",
+    brush=gr.Brush(default_color="white", default_size=20),
+    sources=[],
+    label="Draw a shape",
+)
+demo = gr.Interface(predict, inputs=inputs, outputs="label", live=True)
+demo.launch(debug=True)

model.pth

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

requirements.txt

@@ -1,3 +1,4 @@
-torch
+torch==2.3.1
 torchvision
-Pillow
+gradio
+huggingface_hub