pyCnnImageClassifier/nn_image_classification.ipynb

In [17]:

import numpy as np
from PIL import Image
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

In [6]:

transform = transforms.Compose([transforms.ToTensor(), import numpy as np
from PIL import Image
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as opim
import torchvision
import torchvision.transform as transform

])                              

In [9]:

train_data = torchvision.datasets.CIFAR10(root='./data', train=True, transform=transform, download=True)
test_data = torchvision.datasets.CIFAR10(root='./data', train=False, transform=transform, download=True)

train_loader = torch.utils.data.DataLoader(train_data, batch_size=32, shuffle=True, num_workers=2)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=32, shuffle=True, num_workers=2)

In [10]:

image, label = train_data[0]

In [11]:

image.size()

Out [11]:

torch.Size([3, 32, 32])

In [12]:

class_names = ['plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']

In [ ]:

class NeuralNet(nn.Module):
    def  __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 12, 5) # (12, 28, 28)
        self.pool = nn.MaxPool2d(2, 2) # (12, 14, 14)
        self.conv2 = nn.Conv2d(12, 24, 5) # (24, 10, 10) -> (24, 5, 5) -> Flatten (25 * 5 * 5)
        self.fc1 = nn.Linear(24 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

In [18]:

net = NeuralNet()
loss_function = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

In [19]:

for epoch in range(30):
    print(f'Training epoch {epoch}...')

    running_loss = 0.0

    for i, data in enumerate(train_loader):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()

    print(f'Loss: {running_loss / len(train_loader):0.4f}')

Training epoch 0...
Loss: 2.1758
Training epoch 1...
Loss: 1.7324
Training epoch 2...
Loss: 1.5096
Training epoch 3...
Loss: 1.3996
Training epoch 4...
Loss: 1.3065
Training epoch 5...
Loss: 1.2286
Training epoch 6...
Loss: 1.1649
Training epoch 7...
Loss: 1.1101
Training epoch 8...
Loss: 1.0568
Training epoch 9...
Loss: 1.0140
Training epoch 10...
Loss: 0.9719
Training epoch 11...
Loss: 0.9321
Training epoch 12...
Loss: 0.8969
Training epoch 13...
Loss: 0.8649
Training epoch 14...
Loss: 0.8326
Training epoch 15...
Loss: 0.8014
Training epoch 16...
Loss: 0.7764
Training epoch 17...
Loss: 0.7449
Training epoch 18...
Loss: 0.7199
Training epoch 19...
Loss: 0.6991
Training epoch 20...
Loss: 0.6700
Training epoch 21...
Loss: 0.6438
Training epoch 22...
Loss: 0.6243
Training epoch 23...
Loss: 0.6051
Training epoch 24...
Loss: 0.5802
Training epoch 25...
Loss: 0.5609
Training epoch 26...
Loss: 0.5381
Training epoch 27...
Loss: 0.5162
Training epoch 28...
Loss: 0.4949
Training epoch 29...
Loss: 0.4752

In [20]:

torch.save(net.state_dict(), 'trained_net.pth')

In [23]:

net = NeuralNet()
net.load_state_dict(torch.load('trained_net.pth'))

Out [23]:

<All keys matched successfully>

In [24]:

correct = 0
total = 0
net.eval()
with torch.no_grad():
    for data in test_loader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f'Accuracy: {accuracy}%')

Accuracy: 68.23%

In [32]:

new_transform = transforms.Compose([
    transforms.Resize((32, 32)),  
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

def load_image(image_path):
    image = Image.open(image_path)
    image = new_transform(image)
    image = image.unsqueeze(0)
    return image

image_paths = ['item1.jpg', 'item2.jpg', 'item3.jpg']
images = [load_image(img) for img in image_paths]

net.eval()
with torch.no_grad():
    for image in images:
        output = net(image)
        _, predicted = torch.max(output, 1)
        print(f'Prediction: {class_names[predicted.item()]}')

Prediction: cat
Prediction: deer
Prediction: plane