first test run

main.py

@@ -0,0 +1,42 @@
+import torch
+import gradio as gr
+import torchvision.transforms as transforms
+from neural_network import MNISTNetwork
+
+
+transform = transforms.Compose([
+    transforms.ToTensor(),  # Convert image to tensor
+    transforms.Normalize((0.1307,), (0.3081,))  # Normalize the image
+])
+
+# Load the trained model    
+net = MNISTNetwork()
+net.load_state_dict(torch.load('mnist_net.pth'))
+LABELS = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
+
+def predict(drawing): 
+    if drawing is None: 
+        return "Draw something hoe"
+    
+    input_tensor = transform(drawing)
+    x = input_tensor.view(input_tensor.shape[0], -1)
+    
+    with torch.no_grad():
+        output = net(x)
+
+        probabilities = torch.nn.functional.softmax(output[0], dim=0)
+        values, indices = torch.topk(probabilities, 10)
+        results = {LABELS[i]: v.item() for i, v in zip(indices, values)}
+        
+        return results
+
+
+sketchpad_input = gr.Sketchpad(shape=(28, 28))
+interface = gr.Interface(
+    fn=predict, 
+    inputs=sketchpad_input,
+    outputs="label",
+    live=True
+)
+interface.launch()
+

mnist_net.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:37d782b2adae4525684264e311d73f3bb52e250b6bc13d0a9446e8fb45bc715a
+size 446799

neural_network.py

@@ -0,0 +1,39 @@
+import torch.nn as nn 
+import torch.nn.functional as F
+
+class MNISTNetwork(nn.Module): 
+    
+    def __init__(self): 
+        super().__init__()
+        self.layer1 = nn.Linear(784, 128)
+        self.layer2 = nn.Linear(128, 64)
+        self.layer3 = nn.Linear(64, 32)
+        self.layer4 = nn.Linear(32, 10)
+    
+    def forward(self, x):
+        x = F.relu(self.layer1(x))
+        x = F.relu(self.layer2(x))
+        x = F.relu(self.layer3(x))
+        x = self.layer4(x)
+        return F.log_softmax(x, dim=1)
+
+
+
+# class MNISTNetwork(nn.Module):
+#     def __init__(self):
+#         super().__init__()
+#         self.conv1 = nn.Conv2d(1, 32, kernel_size=5, padding=2)
+#         self.conv2 = nn.Conv2d(32, 64, kernel_size=5, padding=2)
+#         self.fc1 = nn.Linear(64*7*7, 1024)
+#         self.fc2 = nn.Linear(1024, 10)
+
+#     def forward(self, x):
+#         x = nn.functional.relu(self.conv1(x))
+#         x = nn.functional.max_pool2d(x, 2)
+#         x = nn.functional.relu(self.conv2(x))
+#         x = nn.functional.max_pool2d(x, 2)
+#         x = x.view(-1, 64*7*7)
+#         x = nn.functional.relu(self.fc1(x))
+#         x = nn.functional.dropout(x, training=self.training)
+#         x = self.fc2(x)
+#         return nn.functional.log_softmax(x, dim=1)