MNIST practice

app.py

@@ -0,0 +1,55 @@
+import gradio as gr
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from PIL import Image
+from torchvision import transforms
+
+
+class simpleCNN(nn.Module):
+    def __init__(self, num_classes=3):
+        super(simpleCNN, self).__init__()
+        self.name = "simpleCNN"
+        self.conv1 = nn.Conv2d(3, 5, 5)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(5, 10, 5)
+        self.fc1 = nn.Linear(10 * 5 * 5, 32)
+        self.fc2 = nn.Linear(32, num_classes)
+
+    def forward(self, x):
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = x.view(-1, 10 * 5 * 5)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+
+net = simpleCNN(num_classes=3)
+net.load_state_dict(torch.load("./ckpt.pth", map_location=torch.device("cpu")))
+net.eval()
+class_labels = ["other", "car", "truck"]
+
+transform = transforms.Compose(
+    [
+        transforms.Resize((32, 32)),
+        transforms.ToTensor(),
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+    ]
+)
+
+
+@torch.no_grad()
+def predict(img):
+    global net
+    img = Image.fromarray(img.astype("uint8"), "RGB")
+    img = transform(img).unsqueeze(0)
+    pred = net(img).detach().numpy()[0]
+    pred = np.exp(pred) / np.sum(np.exp(pred))
+    return {class_labels[i]: float(pred[i]) for i in range(len(class_labels))}
+
+
+iface = gr.Interface(fn=predict, inputs="image", outputs="label")
+iface.launch()

ckpt.pth

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