Ceretai/task-4-active_learning/2_Active_learning/Race_Classifier.py

import torch

# print(torch.__version__)
from torch.utils.data import Dataset, DataLoader
import torch
import torch.nn as nn
import pandas as pd
from torchvision import transforms, utils
import cv2
import torch.backends.cudnn as cudnn
cudnn.benchmark = False
from torchsummary import summary
# from google.colab import drive
import os
import numpy as np
import torch.optim as optim
from sklearn.metrics import accuracy_score, confusion_matrix, precision_score, recall_score, ConfusionMatrixDisplay
from time import sleep
from tqdm import tqdm


class ClassifierModel(nn.Module):
    def __init__(self):
        super(ClassifierModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        self.relu1 = nn.ReLU()
        self.batch1 = nn.BatchNorm2d(16)
        self.maxpool1 = nn.MaxPool2d(3)
        self.drop1 = nn.Dropout(p=0.25)

        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.relu2 = nn.ReLU()
        self.batch2 = nn.BatchNorm2d(32)
        self.maxpool2 = nn.MaxPool2d(2)
        self.drop2 = nn.Dropout(p=0.25)

        self.conv3 = nn.Conv2d(32, 32, 3, padding=1)
        self.relu3 = nn.ReLU()
        self.batch3 = nn.BatchNorm2d(32)
        self.maxpool3 = nn.MaxPool2d(2)
        self.drop3 = nn.Dropout(p=0.25)

        self.flatten = nn.Flatten()
        self.dense1 = nn.LazyLinear(128)
        self.relu4 = nn.ReLU()
        self.batch4 = nn.BatchNorm1d(128)
        self.drop4 = nn.Dropout(p=0.5)
        self.dense2 = nn.LazyLinear(7)

        self.soft = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.batch1(x)
        x = self.maxpool1(x)
        x = self.drop1(x)

        x = self.conv2(x)
        x = self.relu2(x)
        x = self.batch2(x)
        x = self.maxpool2(x)
        x = self.drop2(x)

        x = self.conv3(x)
        x = self.relu3(x)
        x = self.batch3(x)
        x = self.maxpool3(x)
        x = self.drop3(x)

        x = self.flatten(x)
        x = self.dense1(x)
        x = self.relu4(x)
        x = self.batch4(x)
        x = self.drop4(x)
        x = self.dense2(x)

        x = self.soft(x)
        return x


class custom_transformer():
    def __init__(self, img_size):
        # Get image you want to resize to at intilization
        self.img_size = img_size

    def __call__(self, sample):
        self.img, self.label = sample["image"], sample["label"]

        # resize image
        self.img = cv2.resize(self.img, self.img_size)

        tensor = transforms.ToTensor()
        norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

        ## Transform image into tensor then normalize values
        self.img = tensor(self.img)
        self.img = norm(self.img)

        ## Transform label into tensor then normalize values
        self.label = torch.FloatTensor([self.label]).long()

        return {"image": self.img, "label": self.label}


class custom_dataset():

    def __init__(self, csv_file, root_dir, transform=None):
        # print("Dataset intialized")
        ## Read csv file
        self.images = pd.read_csv(csv_file)
        ## Get data root directory
        self.root_dir = root_dir
        ## get the transform function
        self.transform = transform
        self.races = {"White": 0, "Black": 1, "East Asian": 2, "Indian": 3, "Latino_Hispanic": 4, "Middle Eastern": 5,
                      "Southeast Asian": 6}

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

    def __getitem__(self, idx):
        # print("Get Item")
        if torch.is_tensor(idx):
            idx = idx.tolist()

        img_name = os.path.join(self.root_dir, self.images.iloc[idx, 0])

        image = cv2.imread(img_name)
        # transform race to id using the self.races dictionary
        labels = self.races[self.images.iloc[idx, 3]]

        sample = {'image': image, 'label': labels}

        if self.transform:
            sample = self.transform(sample)
        # print(sample)
        return sample


def eval_model(model, validation_loader, criterion, device):
    model.eval()
    running_loss = 0
    val_label = []
    val_prediction = []
    print("Validation Time !!")
    with tqdm(validation_loader, unit="batch") as vepoch:
        for sample in vepoch:
            vepoch.set_description("Validation Epoch")
            imgs, labels = sample['image'], sample['label']

            imgs = imgs.to(device)
            labels = labels.to(device)

            # Get Loss for validation
            logits = model(imgs)
            loss = criterion(logits, torch.reshape(labels, (-1,)))
            running_loss += loss.item()

            # Get All Preditions and Labels in the validation set
            val_label = val_label + labels.tolist()
            argmax_logits = torch.argmax(logits, dim=1)
            val_prediction = val_prediction + argmax_logits.tolist()
            vepoch.set_postfix(loss=loss.item())
            sleep(0.1)

    val_label = np.array(val_label)
    val_prediction = np.array(val_prediction)

    val_label = torch.tensor(val_label.flatten())
    val_prediction = torch.tensor(val_prediction.flatten())

    model.train()
    return running_loss / len(validation_loader), accuracy_score(val_label, val_prediction), recall_score(val_label,
                                                                                                          val_prediction,
                                                                                                          average='micro'), precision_score(
        val_label, val_prediction, average='micro'), confusion_matrix(val_label, val_prediction)


def train_model(model, device, train_dataloader, validation_dataloader, n_epoch, save_path):
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.0001)
    torch.backends.cudnn.benchmark = False
    for epoch in range(n_epoch):
        print("Epoch number:", epoch)
        model.train()
        running_loss = 0
        train_label = []
        train_prediction = []

        with tqdm(train_dataloader, unit="batch") as tepoch:

            for sample in tepoch:
                tepoch.set_description(f"Epoch {epoch}")
                ## Get Images and Labels
                imgs = sample['image'].to(device)
                labels = sample['label'].to(device)

                # Prepare Optimzer
                optimizer.zero_grad()

                # Get Prediction in a form of (batch_size * num_classes)
                prediction = model(imgs)

                ## Append all real labels into a list to be able to calculate training accuracy
                train_label = train_label + labels.tolist()

                ## Get indices with highst probability and append it to calculate training accuracy
                argmax_prediction = torch.argmax(prediction, dim=1)
                train_prediction = train_prediction + argmax_prediction.tolist()

                ## Calculate loss

                loss = criterion(prediction, torch.reshape(labels, (-1,)))
                loss.backward()
                ## Optimizer step to update weights
                optimizer.step()

                ## Accumelate training loss
                running_loss += loss.item()
                tepoch.set_postfix(loss=loss.item())
                sleep(0.1)

            if epoch % 2 == 0 or epoch == n_epoch - 1:
                ## Save model
                torch.save(model.state_dict(), os.path.join(save_path, "epoch-" + str(epoch) + ".pth"))
                print('Model Saved !')

                # Calculate training loss
                train_l = running_loss / len(train_dataloader)

                # convert train_label and train_prediction to numpy arrays to be able to flatten them
                train_label = np.array(train_label)
                train_prediction = np.array(train_prediction)

                ## Convert them back to tensors to use sckit learn
                train_label = torch.tensor(train_label.flatten())
                train_prediction = torch.tensor(train_prediction.flatten())

                # Get training accuracy
                train_acc = accuracy_score(train_label, train_prediction)

                ## Evaluation(Validation) loop
                val_l, val_acc, val_recall, val_precision, conf_matrix = eval_model(model, validation_dataloader,
                                                                                    criterion, device)

                # Print Evaluation
                print(
                    'epoch %d train loss: %.3f  |  train acc: %.3f |  val loss: %.3f  |  val acc: %.3f |  val recall: '
                    '%.3f |  val precision: %.3f '
                    % (epoch, train_l, train_acc, val_l, val_acc, val_recall, val_precision))
                print(conf_matrix)

def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    device = "cpu"
    print(device)

    train_csv_file = "E:/My Drive/FairFace/fairface_label_train.csv"
    val_csv_file = "E:/My Drive/FairFace/fairface_label_val.csv"

    root_path = "E:/My Drive/FairFace/fairface-img-margin025-trainval/"

    training_data = custom_dataset(csv_file=train_csv_file, root_dir=root_path, transform=custom_transformer((200, 200)))
    validation_data = custom_dataset(csv_file=val_csv_file, root_dir=root_path, transform=custom_transformer((200, 200)))

    train_dataloader = DataLoader(training_data, batch_size=16, shuffle=True)
    validation_dataloader = DataLoader(validation_data, batch_size=16, shuffle=True)
    model = ClassifierModel().to(device)
    train_model(model, device, train_dataloader, validation_dataloader, 100,"C:/Users/Ahmed Hamdi/PycharmProjects/RaceClassifier")

if __name__ == '__main__':
    main()