Ants and Bees using VGG16

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

Loading Data

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

Make Class to load data for model

class ants_bees(Dataset):
    def __init__(self, folder):
        ants = glob(folder+'/ants/*.jpg')
        bees = glob(folder+'/bees/*.jpg')
        self.fpaths = ants + bees
        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)
        self.targets = [fpath.split('/')[-2].startswith('bees') for fpath in self.fpaths] 
    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))
        im = torch.tensor(im/255)
        im = im.permute(2,0,1)
        im = self.normalize(im) 
        return im.float().to(device), torch.tensor([target]).float().to(device)

Load data and display image

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

Output

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, 1),
    nn.Sigmoid())
    loss_fn = nn.BCELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr= 1e-3)
    return model.to(device), loss_fn, optimizer

Summarize

!pip install torch_summary
from torchsummary import summary
model, criterion, optimizer = get_model()
summary(model, torch.zeros(1,3,224,224))

Function to train model

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

Function to calculate Accuracy and validation loss

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

Defining batch and function that load files batch-wise

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

Get model and get data

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

Train model for 5 epochs

train_losses, train_accuracies = [], []
val_losses, val_accuracies = [], []
for epoch in range(5):
    print(f" epoch {epoch + 1}/5")
    train_epoch_losses, train_epoch_accuracies = [], []
    val_epoch_accuracies = []

    for ix, batch in enumerate(iter(trn_dl)):
        x, y = batch
        batch_loss = train_batch(x, y, model, optimizer, loss_fn)
        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)

Plotting accuracy and loss

epochs = np.arange(5)+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('.../Trained_model/ants_bees_model_epoch_10_VGG16_scripted.pt') 

Confusion Matrix

from sklearn.metrics import confusion_matrix
import seaborn as sn
import pandas as pd

y_pred = []
y_true = []
count = 0
data = ants_bees(train_data_dir)
print(len(data))
# iterate over test data
for inputs, labels in data:

  im = torch.tensor(inputs).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)
  output = model(img3) # Feed Network
  
  
  output = output.cpu().detach().numpy()
  
  output = np.round(output[0][0],decimals=0)
 

  y_pred.append(output) # Save Prediction

  labels = labels.data.cpu().numpy()
  
  y_true.extend(labels) # Save Truth

  count=count+1

# constant for classes
classes = ('Ants','Bees')

# Build confusion matrix
cf_matrix = confusion_matrix(y_true, y_pred)

df_cm = pd.DataFrame(cf_matrix, index = [i for i in classes],
                     columns = [i for i in classes])
plt.figure(figsize = (12,7))
sn.heatmap(df_cm, annot=True)
plt.savefig('output.png')

Open a new file and load model

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('.../ants_bees_model_epoch_10_VGG16_scripted.pt')
model.eval()

Define function for prediction

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

Load Image and predict

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

More Predictions