# COVID-CT-Mask-Net

# Torchvision detection package is locally re-implemented

# Transformed into a classification model with Mask R-CNN backend

# by Alex Ter-Sarkisov@City, University of London

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

# 2020

import os

import re

import sys

import time

import config_classifier as config

import cv2

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

import models.mask_net as mask_net

import numpy as np

import torch

import torchvision

import torchvision.transforms as transforms

import utils

from PIL import Image as PILImage

from models.mask_net.faster_rcnn import FastRCNNPredictor, TwoMLPHead

from models.mask_net.rpn import AnchorGenerator

def main(config, 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

    pretrained_model, model_name, test_data_dir, device, rpn_nms, roi_nms, backbone_name, truncation, roi_batch_size, n_c, s_features\

              = config.ckpt, config.model_name, config.test_data_dir, config.device, config.rpn_nms_th, \

                config.roi_nms_th, config.backbone_name, config.truncation, config.roi_batch_size, config.num_classes, config.s_features

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

        device = torch.device('cuda')

    else:

        device = torch.device('cpu')

    # either 2+1 or 1+1 classes

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

    # 'box_detections_per_img': batch size input in module S

    # 'box_score_thresh': negative to accept all predictions

    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}

    print(covid_mask_net_args)

    # extract anchor generator from the checkpoint

    sizes = ckpt['anchor_generator'].sizes

    aspect_ratios = ckpt['anchor_generator'].aspect_ratios

    anchor_generator = AnchorGenerator(sizes, aspect_ratios)

    # Faster R-CNN interfaces, masks not implemented at this stage

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

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

    # Mask prediction is not necessary, keep it for future extensions

    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

    # representation size of the S classification module

    # these should be provided in the config

    covid_mask_net_args['s_representation_size'] = s_features

    # Instance of the model, copy weights

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

    covid_mask_net_model.load_state_dict(ckpt['model_weights'])

    covid_mask_net_model.eval().to(device)

    print(covid_mask_net_model)

    # confusion matrix

    confusion_matrix = torch.zeros(3, 3, dtype=torch.int32).to(device)

    for idx, f in enumerate(os.listdir(test_data_dir)):

        step(f, covid_mask_net_model, test_data_dir, device, confusion_matrix)

    print("------------------------------------------")

    print("Confusion Matrix for 3-class problem:")

    print("0: Control, 1: Normal Pneumonia, 2: COVID")

    print(confusion_matrix)

    print("------------------------------------------")

    # confusion matrix

    cm = confusion_matrix.float()

    cm[0, :].div_(cm[0, :].sum())

    cm[1, :].div_(cm[1, :].sum())

    cm[2, :].div_(cm[2, :].sum())

    print("------------------------------------------")

    print("Class Sensitivity:")

    print(cm)

    print("------------------------------------------")

    print('Overall accuracy:')

    print(confusion_matrix.diag().float().sum().div(confusion_matrix.sum()))

    end_time = time.time()

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

def test_step(im_input, model, source_dir, device, c_matrix):

    # CNCB NCOV datasets: the first integer is the correct class:

    # 0: control

    # 1: pneumonia

    # 2: COVID

    # extract the correct class from the file name

    correct_class = int(im_input.split('/')[-1].split('_')[0])

    im = PILImage.open(os.path.join(source_dir, im_input))

    if im.mode != 'RGB':

        im = im.convert(mode='RGB')

    # get rid of alpha channel

    img = np.array(im)

    # print(img)

    if img.shape[2] > 3:

        img = img[:, :, :3]

    t_ = transforms.Compose([

        transforms.ToPILImage(),

        transforms.Resize(512),

        transforms.ToTensor()])

    img = t_(img)

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

        img = img.to(device)

    out = model([img])

    pred_class = out[0]['final_scores'].argmax().item()

    # get confusion matrix

    c_matrix[correct_class, pred_class] += 1

# run the inference

if __name__ == '__main__':

    config_test = config.get_config_pars_classifier("test")

    main(config_test, test_step)