Update app.py

app.py

@@ -7,51 +7,49 @@ import matplotlib.pyplot as plt
 from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
 import PIL.Image
 
-# Load the model cause it is necessary for pytorch 
-
+# Load the model
 class MyMnist_ModelV0(nn.Module):
-  def __init__(self,  input_shape: int, hidden_units: int, hidden_units2: int, output_shape: int):
+    def __init__(self, input_shape: int, hidden_units: int, hidden_units2: int, output_shape: int):
         super().__init__()
         self.layer_stack = nn.Sequential(
-            nn.Flatten(), # neural networks like their inputs in vector form
-            nn.Linear(in_features=input_shape, out_features=hidden_units),  # in_features = number of features in a data sample (784 pixels)
+            nn.Flatten(), 
+            nn.Linear(in_features=input_shape, out_features=hidden_units), 
             nn.ReLU(),
             nn.Linear(in_features=hidden_units, out_features=hidden_units2),
             nn.ReLU(),
             nn.Linear(in_features=hidden_units2, out_features=output_shape)
         )
 
-  def forward(self, x):
-      return self.layer_stack(x)
+    def forward(self, x):
+        return self.layer_stack(x)
 
-# instance of the model
+# Load the pre-trained model
 load_model = MyMnist_ModelV0(input_shape=784,
-                          hidden_units=256,
-                          hidden_units2=128,
-                          output_shape=10
-)
+                              hidden_units=256,
+                              hidden_units2=128,
+                              output_shape=10)
 
-PATH = "state_dict_model.pth" # PATH where you load the model trained
+PATH = "state_dict_model.pth"  # Path to the trained model
 
 load_model.load_state_dict(torch.load(PATH))
 load_model.eval()
+
+# Function to recognize digit
 def recognize_digit(image):
     if image is not None:
-        # Preprocess of the image
-        image = Image.fromarray(image).resize((28, 28))
-        transform = transforms.Compose([
-                        transforms.ToTensor(),
-                        transforms.Normalize((0.5,), (0.5,))
-                    ])
-        image = PIL.ImageOps.invert(image)  # Invert colors
-        image = transform(image)
-        with torch.inference_mode(): # inference mode of pytoroch
+        # Convert image to grayscale
+        image = np.array(image.convert("L"))
+        # Resize image to 28x28
+        image = torch.tensor(image / 255.0, dtype=torch.float32)
+        # Perform inference
+        with torch.inference_mode():  
             prediction = load_model(image)
         prediction = torch.softmax(prediction, dim=1)
         return {str(i): float(prediction[0][i]) for i in range(10)}
     else:
         return ""
 
+# Function to create a canvas for drawing
 def create_canvas():
     fig, ax = plt.subplots()
     ax.set_title("Draw your digit")
@@ -60,10 +58,13 @@ def create_canvas():
     canvas = FigureCanvas(fig)
     return canvas
 
+# Create canvas
 canvas = create_canvas()
 
+# Define Gradio interface
+demo = gr.Interface(fn=recognize_digit,
+                    inputs=gr.inputs.Image(canvas=canvas),
+                    outputs=gr.outputs.Label(num_top_classes=1))
 
-demo = gr.Interface(fn=recognize_digit, 
-                    inputs=gr.Image(image_mode="L", source="canvas", canvas=canvas), 
-                    outputs=gr.Label(num_top_classes=1))
-demo.launch(share=True)
+# Launch the interface
+demo.launch(share=True)