# COVID-CT-Mask-Net

# I re-implemented Torchvision's detection library (Faster and Mask R-CNN) as a classifier

# Alex Ter-Sarkisov @ City, University of London

# alex.ter-sarkisov@city.ac.uk

#

import os

import pickle

import sys

import sys

import time

import config_classifier

import cv2

import datasets.dataset_classifier as dataset

# IMPORT LOCAL IMPLEMENTATION OF TORCHVISION'S DETECTION LIBRARY

import numpy as np

import torch

import torch.nn.functional as F

import torchvision

import utils

from PIL import Image as PILImage

# IMPORT LOCAL IMPLEMENTATION OF TORCHVISION'S DETECTION LIBRARY

# Faster R-CNN interface

import models.mask_net as mask_net

from models.mask_net.faster_rcnn import FastRCNNPredictor, TwoMLPHead

from models.mask_net.rpn import AnchorGenerator

from torch.utils import data

from torchvision import transforms

# main method

def main(config, main_step):

    torch.manual_seed(time.time())

    start_time = time.time()

    devices = ['cpu', 'cuda']

    backbones = ['resnet50', 'resnet34', 'resnet18']

    truncation_levels = ['0','1','2']

    assert config.device in devices

    assert config.backbone_name in backbones

    assert config.truncation in truncation_levels

    start_epoch, pretrained_classifier, pretrained_segment, model_name, num_epochs, save_dir, train_data_dir, val_data_dir, \

    batch_size, device, save_every, lrate, rpn_nms, roi_nms, backbone_name, truncation, roi_batch_size, n_c, s_features = \

                                            config.start_epoch, config.pretrained_classification_model, \

                                            config.pretrained_segmentation_model, \

                                            config.model_name, config.num_epochs, config.save_dir, \

                                            config.train_data_dir, config.val_data_dir, \

                                            config.batch_size, config.device, config.save_every, \

                                            config.lrate, config.rpn_nms_th, config.roi_nms_th, \

                                            config.backbone_name, config.truncation, \

                                            config.roi_batch_size, config.num_classes, config.s_features

    if pretrained_classifier is not None and pretrained_segment is not None:

        print("Not clear which model to use, switching to the classifier")

        pretrained_model = pretrained_classifier

    elif pretrained_classifier is not None and pretrained_segment is None:

        pretrained_model = pretrained_classifier

    else:

        pretrained_model = pretrained_segment

    if device == 'cuda' and torch.cuda.is_available():

        device = torch.device('cuda')

    else:

        device = torch.device('cpu')

    ##############################################################################################

    # DATASETS+DATALOADERS

    # Alex: could be added in the config file in the future

    # parameters for the dataset

    # 512x512 is the recommended image size input

    dataset_covid_pars_train_cl = {'stage': 'train', 'data': train_data_dir, 'img_size': (512,512)}

    datapoint_covid_train_cl = dataset.COVID_CT_DATA(**dataset_covid_pars_train_cl)

    #

    dataset_covid_pars_eval_cl = {'stage': 'eval', 'data': val_data_dir, 'img_size': (512,512)}

    datapoint_covid_eval_cl = dataset.COVID_CT_DATA(**dataset_covid_pars_eval_cl)

    #

    dataloader_covid_pars_train_cl = {'shuffle': True, 'batch_size': batch_size}

    dataloader_covid_train_cl = data.DataLoader(datapoint_covid_train_cl, **dataloader_covid_pars_train_cl)

    #

    dataloader_covid_pars_eval_cl = {'shuffle': True, 'batch_size': batch_size}

    dataloader_covid_eval_cl = data.DataLoader(datapoint_covid_eval_cl, **dataloader_covid_pars_eval_cl)

    #

    ##### LOAD PRETRAINED WEIGHTS FROM MASK R-CNN MODEL

    # This must be the full path to the checkpoint with the anchor generator and model weights

    # Assumed that the keys in the checkpoint are model_weights and anchor_generator

    ckpt = torch.load(pretrained_model, map_location=device)

    # keyword arguments

    # box_score_threshold:negative!

    # set both NMS thresholds to 0.75 to get adjacent RoIs

    # Box detections/image: batch size for the classifier

    #

    covid_mask_net_args = {'num_classes': None, 'min_size': 512, 'max_size': 1024, 'box_detections_per_img': roi_batch_size,

                           'box_nms_thresh': roi_nms, 'box_score_thresh': -0.01, 'rpn_nms_thresh': rpn_nms}

    # copy the anchor generator parameters, create a new one to avoid implementations' clash

    sizes = ckpt['anchor_generator'].sizes

    aspect_ratios = ckpt['anchor_generator'].aspect_ratios

    anchor_generator = AnchorGenerator(sizes, aspect_ratios)

    # out_channels:256, FPN

    # num_classes:3 (1+2)

    box_head = TwoMLPHead(in_channels=256*7*7, representation_size=128)

    box_predictor = FastRCNNPredictor(in_channels=128, num_classes=n_c)

    covid_mask_net_args['rpn_anchor_generator'] = anchor_generator

    covid_mask_net_args['box_predictor'] = box_predictor

    covid_mask_net_args['box_head'] = box_head

    covid_mask_net_args['s_representation_size'] = s_features

    # Instantiate the model

    covid_mask_net_model = mask_net.fasterrcnn_resnet_fpn(backbone_name, truncation, **covid_mask_net_args)

    # which parameters to train?

    trained_pars = []

    # if the weights are loaded from the segmentation model:

    if pretrained_classifier is None:

        for _n, _par in covid_mask_net_model.state_dict().items():

            if _n in ckpt['model_weights']:

                print('Loading parameter', _n)

                _par.copy_(ckpt['model_weights'][_n])

    # if the weights are loaded from the classification model

    else:

        covid_mask_net_model.load_state_dict(ckpt['model_weights'])

        if 'epoch' in ckpt.keys():

            start_epoch = int(ckpt['epoch']) + 1

        if 'model_name' in ckpt.keys():

            model_name = str(ckpt['model_name'])

    # Evaluation mode, no labels!

    covid_mask_net_model.eval()

    # set the model to training mode without triggering the 'training' mode of Mask R-CNN

    # set up the optimizer

    utils.switch_model_on(covid_mask_net_model, ckpt, trained_pars)

    utils.set_to_train_mode(covid_mask_net_model)

    print(covid_mask_net_model)

    covid_mask_net_model = covid_mask_net_model.to(device)

    total_trained_pars = sum([x.numel() for x in trained_pars])

    print("Total trained pars {0:d}".format(total_trained_pars))

    optimizer_pars = {'lr': lrate, 'weight_decay': 1e-3}

    optimizer = torch.optim.Adam(trained_pars, **optimizer_pars)

    if pretrained_classifier is not None and 'optimizer_state' in ckpt.keys():

        optimizer.load_state_dict(ckpt['optimizer_state'])

    if start_epoch>0:

       num_epochs += start_epoch

    print("Start training, epoch = {:d}".format(start_epoch))

    for e in range(start_epoch, num_epochs):

        train_loss_epoch = main_step("train", e, dataloader_covid_train_cl, optimizer, device, covid_mask_net_model,

                                     save_every, lrate, model_name, None, None)

        eval_loss_epoch = main_step("eval", e, dataloader_covid_eval_cl, optimizer, device, covid_mask_net_model,

                                    save_every, lrate, model_name, anchor_generator, save_dir)

        print(

            "Epoch {0:d}: train loss = {1:.3f}, validation loss = {2:.3f}".format(e, train_loss_epoch, eval_loss_epoch))

    end_time = time.time()

    print("Training took {0:.1f} seconds".format(end_time - start_time))

def step(stage, e, dataloader, optimizer, device, model, save_every, lrate, model_name, anchors, save_dir):

    epoch_loss = 0

    for id, b in enumerate(dataloader):

        optimizer.zero_grad()

        X, y = b

        if device == torch.device('cuda'):

            X, y = X.to(device), y.to(device)

        # some batches are less than batch_size

        batch_s = X.size()[0]

        batch_scores = []

        # input all images in the batch into COVID-Mask-Net to get B scores

        for id in range(batch_s):

            image = [X[id]]  # remove the batch dimension

            predict_scores = model(image)

            batch_scores.append(predict_scores[0]['final_scores'])

        # batchify scores/image and compute binary cross-entropy loss

        batch_scores = torch.stack(batch_scores)

        batch_loss = F.binary_cross_entropy_with_logits(batch_scores, y)

        if stage == "train":

            batch_loss.backward()

            optimizer.step()

        else:

            pass

        epoch_loss += batch_loss.clone().detach().cpu().numpy()

    epoch_loss = epoch_loss / len(dataloader)

    if not (e+1) % save_every and stage == "eval":

        model.eval()

        state = {'epoch': str(e+1), 'model_weights': model.state_dict(),

                 'optimizer_state': optimizer.state_dict(), 'lrate': lrate, 'anchor_generator': anchors,

                 'model_name': model_name}

        if model_name is None:

            torch.save(state, os.path.join(save_dir, "covid_ct_mask_net_ckpt_" + str(e+1) + ".pth"))

        else:

            torch.save(state, os.path.join(save_dir, model_name + "_ckpt_" + str(e+1) + ".pth"))

        utils.set_to_train_mode(model)

    return epoch_loss

# run the training

if __name__ == '__main__':

    config_train = config_classifier.get_config_pars_classifier("trainval")

    if config_train.pretrained_classification_model is None and config_train.pretrained_segmentation_model is None:

        print("You must have at least one pretrained model!")

        sys.exit(0)

    else:

        main(config_train, step)