Update app.py

app.py

@@ -7,6 +7,8 @@
 #iface = gr.Interface(fn=greet, inputs="text", outputs="text")
 #iface.launch()
 
+# Setting up the Sketch Recognition Model
+
 import torch
 from torch import nn
 
@@ -28,3 +30,27 @@ model = nn.Sequential(
 state_dict = torch.load('pytorch_model.bin',    map_location='cpu')
 model.load_state_dict(state_dict, strict=False)
 model.eval()
+
+# Defining a predict function
+
+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
+
+import gradio as gr
+
+gr.Interface(fn=predict, 
+             inputs="sketchpad",
+             outputs="label",
+             live=True).launch()