Update app.py

app.py

@@ -1,10 +1,8 @@
-# train_model.py
-
 import torch
 import torch.nn as nn
-import torch.optim as optim
-from torchvision import datasets, transforms
-from torch.utils.data import DataLoader
+import torchvision.transforms as transforms
+import gradio as gr
+from PIL import Image
 
 # Define the neural network model
 class Net(nn.Module):
@@ -21,33 +19,35 @@ class Net(nn.Module):
         x = self.fc3(x)
         return torch.log_softmax(x, dim=1)
 
-# Load and preprocess the MNIST dataset
+# Load the trained model
+model = Net()
+model.load_state_dict(torch.load('mnist_model.pth', map_location=torch.device('cpu')))
+model.eval()
+
+# Define the transform to preprocess the input image
 transform = transforms.Compose([
+    transforms.Grayscale(num_output_channels=1),
+    transforms.Resize((28, 28)),
     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)
-
-# Initialize the model, loss function, and optimizer
-model = Net()
-criterion = nn.CrossEntropyLoss()
-optimizer = optim.Adam(model.parameters(), lr=0.001)
-
-# Train the model
-def train(model, train_loader, criterion, optimizer, epochs=5):
-    model.train()
-    for epoch in range(epochs):
-        for data, target in train_loader:
-            optimizer.zero_grad()
-            output = model(data)
-            loss = criterion(output, target)
-            loss.backward()
-            optimizer.step()
-
-train(model, train_loader, criterion, optimizer)
-
-# Save the trained model
-torch.save(model.state_dict(), 'mnist_model.pth')
-print("Model saved as 'mnist_model.pth'")
+# Define the prediction function
+def predict(image):
+    image = transform(image).unsqueeze(0)  # Add batch dimension
+    with torch.no_grad():
+        output = model(image)
+        prediction = torch.argmax(output, dim=1).item()
+    return prediction
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=predict, 
+    inputs=gr.inputs.Image(shape=(28, 28), image_mode='L', invert_colors=False), 
+    outputs="label",
+    live=True
+)
+
+# Launch the Gradio interface
+if __name__ == "__main__":
+    iface.launch()