bt_classifier.py

# -*- coding: utf-8 -*-
"""bt_classifier.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1XrA2NhcdyB5_7sJ9Mt0EOwQE-BqJsPAQ

Brain Tumor Classification
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms, models
from torchvision.utils import make_grid
import os
import random
import numpy as np
import pandas as pd
import pickle
import time
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.metrics import confusion_matrix, classification_report, jaccard_similarity_score
from google.colab import drive

drive.mount('/content/drive')

"""Empty GPU's memory/cache for training"""

torch.cuda.empty_cache()

class BrainTumorDataset(Dataset):
  def __init__(self, images, labels):
    self.X = images
    self.y = labels
    # Transformation for converting original image array to an image and then convert it to a tensor
    self.transform = transforms.Compose([transforms.ToPILImage(),
        transforms.ToTensor()])
    # Transformation for converting original image array to an image, rotate it randomly between -90 degrees and 90 degrees, and then convert it to a tensor
    self.transform2 = transforms.Compose([
        transforms.ToPILImage(),
        transforms.RandomRotation(90),
        transforms.ToTensor()])

    self.transform4 = transforms.Compose([
        transforms.ToPILImage(),
        transforms.RandomRotation(180),
        transforms.ToTensor()])

    self.transform6 = transforms.Compose([
        transforms.ToPILImage(),
        transforms.RandomRotation(300),
        transforms.ToTensor()])

  def __len__(self):
    return len(self.X)

  def __getitem__(self, idx):
    data = self.transform(self.X[idx])
    # Augmented image at 90 degrees as a tensor
    aug90 = self.transform2(self.X[idx])
    aug180 = self.transform4(self.X[idx])
    aug300 = self.transform6(self.X[idx])
    new_batch = [data, aug90, aug180, aug300]
    labels = torch.zeros(4, dtype=torch.float32)
    labels[int(self.y[idx])] = 1.0
    new_labels = [ labels, labels, labels, labels]

    return (torch.stack(new_labels), torch.stack(new_batch))

"""## Load the Dataset"""

training_data = pickle.load(open('/content/drive/MyDrive/brain tumor/new_dataset/lbp_data1.pickle', 'rb'))

from google.colab.patches import cv2_imshow
cv2_imshow(training_data[900][0])
print(training_data[900][1])

"""Create empty lists for storing our data"""

Xt = []
yt = []
features = None
labels = None
label = []

for features,labels in training_data:
  Xt.append(features)
  yt.append(labels)

# 70 % training, 15% validating, 15% testing
X_train, X_test, y_train, y_test = train_test_split(Xt, yt, test_size=0.3, shuffle=True)  # 70% training, 30% testing
X_valid, X_test, y_valid, y_test = train_test_split(X_test, y_test, test_size=0.5, shuffle=True)  # split testing set into 50% validation , 50% testing

"""Empty the previously used lists and arrays to free up RAM / Cache"""

Xt = None
yt = None
features = None
labels = None
label = None
training_data = None

train_set = BrainTumorDataset(X_train, y_train)
valid_set = BrainTumorDataset(X_valid, y_valid)
test_set = BrainTumorDataset(X_test, y_test)

"""Print original number of samples in each set"""

print(f"Number of training samples: {len(X_train)}")
print(f"Number of validation samples: {len(X_valid)}")
print(f"Number of testing samples: {len(X_test)}")

"""Print augmented number of samples in each set"""

print(f"Number of augmented training samples: {len(X_train) * 4}")
print(f"Number of augmented validation samples: {len(X_valid)* 4}")
print(f"Number of augmented testing samples: {len(X_test)* 4}")

"""Create a DataLoader for each set with batch """

train_gen = DataLoader(train_set, batch_size=1, shuffle=True, pin_memory=True, num_workers=8)
valid_gen = DataLoader(valid_set, batch_size=1, shuffle=True, pin_memory=True, num_workers=8)
test_gen = DataLoader(test_set, batch_size=4, shuffle=True, pin_memory=True, num_workers=8)

device_name = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device_name)

""" Building the Model"""

# instantiate transfer learning model
resnet_model = models.resnet50(pretrained=True)
for param in resnet_model.parameters():
    param.requires_grad = True
n_inputs = resnet_model.fc.in_features
resnet_model.fc = nn.Sequential(nn.Linear(n_inputs, 2048),
                                nn.SELU(),
                                nn.Dropout(p=0.4),
                                nn.Linear(2048, 2048),
                                nn.SELU(),
                                nn.Dropout(p=0.4),
                                nn.Linear(2048, 4),
                                nn.LogSigmoid())

for name, child in resnet_model.named_children():
  for name2, params in child.named_parameters():
    params.requires_grad = True
resnet_model.to(device)
resnet_model

"""## Training Configuration

"""

criterion = nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.SGD(resnet_model.parameters(), momentum=0.9, lr=3e-4)
epochs = 30

train_losses = []
test_losses = []
train_correct = []
test_correct = []

def save_checkpoint(state, is_best, filename='/content/drive/MyDrive/bt_resnet50_lbp.pth.tar'):
    torch.save(state, filename)

"""## Train the model"""

start_time = time.time()
best_prec1 = 2
b = None
train_b = None
test_b = None

# start training
for i in range(epochs):
    trn_corr = 0
    tst_corr = 0
    e_start = time.time()
    
    # train in batches
    for b, (y, X) in enumerate(train_gen):
        X, y = X.to(device), y.to(device)
        y_pred = resnet_model(X.view(-1, 3, 512, 512))
        # calculate loss
        loss = criterion(y_pred.float(), torch.argmax(y.view(-1, 4), dim=1).long())
        predicted = torch.argmax(y_pred, dim=1).data
        batch_corr = (predicted == torch.argmax(y.view(-1, 4), dim=1)).sum()
        trn_corr += batch_corr
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    e_end = time.time()
    print(f'Epoch {(i+1)} Batch {(b+1)*4}\nAccuracy: {trn_corr.item()*100/(1*4*b):2.2f} %  Loss: {loss.item():2.4f}  Duration: {((e_end-e_start)/60):.2f} minutes') # 4 images per batch * 8 augmentations per image * batch length
    train_b = b
    train_losses.append(loss)
    train_correct.append(trn_corr)

    X, y = None, None
    # validate using validation generator
    with torch.no_grad():
        for b, (y, X) in enumerate(valid_gen):
            X, y = X.to(device), y.to(device)
            y_val = resnet_model(X.view(-1, 3, 512, 512))
            predicted = torch.argmax(y_val, dim=1).data
            tst_corr += (predicted == torch.argmax(y.view(-1, 4), dim=1)).sum()

    # get loss of validation set
    loss = criterion(y_val.float(), torch.argmax(y.view(-1, 4), dim=1).long())
    print(f'Validation Accuracy {tst_corr.item()*100/(1*4*b):2.2f} Validation Loss: {loss.item():2.4f}\n')
    is_best = loss < best_prec1
    best_prec1 = min(loss, best_prec1)
    save_checkpoint({
            'epoch': i + 1,
            'state_dict': resnet_model.state_dict(),
            'best_prec1': best_prec1,
        }, is_best)
    test_b  = b
    test_losses.append(loss)
    test_correct.append(tst_corr)
end_time = time.time() - start_time    

# print training summary
print("\nTraining Duration {:.2f} minutes".format(end_time/60))
print("GPU memory used : {} kb".format(torch.cuda.memory_allocated()))
print("GPU memory cached : {} kb".format(torch.cuda.memory_cached()))

torch.cuda.empty_cache()

"""## Save the model

Save the model after the training is completed
"""

torch.save(resnet_model.state_dict(), '/content/drive/My Drive/bt_resnet50_model.pt')

"""## Evaluation"""

print(f'Validation accuracy: {test_correct[-1].item()*100/(test_b*4*1):.2f}%')

plt.plot(train_losses, label='Training loss')
plt.plot(test_losses, label='Validation loss')
plt.title('Loss Metrics')
plt.ylabel('Loss')
plt.xlabel('Epochs')
plt.legend()
plt.show()

plt.plot([t/85 for t in train_correct], label='Training accuracy')
plt.plot([t/18 for t in test_correct], label='Validation accuracy')
plt.title('Accuracy Metrics')
plt.ylabel('Accuracy')
plt.xlabel('Epochs')
plt.legend()
plt.show()

# resnet_model.load_state_dict(torch.load('/content/drive/My Drive/bt_resnet_torch.pt'))
train_gen = None
valid_gen = None
train_set = None
valid_set = None

resnet_model.eval()
with torch.no_grad():
    correct = 0
    test_loss = []
    test_corr = []
    labels = []
    pred = []
    # perform test set evaluation batch wise
    for (y, X) in test_gen:
        X, y = X.to(device), y.to(device)
        labels.append(torch.argmax(y.view(-1, 4), dim=1).data)
        y_val = resnet_model(X.view(-1, 1, 512, 512))
        predicted = torch.argmax(y_val, dim=1).data
        pred.append(predicted)

        # calculate loss
        loss = criterion(y_val.float(), torch.argmax(y.view(-1, 4), dim=1).long())
        correct += (predicted == torch.argmax(y.view(-1, 4), dim=1)).sum()

        # append correct samples labels and losses
        test_corr.append(correct)
        test_loss.append(loss)
        
print(f"Test Loss: {test_loss[-1].item():.4f}")

"""Print the test accuracy

"""

print(f'Test accuracy: {test_corr[-1].item()*100/(460*4):.2f}%')

labels = torch.stack(labels)
pred = torch.stack(pred)

LABELS = ['Meningioma', 'Glioma', 'Pitutary']

arr = confusion_matrix(pred.view(-1).cpu(), labels.view(-1).cpu())
df_cm = pd.DataFrame(arr, LABELS, LABELS)
plt.figure(figsize = (9,6))
sns.heatmap(df_cm, annot=True, fmt="d", cmap='viridis')
plt.xlabel("Prediction")
plt.ylabel("Target")
plt.show()