Add app.py & examples

app.py

@@ -1,7 +1,47 @@
+import torch
 import gradio as gr
+from PIL import Image
+from torch import nn
+from torchvision import transforms
 
-def greet(name):
-    return "Hello " + name + "!!"
+classes = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+
+class NeuralNetwork(nn.Module):
+    def __init__(self):
+        super(NeuralNetwork, self).__init__()
+        self.flatten = nn.Flatten()
+        self.linear_relu_stack = nn.Sequential(
+            nn.Linear(28 * 28, 784),
+            nn.ReLU(),
+            nn.Linear(784, 784),
+            nn.ReLU(),
+            nn.Linear(784, 10)
+        )
+
+    def forward(self, x):
+        x = self.flatten(x)
+        logits = self.linear_relu_stack(x)
+        return logits
+
+
+model = NeuralNetwork()
+model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')))
+model.eval()
+
+
+def image_classifier(img_input):
+    img = Image.fromarray(img_input.astype('uint8'), 'RGB')
+    img = transforms.ToTensor()(img)
+
+    with torch.no_grad():
+        pred = model(img)[0]
+        pred -= pred.min()
+        pred /= pred.max()
+        return {classes[i]: float(pred[i]) for i in range(10)}
+
+
+gr.Interface(fn=image_classifier,
+             inputs=gr.Image(shape=(28, 28)),
+             outputs="label",
+             examples=["mnist_0.png", "mnist_2.png", "mnist_3.png"]).launch()