Update app.py

app.py

@@ -9,9 +9,15 @@
 
 # Setting up the Sketch Recognition Model
 
+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(),
@@ -27,30 +33,22 @@ model = nn.Sequential(
     nn.ReLU(),
     nn.Linear(256, len(LABELS)),
 )
-state_dict = torch.load('pytorch_model.bin',    map_location='cpu')
+state_dict = torch.load('pytorch_model.bin', map_location='cpu')
 model.load_state_dict(state_dict, strict=False)
 model.eval()
 
-# Defining a predict function
+def predict(im):
+    x = torch.tensor(im, dtype=torch.float32).unsqueeze(0).unsqueeze(0) / 255.
 
-from pathlib import Path
-
-LABELS = Path('class_names.txt').read_text().splitlines()
-
-def predict(img):
-    x = torch.tensor(img, 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)
-    confidences = {LABELS[i]: v.item() for i, v in zip(indices, values)}
-    return confidences
 
-# Creating a Gradio Interface
+    return {LABELS[i]: v.item() for i, v in zip(indices, values)}
 
-import gradio as gr
 
-gr.Interface(fn=predict, 
-             inputs="sketchpad",
-             outputs="label",
-             live=True).launch()
+interface = gr.Interface(predict, inputs='sketchpad', outputs='label', live=True)
+interface.launch(debug=True)