number-detector.py

import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.datasets as dsets 
from torch.autograd import Variable
import numpy as np

from PIL import Image
import io
# import cv2

testWithRealImage = True
IMG_URL = "./test-img.jpg"

# ----(1)Get Data----
train_dataset = dsets.MNIST(root='./data',
                            train=True,
                            transform=transforms.ToTensor(),
                            download=True)

test_dataset = dsets.MNIST(root='./data',
                            train=False,
                            transform=transforms.ToTensor())

# amount of images loaded at one time
batch_size = 100
# number if iterations
n_iters = 3000
# bass in number of images, then every epoch we go back and update the model
num_epochs = n_iters / (len(train_dataset) / batch_size)
num_epochs = int(num_epochs)

# make data iterable
# Shuffle is true so we don't go through the data in a certain order
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                            batch_size=batch_size,
                                            shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                            batch_size=batch_size,
                                            shuffle=False)

# Create the model class
class FeedForwardNeuralNetModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(FeedForwardNeuralNetModel, self).__init__()
        # Linear function
        self.fcl = nn.Linear(input_dim, hidden_dim)
        # Non-Linear function
        self.sigmoid = nn.Sigmoid()
        # Linear function
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    # x is the image
    def forward(self, x):
        # Linear function
        out = self.fcl(x)
        # Non-Linear function
        out = self.sigmoid(out)
        # Linear function
        out = self.fc2(out)
        return out

# size of image
input_dim = 28*28
# Number of neurons
hidden_dim = 1000
# because we are detecting numbers from 1-10
output_dim = 10

# instantiate the model
model = FeedForwardNeuralNetModel(input_dim, hidden_dim, output_dim)

# instantiate loss class
# Doss soft max and cross entropy
# gives us our gradients
criteron = nn.CrossEntropyLoss()

learning_rate = 0.1
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

if testWithRealImage == True:
    images = torch.Tensor(1, 1, 28, 28)
    img = Image.open( IMG_URL )
    img = img.convert('1')
    data = np.array( img, dtype='uint8' )
    # print(data)
    images[0][0] = torch.from_numpy(data)
    tensor = Variable(images.view(-1, 28*28))
    # print(tensor)
    outputs = model.forward(tensor)
    _, predicted = torch.max(outputs.data, 1)
    print(predicted)
    exit()

iter = 0
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        # Load images as Varaible
        images = Variable(images.view(-1, 28*28))
        labels = Variable(labels)

        # clear gradients
        optimizer.zero_grad()

        # forward pass
        outputs = model(images)

        loss = criteron(outputs, labels)
        loss.backward()

        # update parameters
        optimizer.step()

        iter += 1

        if iter % 500 == 0:
            # Calculate accuracy
            correct = 0
            total = 0
            # Iterate through the dataset
            for images, labels in test_loader:
                images = Variable(images.view(-1, 28*28))
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum()
            
            accuracy = 100 * correct / total

            print('Iteration: {}. Loss: {}. Accuracy: {}'.format(iter, loss.data[0], accuracy))

# save model
# torch.save(model.state_dict(), 'feedforward-numrec.pkl')