updates desc

app.py

@@ -1,12 +1,10 @@
 import gradio as gr
-import pickle
 import torch
 import numpy as np
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 with open('cnn_model.bin', 'rb') as f:
-    # nn = pickle.load(f)
     nn = torch.load(f, map_location=torch.device('cpu'))
 
 nn.to(device)
@@ -24,10 +22,16 @@ def predict(input):
     return dict(enumerate(p.tolist()))
 
 
+desc = """\
+This project uses a Convolutional Neural Network to classify handwritten digits.
+Trained on the MNIST dataset.
+Use most of the drawing area for better results.
+"""
+
 demo = gr.Interface(
     fn=predict,
     title='ConvNet for handwritten digits classification',
-    description='Created with PyTorch.',
+    description=desc,
     inputs=[
         gr.Sketchpad(
             shape=(28, 28),

networktorch.py

@@ -1,50 +0,0 @@
-from torch import nn
-
-
-class NeuralNetworkTorch(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-        self.stack = nn.Sequential(
-            nn.Linear(784, 64),
-            nn.Sigmoid(),
-
-            nn.Linear(64, 10),
-            nn.Sigmoid()
-        )
-
-    def forward(self, x):
-        return self.stack(x)
-
-
-class ConvNeuralNetworkTorch(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-        self.conv = nn.Sequential(
-            nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1),
-            nn.ReLU(),
-
-            nn.MaxPool2d(kernel_size=2, stride=2),
-
-            nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1),
-            nn.ReLU(),
-
-            # nn.MaxPool2d(kernel_size=2, stride=2),
-        )
-
-        self.fc = nn.Sequential(
-            nn.Linear(16 * 14 * 14, 10),
-            nn.Sigmoid(),
-        )
-
-    def forward(self, x):
-        # we do some reshaping here simply to avoid making changes to the caller
-        # so it continues to work with the fully conected network above
-        x = x.reshape(-1, 1, 28, 28) / 255
-
-        conv_output = self.conv(x)
-        flat = conv_output.reshape(len(x), -1)
-        final_output = self.fc(flat)
-
-        return final_output