Update app.py

app.py

@@ -13,23 +13,32 @@ from PIL import Image
 class Net(nn.Module):
     def __init__(self):
         super(Net, self).__init__()
+        # Define layers of the neural network
         self.fc1 = nn.Linear(28 * 28, 128)
         self.fc2 = nn.Linear(128, 64)
         self.fc3 = nn.Linear(64, 10)
 
     def forward(self, x):
-        x = x.view(-1, 28 * 28)  # Flatten the input
+        # Flatten the input tensor
+        x = x.view(-1, 28 * 28)
+        # Apply ReLU activation function
         x = F.relu(self.fc1(x))
         x = F.relu(self.fc2(x))
+        # Output layer with log softmax activation
         x = self.fc3(x)
         return F.log_softmax(x, dim=1)
 
 # Load and preprocess the MNIST dataset
-transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
+transform = transforms.Compose([
+    transforms.ToTensor(), 
+    transforms.Normalize((0.5,), (0.5,))
+])
 
+# Download and load training dataset
 train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
 train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
 
+# Download and load test dataset
 test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
 test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
 
@@ -38,8 +47,10 @@ model = Net()
 criterion = nn.CrossEntropyLoss()
 optimizer = optim.Adam(model.parameters(), lr=0.001)
 
-# Load the model file
-model_path = 'mnist_model.pth'  # Adjust this path if needed
+# Path to the model file
+model_path = 'mnist_model.pth'
+
+# Check if the model file exists
 if not os.path.isfile(model_path):
     raise FileNotFoundError(f"The model file '{model_path}' was not found.")
 
@@ -49,29 +60,29 @@ model.eval()
 
 # Define the predict function
 def predict_image(img):
-    # Preprocess the image
-    img = img.convert('L')
-    img = img.resize((28, 28))
-    img = np.array(img).astype('float32') / 255.0
-    img = (img - 0.5) / 0.5  # Normalize
+    # Preprocess the uploaded image
+    img = img.convert('L')  # Convert image to grayscale
+    img = img.resize((28, 28))  # Resize image to 28x28 pixels
+    img = np.array(img).astype('float32') / 255.0  # Normalize pixel values
+    img = (img - 0.5) / 0.5  # Normalize to range [-1, 1]
     img = torch.tensor(img).unsqueeze(0).unsqueeze(0)  # Add batch and channel dimensions
 
     # Make a prediction
     with torch.no_grad():
-        output = model(img)
-        predicted_digit = output.argmax(dim=1, keepdim=True).item()
+        output = model(img)  # Forward pass through the model
+        predicted_digit = output.argmax(dim=1, keepdim=True).item()  # Get the predicted digit
 
     return predicted_digit
 
 # Create the Gradio interface
 iface = gr.Interface(
-    fn=predict_image,
-    inputs=gr.inputs.Image(shape=(28, 28), image_mode='L', invert_colors=False),
-    outputs='label',
-    live=True,
-    description="Upload an image of a handwritten digit, and the model will predict the digit."
+    fn=predict_image,  # Function to be called on image upload
+    inputs=gr.inputs.Image(shape=(28, 28), image_mode='L', invert_colors=False),  # Input format
+    outputs='label',  # Output format
+    live=True,  # Live update
+    description="Upload an image of a handwritten digit, and the model will predict the digit."  # Description of the interface
 )
 
-# Launch the interface
+# Launch the Gradio interface
 if __name__ == '__main__':
     iface.launch()