Upload with huggingface_hub

app.py

@@ -5,32 +5,55 @@ from PIL import ImageOps
 import os
 from dotenv import load_dotenv
 
+from torch import nn
+import torch.nn.functional as F
+
+class SimpleLenet(nn.Module):
+    def __init__(self, args=None):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 6, 5, padding=2) # -> 6 channels, 28x28
+        self.pool = nn.MaxPool2d(2) # -> 6 channels, 14x14
+        self.conv2 = nn.Conv2d(6, 120, 14) #-> 120 channels, 1x1
+        self.fc1 = nn.Linear(120, 10)
+        self.fc2 = nn.Linear(10, 10)
+
+    def __call__(self, x):
+        xx = F.relu(self.conv1(x))
+        xx = F.relu(self.pool(xx))
+        xx = F.relu(self.conv2(xx))
+        xx = xx.flatten(1)
+        xx = F.relu(self.fc1(xx))
+        return self.fc2(xx)
+
 load_dotenv()
 
 hf_writer = gr.HuggingFaceDatasetSaver(os.getenv('HF_TOKEN'), "simple-mnist-flagging")
 
 def load_model():
-    model_dict = torch.load('linear_model.pt')
-    return model_dict
+    model = SimpleLenet()
+    model.load_state_dict(torch.load('model.pt'))
+    model.eval()
+    return model
 
 model = load_model()
 convert_tensor = transforms.ToTensor()
 
 def predict(img):
     img =  ImageOps.grayscale(img).resize((28,28))
-    image_tensor = convert_tensor(img).view(28*28)
-    res = image_tensor @ model['weights'] + model['bias']
-    res = res.sigmoid()
-    return {"It's 3": float(res), "It's 7": float(1-res)}
+    image_tensor = convert_tensor(img).view(1, 1, 28, 28)
+    logits = model(image_tensor)
+    pred = torch.argmax(logits, dim=1)
+    return pred.tolist()[0]
 
-title = "Is it 7 or 3"
-description = '<p><center>Write a number, 7 or 3, in the middle.</center></p>'
+title = "Handwritten digit recognition"
+description = '<p><center>Write a single digit in the middle of the canvas</center></p>'
 
 gr.Interface(fn=predict, 
              inputs=gr.Paint(type="pil", invert_colors=True),
-             outputs=gr.Label(num_top_classes=2),
+             outputs="text",
              title=title,
              flagging_options=["incorrect","ambiguous"],
              flagging_callback=hf_writer,
              description=description,
-             allow_flagging='manual').launch()
+             allow_flagging='manual').launch()
+

model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:922abc05756421cbdffc98dc27a3eef6abf476761170e77656e7bb5477a45b10
+size 573263