Update app.py

app.py

@@ -32,15 +32,19 @@ PATH = "state_dict_model.pth" # PATH where you load the model trained
 load_model.load_state_dict(torch.load(PATH))
 load_model.eval()
 def recognize_digit(image):
-    if image is not None:
-        # Preprocess of the image
+    if sketch is not None:
+        # Procesamiento del dibujo
+        sketch = transforms.ToPILImage()(sketch).convert("L").resize((28, 28))
+        sketch = PIL.ImageOps.invert(sketch)
         transform = transforms.Compose([
-                        transforms.ToTensor(),
-                        transforms.Normalize((0.5,), (0.5,))
-                    ])
-        image = transform(image)
-        with torch.inference_mode(): # inference mode of pytoroch
-            prediction = load_model(image)
+            transforms.ToTensor(),
+            transforms.Normalize((0.5,), (0.5,))
+        ])
+        sketch = transform(sketch).unsqueeze(0)
+        
+        with torch.no_grad():
+            prediction = load_model(sketch)
+        
         prediction = torch.softmax(prediction, dim=1)
         return {str(i): float(prediction[0][i]) for i in range(10)}
     else: