Spaces:
Sleeping
Sleeping
File size: 1,831 Bytes
c5fe2f1 c58a525 1210b12 c5fe2f1 c58a525 c5fe2f1 1210b12 7355da1 c58a525 c5fe2f1 c58a525 7355da1 1210b12 c58a525 c5fe2f1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import gradio as gr
import numpy as np
from PIL import Image
# Define the CNN
class SimpleCNN(nn.Module):
def __init__(self):
super(SimpleCNN, self).__init__()
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
self.fc1 = nn.Linear(64 * 7 * 7, 128)
self.fc2 = nn.Linear(128, 10)
self.relu = nn.ReLU()
def forward(self, x):
x = self.pool(self.relu(self.conv1(x))) # Output: 32x14x14
x = self.pool(self.relu(self.conv2(x))) # Output: 64x7x7
x = x.view(-1, 64 * 7 * 7) # Flattened to: 3136
x = self.relu(self.fc1(x)) # Output: 128
x = self.fc2(x) # Output: 10 logits
return x
# Load the trained model
model = SimpleCNN()
model.load_state_dict(torch.load('mnist_cnn.pth', map_location=torch.device('cpu')))
model.eval()
# Define the transformation for the input image
transform = transforms.Compose([
transforms.Grayscale(num_output_channels=1),
transforms.Resize((28, 28)),
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
# Prediction function
# Prediction function
def predict(image):
image = transform(image).unsqueeze(0) # Add batch dimension
with torch.no_grad():
output = model(image)
probabilities = nn.Softmax(dim=1)(output)
predicted_class = torch.argmax(probabilities, dim=1)
return {str(i): probabilities[0][i].item() for i in range(10)}
# Create the Gradio interface
interface = gr.Interface(
fn=predict,
inputs=gr.Sketchpad(type='pil'),
outputs=gr.Label(num_top_classes=10)
)
# Launch the interface
interface.launch()
|