Multi class Classification

We have three classes of images : 7, 8 and 9. The following codes uses VGG16 and CNN to build a model which classifies the images of 7, 8 and 9.

Import necessary libraries as usual

import torchvision
import torch.nn as nn
import torch
import torch.nn.functional as F
from torchvision import transforms,models,datasets
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
from torch import optim
device = 'cuda' if torch.cuda.is_available() else 'cpu'
import cv2, glob, numpy as np, pandas as pd
import matplotlib.pyplot as plt
from glob import glob
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset

Load Data: Source of data 7_8_9

train_data_dir = '.../7_8_9/Data/7_8_9/Train'
val_data_dir = '.../7_8_9/Data/7_8_9/Val'

Class to load data and make targets

class s_e_n(Dataset):
    def __init__(self, folder):
        seven = glob(folder+'/7/*.PNG')
        eight = glob(folder+'/8/*.PNG')
        nine = glob(folder+'/9/*.PNG')
        self.fpaths = seven + eight +nine
        self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
        from random import shuffle, seed; seed(10); shuffle(self.fpaths)
       
        targets = []
        for fpath in self.fpaths:
          if(fpath.split('/')[-2].startswith('7')):
            target = 0
          elif ((fpath.split('/')[-2].startswith('8'))):
            target = 1
          elif ((fpath.split('/')[-2].startswith('9'))):
            target = 2
        
          targets.append(target)

        self.targets = targets

        
    def __len__(self): return len(self.fpaths)
    def __getitem__(self, ix):
        f = self.fpaths[ix]
        target = self.targets[ix]

        #im = (cv2.imread(f)[:,:,::-1])
        im = cv2.imread(f)
       
        im = im[:,:,::-1]
       
        im = cv2.resize(im, (224,224))
       
        im = torch.tensor(im/255)

      
        im = im.permute(2,0,1)

      
        im = self.normalize(im) 

     
        return im.float().to(device), torch.tensor([target]).to(device)

Load data and display

data = s_e_n(train_data_dir)
im, label = data[2]
plt.imshow(im.permute(1,2,0).cpu())
print(label)

Vgg16 model

def get_model():
    model = models.vgg16(pretrained=True)
    for param in model.parameters():
        param.requires_grad = False
    model.avgpool = nn.AdaptiveAvgPool2d(output_size=(1,1))
    model.classifier = nn.Sequential(nn.Flatten(),
    nn.Linear(512, 128),
    nn.ReLU(),
    nn.Dropout(0.2),
    nn.Linear(128, 3),
    nn.Sigmoid())
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr= 1e-3)
    return model.to(device), loss_fn, optimizer

Summary of model

#Printing summary of model
from torchsummary import summary
model, loss_fn, optimizer = get_model()
summary(model, input_size=(3, 224, 224))

Function to train images per batch, calculate accuracy, load data and calculate validation loss

def train_batch(x, y, model, opt, loss_fn):
    model.train()
    prediction = model(x)
    batch_loss = loss_fn(prediction, y)
    batch_loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    return batch_loss.item()

@torch.no_grad()
def accuracy(x, y, model):
    z = model(x) 
    y_hat = torch.max(z.data,1)
    is_correct = y_hat[1][0] == label
    return is_correct.cpu().numpy().tolist()

def get_data():     
    train = s_e_n(train_data_dir)
    trn_dl = DataLoader(train, batch_size=32, shuffle=True, drop_last = True)
    val = s_e_n(val_data_dir)
    val_dl = DataLoader(val, batch_size=32, shuffle=True, drop_last = True)
    return trn_dl, val_dl

@torch.no_grad()
def val_loss(x, y, model):
    prediction = model(x)
    val_loss = loss_fn(prediction, y)
    return val_loss.item()

Load data and models

trn_dl, val_dl = get_data()
model, loss_fn, optimizer = get_model()

Train model up to 10 epochs

train_losses, train_accuracies = [], []
val_losses, val_accuracies = [], []
for epoch in range(10):
    
    print(epoch)
    train_epoch_losses, train_epoch_accuracies = [], []
    val_epoch_accuracies = []
    for ix, batch in enumerate(iter(trn_dl)):
        #print(ix)
        x, y = batch
        #print(y)
        #print(y.shape)
        y = y.reshape(-1)
        #y = y.type(torch.LongTensor)
       # print(batch)
        batch_loss = train_batch(x, y, model, optimizer, loss_fn)
        #print(y)
        train_epoch_losses.append(batch_loss)        
    train_epoch_loss = np.array(train_epoch_losses).mean()

    for ix, batch in enumerate(iter(trn_dl)):
        x, y = batch
        is_correct = accuracy(x, y, model)
        train_epoch_accuracies.extend(is_correct)
    train_epoch_accuracy = np.mean(train_epoch_accuracies)

    for ix, batch in enumerate(iter(val_dl)):
        x, y = batch
        val_is_correct = accuracy(x, y, model)
        val_epoch_accuracies.extend(val_is_correct)
        y = y.reshape(-1)
        validation_loss = val_loss(x, y, model)
    val_epoch_accuracy = np.mean(val_epoch_accuracies)

    train_losses.append(train_epoch_loss)
    train_accuracies.append(train_epoch_accuracy)
    val_losses.append(validation_loss)
    val_accuracies.append(val_epoch_accuracy)

Plot loss and accuracy

epochs = np.arange(10)+1
import matplotlib.ticker as mtick
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
%matplotlib inline
plt.rcParams["figure.figsize"] = (10,5)
plt.subplot(211)
plt.plot(epochs, train_losses, 'bo', label='Training loss')
plt.plot(epochs, val_losses, 'r', label='Validation loss')
plt.gca().xaxis.set_major_locator(mticker.MultipleLocator(1))
plt.title('Training and validation loss with CNN')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.grid('off')
plt.show()
plt.subplot(212)
plt.plot(epochs, train_accuracies, 'bo', label='Training accuracy')
plt.plot(epochs, val_accuracies, 'r', label='Validation accuracy')
plt.gca().xaxis.set_major_locator(mticker.MultipleLocator(1))
plt.title('Training and validation accuracy with CNN')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.gca().set_yticklabels(['{:.0f}%'.format(x*100) for x in plt.gca().get_yticks()]) 
plt.legend()
plt.grid('off')
plt.show()

Save model

#save model
model_scripted = torch.jit.script(model) # Export to TorchScript
model_scripted.save('.../789_model_epoch_10_scripted.pt') 

Load model and test

import torch
device = 'cuda' if torch.cuda.is_available() else 'cpu'
import cv2, glob, numpy as np, pandas as pd
import matplotlib.pyplot as plt
%matplotlib inline
from glob import glob
from google.colab import files

import warnings
warnings.filterwarnings('ignore')

model = torch.jit.load('.../789_model_epoch_10_scripted.pt')
model.eval()

Function to predict

def predict(img_file):

  im = (cv2.imread(file_name)[:,:,::-1])

  im = cv2.resize(im, (224,224))
  im = torch.tensor(im/255).permute(2,0,1).to(device).float()
  img = im
  img = img.view(3,224, 224)
  img2  = torch.tensor(img).to(device).float()
  img3 = torch.tensor(img2).view(-1,3,224,224)
  np_output = model(img3)
  print(np_output)

  y_hat = torch.max(np_output.data,1)
  print(y_hat[1][0])
  print(type(y_hat[1][0]))
  
  if((y_hat[1][0])==0):
    prediction = 'Seven'
  elif ((y_hat[1][0])==1):
    prediction = 'Eight'
  else:
    prediction = 'Nine'
  plt.imshow(im.permute(1,2,0).cpu())

  return prediction

Load image

uploaded = files.upload()
file_name = uploaded.keys()
file_name = list(file_name)[0]
predict(file_name)

When model got it wrong