Create app.py

app.py

@@ -0,0 +1,90 @@
+# app.py
+
+# Import necessary libraries
+import numpy as np
+import gradio as gr
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+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
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        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
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        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,))])
+
+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()
+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')
+
+# Load the trained model
+model.load_state_dict(torch.load('mnist_model.pth'))
+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
+    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()
+
+    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."
+)
+
+# Launch the interface
+if __name__ == '__main__':
+    iface.launch()