Ants and Bees using CNN

We are trying to build a model using CNN with the help of PyTorch library that can differentiate between images of Ants and Bees.

Data used in this task is available at : https://www.kaggle.com/watanabe2362/hymenoptera

This process is divided into two parts:

• Part I : Building the model and saving it.
• Part II: Loading the model and using it.

Part I:

Importing necessary libraries

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

#connecting python script to gpu
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
!pip install torchsummary

Load Data

train_data_dir = '.../Data/hymenoptera_data/train'
val_data_dir = '.../Data/hymenoptera_data/val'

Making a class which will help to load and label data

from torch.utils.data import DataLoader, Dataset
class ants_bees(Dataset):
  def __init__(self, folder):
    ants = glob(folder+'/ants/*.jpg')
    bees = glob(folder+'/bees/*.jpg')
    self.fpaths = ants + bees
    from random import shuffle, seed; seed(10); shuffle(self.fpaths)
    self.targets = [fpath.split('/')[-2].startswith('bees') for fpath in self.fpaths] # ants=0 & bees=1
  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.resize(im, (224,224))
    return torch.tensor(im/255).permute(2,0,1).to(device).float(), torch.tensor([target]).float().to(device)

Checking and Visualizing Data

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

Defining Neural Network Layers

#Defining convolution layer
def conv_layer(ni,no,kernel_size,stride=1):
    return nn.Sequential(
    nn.Conv2d(ni, no, kernel_size, stride),
    nn.ReLU(),
    nn.BatchNorm2d(no),
    nn.MaxPool2d(2)
    )

#Defining deep CNN model
def get_model():
    model = nn.Sequential(
    conv_layer(3, 64, 3),
    conv_layer(64, 512, 3),
    conv_layer(512, 512, 3),
    conv_layer(512, 512, 3),
    conv_layer(512, 512, 3),
    conv_layer(512, 512, 3),
    nn.Flatten(),
    nn.Linear(512, 1),
    nn.Sigmoid(),
    ).to(device)

    #Binary cross entropy is used to calculate loss
    loss_fn = nn.BCELoss()
    #Adam us used as an optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr= 1e-3)
    return model, loss_fn, optimizer

Summary of the Model

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

Defining functions for loading training and accuracy

def train_batch(x, y, model, opt, loss_fn):
    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):
    prediction = model(x)
    is_correct = (prediction > 0.5) == y
    return is_correct.cpu().numpy().tolist()

Functions to load data and to calculate validation loss

def get_data():     
    train = ants_bees(train_data_dir)
    trn_dl = DataLoader(train, batch_size=32, shuffle=True, drop_last = True)
    val = ants_bees(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()

Loading Data and model along with its parameters

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

Training for 10 epochs; Calculation train accuracy, validation accuracy, train loss and validation loss

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
        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)
        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)

Loss and Accuracy after training for 10 epochs

print(f'Training loss after 10 epochs: {train_epoch_loss}\n')
print(f'Validation loss after 10 epochs: {validation_loss}\n')
print(f'Training accuracy after 10 epochs: {train_epoch_accuracy*100}\n')
print(f'Validation accuracy after 10 epochs: {val_epoch_accuracy*100}\n')

Plotting graph of accuracy and losses

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('.../ants_bees_model_epoch_10_scripted.pt') 

Part II: Loading model and using it

Importing Files

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')

Loading model

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

Define Function which takes image as an input and predicts whether its an ant or a bee

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).cpu().detach().numpy()
  print(np_output)
  if((np.round(np_output[0][0],decimals=0))==1):
    prediction = 'This is a bee'
  else:
    prediction = 'This is an ant'
  plt.imshow(im.permute(1,2,0).cpu())

  return prediction

Loading The image

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

Final Output

When model got it wrong