Update app.py

app.py

@@ -1,14 +1,15 @@
-# Import necessary libraries
-import numpy as np
 import os
+import numpy as np
 import gradio as gr
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torchvision import transforms
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
 from PIL import Image
 
-# Define the neural network model using PyTorch
+# Define the neural network model
 class Net(nn.Module):
     def __init__(self):
         super(Net, self).__init__()
@@ -23,26 +24,43 @@ class Net(nn.Module):
         x = self.fc3(x)
         return F.log_softmax(x, dim=1)
 
-# Initialize the model and load the trained weights
+# Load and preprocess the MNIST dataset
+transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
+
+train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
+train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+
+test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
+test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
+
+# Initialize the model, loss function, and optimizer
 model = Net()
-model.load_state_dict(torch.load('mnist_model.pth'))
-model.eval()
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=0.001)
 
-# Define the image transformations
-transform = transforms.Compose([
-    transforms.Resize((28, 28)),  # Resize image to 28x28
-    transforms.Grayscale(),       # Convert to grayscale
-    transforms.ToTensor(),        # Convert to tensor
-    transforms.Normalize((0.5,), (0.5,))  # Normalize
-])
+# Check if the model file exists
+model_path = 'mnist_model.pth'
+if not os.path.isfile(model_path):
+    raise FileNotFoundError(f"The model file '{model_path}' was not found.")
+
+# Load the model state dict
+model.load_state_dict(torch.load(model_path, weights_only=True))
+model.eval()
 
-# Define the prediction function
+# Define the predict function
 def predict_image(img):
-    img = transform(img)  # Apply transformations
-    img = img.unsqueeze(0)  # Add batch dimension
+    # 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
+    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).item()
+        predicted_digit = output.argmax(dim=1, keepdim=True).item()
+
     return predicted_digit
 
 # Create the Gradio interface
@@ -54,18 +72,6 @@ iface = gr.Interface(
     description="Upload an image of a handwritten digit, and the model will predict the digit."
 )
 
-# Check if the file exists
-if not os.path.isfile('mnist_model.pth'):
-    raise FileNotFoundError("The model file 'mnist_model.pth' was not found.")
-else:
-    print("Model file found, proceeding with loading.")
-
-# Load the model state dict
-model.load_state_dict(torch.load('mnist_model.pth'))
-
-model.load_state_dict(torch.load('mnist_model.pth', weights_only=True))
-
-
-# Launch the Gradio interface
+# Launch the interface
 if __name__ == '__main__':
     iface.launch()