Update app.py

app.py

@@ -32,25 +32,22 @@ 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 sketch is not None:
-        # Procesamiento del dibujo
-        sketch = transforms.ToPILImage()(sketch).convert("L").resize((28, 28))
-        sketch = PIL.ImageOps.invert(sketch)
+    if image is not None:
+        # Preprocess of the image
         transform = transforms.Compose([
-            transforms.ToTensor(),
-            transforms.Normalize((0.5,), (0.5,))
-        ])
-        sketch = transform(sketch).unsqueeze(0)
-        
-        with torch.no_grad():
-            prediction = load_model(sketch)
-        
+                        transforms.ToTensor(),
+                        transforms.Normalize((0.5,), (0.5,))
+                    ])
+        image = transform(image)
+        with torch.inference_mode(): # inference mode of pytoroch
+            prediction = load_model(image)
         prediction = torch.softmax(prediction, dim=1)
         return {str(i): float(prediction[0][i]) for i in range(10)}
     else:
         return ""
+    
 
 demo = gr.Interface(fn=recognize_digit, 
-                    inputs="sketchpad", 
-                    outputs=gr.outputs.Label(num_top_classes=3)
-demo.launch(share=True)
+                    inputs=gr.Image(shape=(28,28), image_mode="L", invert_colors=True, source="canvas"), 
+                    outputs=gr.Label(num_top_classes=1))
+demo.launch(True)