Upload 2 files

demo/app.py

@@ -0,0 +1,61 @@
+import torch
+from torch import nn
+from torchvision import transforms
+from PIL import Image
+
+# Define the model architecture
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 32, 3, 1)
+        self.conv2 = nn.Conv2d(32, 64, 3, 1)
+        self.dropout1 = nn.Dropout(0.25)
+        self.dropout2 = nn.Dropout(0.5)
+        self.fc1 = nn.Linear(9216, 128)
+        self.fc2 = nn.Linear(128, 10)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = torch.relu(x)
+        x = self.conv2(x)
+        x = torch.relu(x)
+        x = torch.max_pool2d(x, 2)
+        x = self.dropout1(x)
+        x = torch.flatten(x, 1)
+        x = self.fc1(x)
+        x = torch.relu(x)
+        x = self.dropout2(x)
+        x = self.fc2(x)
+        output = torch.log_softmax(x, dim=1)
+        return output
+
+# Load the trained model
+model = Net()
+#model.load_state_dict(torch.load('mnist-cnn.pth'))  # Load weights
+model.load_state_dict(torch.load('mnist-cnn.pth', weights_only=True))  # Load weights
+
+# Set the model to evaluation mode
+model.eval()
+
+# Function to load and preprocess the image
+def preprocess_image(image_path):
+    img = Image.open(image_path).convert('L')  # Convert to grayscale
+    transform = transforms.Compose([
+        transforms.Resize((28, 28)),
+        transforms.ToTensor(),
+        transforms.Normalize((0.1307,), (0.3081,))
+    ])
+    img_tensor = transform(img).unsqueeze(0)  # Add batch dimension
+    return img_tensor
+
+# Load and preprocess a new image
+image_path = "test_digit.jpg"  # Replace with your image file path
+input_image = preprocess_image(image_path)
+
+# Make the prediction
+with torch.no_grad():
+    outputs = model(input_image)
+    predicted_class = torch.argmax(outputs, dim=1).item()
+
+# Output the predicted digit
+print(f"Predicted digit: {predicted_class}")