# Mask R-CNN model for lesion segmentation in chest CT scans
# Torchvision detection package is locally re-implemented
# by Alex Ter-Sarkisov@City, University of London
# alex.ter-sarkisov@city.ac.uk
# 2020
import argparse
import time
import pickle
import torch
import torchvision
import numpy as np
import os, sys
import cv2
import models.mask_net as mask_net
from models.mask_net.rpn_segmentation import AnchorGenerator
from models.mask_net.faster_rcnn import FastRCNNPredictor, TwoMLPHead
from models.mask_net.covid_mask_net import MaskRCNNHeads, MaskRCNNPredictor
from torch.utils import data
import torch.utils as utils
import datasets.dataset_segmentation as dataset
from PIL import Image as PILImage
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.patches import Rectangle
import utils
import config_segmentation as config
# main method
def main(config, main_step):
devices = ['cpu', 'cuda']
mask_classes = ['both', 'ggo', 'merge']
truncation_levels = ['0','1','2']
backbones = ['resnet50', 'resnet34', 'resnet18']
assert config.backbone_name in backbones
assert config.mask_type in mask_classes
assert config.truncation in truncation_levels
# import arguments from the config file
start_epoch, model_name, use_pretrained_resnet_backbone, num_epochs, save_dir, train_data_dir, val_data_dir, imgs_dir, gt_dir, batch_size, device, save_every, lrate, rpn_nms, mask_type, backbone_name, truncation = \
config.start_epoch, config.model_name, config.use_pretrained_resnet_backbone, config.num_epochs, config.save_dir, \
config.train_data_dir, config.val_data_dir, config.imgs_dir, config.gt_dir, config.batch_size, config.device, config.save_every, config.lrate, config.rpn_nms_th, config.mask_type, config.backbone_name, config.truncation
assert device in devices
if not save_dir in os.listdir('.'):
os.mkdir(save_dir)
if batch_size > 1:
print("The model was implemented for batch size of one")
if device == 'cuda' and torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
print(device)
# Load the weights if provided
if config.pretrained_model is not None:
pretrained_model = torch.load(config.pretrained_model, map_location = device)
use_pretrained_resnet_backbone = False
else:
pretrained_model=None
torch.manual_seed(time.time())
##############################################################################################
# DATASETS + DATALOADERS
# Alex: could be added in the config file in the future
# parameters for the dataset
dataset_covid_pars_train = {'stage': 'train', 'gt': os.path.join(train_data_dir, gt_dir),
'data': os.path.join(train_data_dir, imgs_dir), 'mask_type':mask_type, 'ignore_small':True}
datapoint_covid_train = dataset.CovidCTData(**dataset_covid_pars_train)
dataset_covid_pars_eval = {'stage': 'eval', 'gt': os.path.join(val_data_dir, gt_dir),
'data': os.path.join(val_data_dir, imgs_dir), 'mask_type':mask_type, 'ignore_small':True}
datapoint_covid_eval = dataset.CovidCTData(**dataset_covid_pars_eval)
###############################################################################################
dataloader_covid_pars_train = {'shuffle': True, 'batch_size': batch_size}
dataloader_covid_train = data.DataLoader(datapoint_covid_train, **dataloader_covid_pars_train)
#
dataloader_covid_pars_eval = {'shuffle': False, 'batch_size': batch_size}
dataloader_covid_eval = data.DataLoader(datapoint_covid_eval, **dataloader_covid_pars_eval)
###############################################################################################
# MASK R-CNN model
# Alex: these settings could also be added to the config
if mask_type == "both":
n_c = 3
else:
n_c = 2
maskrcnn_args = {'min_size': 512, 'max_size': 1024, 'rpn_batch_size_per_image': 256, 'rpn_positive_fraction': 0.75,
'box_positive_fraction': 0.75, 'box_fg_iou_thresh': 0.75, 'box_bg_iou_thresh': 0.5,
'num_classes': None, 'box_batch_size_per_image': 256, 'rpn_nms_thresh': rpn_nms}
# Alex: for Ground glass opacity and consolidatin segmentation
# many small anchors
# use all outputs of FPN
# IMPORTANT!! For the pretrained weights, this determines the size of the anchor layer in RPN!!!!
# pretrained model must have anchors
if pretrained_model is None:
anchor_generator = AnchorGenerator(
sizes=tuple([(2, 4, 8, 16, 32) for r in range(5)]),
aspect_ratios=tuple([(0.1, 0.25, 0.5, 1, 1.5, 2) for rh in range(5)]))
else:
print("Loading the anchor generator")
sizes = pretrained_model['anchor_generator'].sizes
aspect_ratios = pretrained_model['anchor_generator'].aspect_ratios
anchor_generator = AnchorGenerator(sizes=sizes, aspect_ratios=aspect_ratios)
print(anchor_generator, anchor_generator.num_anchors_per_location())
# num_classes:3 (1+2)
# in_channels
# 256: number if channels from FPN
# For the ResNet50+FPN: keep the torchvision architecture, but with 128 features
# For lightweights models: re-implement MaskRCNNHeads with a single layer
box_head = TwoMLPHead(in_channels=256*7*7,representation_size=128)
if backbone_name == 'resnet50':
maskrcnn_heads = None
box_predictor = FastRCNNPredictor(in_channels=128, num_classes=n_c)
mask_predictor = MaskRCNNPredictor(in_channels=256, dim_reduced=256, num_classes=n_c)
else:
#Backbone->FPN->boxhead->boxpredictor
box_predictor = FastRCNNPredictor(in_channels=128, num_classes=n_c)
maskrcnn_heads = MaskRCNNHeads(in_channels=256, layers=(128,), dilation=1)
mask_predictor = MaskRCNNPredictor(in_channels=128, dim_reduced=128, num_classes=n_c)
maskrcnn_args['box_head'] = box_head
maskrcnn_args['rpn_anchor_generator'] = anchor_generator
maskrcnn_args['mask_head'] = maskrcnn_heads
maskrcnn_args['mask_predictor'] = mask_predictor
maskrcnn_args['box_predictor'] = box_predictor
# Instantiate the segmentation model
maskrcnn_model = mask_net.maskrcnn_resnet_fpn(backbone_name, truncation, pretrained_backbone=use_pretrained_resnet_backbone, **maskrcnn_args)
# pretrained?
print(maskrcnn_model.backbone.out_channels)
if pretrained_model is not None:
print("Loading pretrained weights")
maskrcnn_model.load_state_dict(pretrained_model['model_weights'])
if pretrained_model['epoch']:
start_epoch = int(pretrained_model['epoch'])+1
if 'model_name' in pretrained_model.keys():
model_name = str(pretrained_model['model_name'])
# Set to training mode
print(maskrcnn_model)
maskrcnn_model.train().to(device)
optimizer_pars = {'lr': lrate, 'weight_decay': 1e-3}
optimizer = torch.optim.Adam(list(maskrcnn_model.parameters()), **optimizer_pars)
if pretrained_model is not None and 'optimizer_state' in pretrained_model.keys():
optimizer.load_state_dict(pretrained_model['optimizer_state'])
start_time = time.time()
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, optimizer, device, maskrcnn_model, save_every,
lrate, model_name, None, None)
eval_loss_epoch = main_step("eval", e, dataloader_covid_eval, optimizer, device, maskrcnn_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 b in dataloader:
optimizer.zero_grad()
X, y = b
if device == torch.device('cuda'):
X, y['labels'], y['boxes'], y['masks'] = X.to(device), y['labels'].to(device), y['boxes'].to(device), y[
'masks'].to(device)
images = [im for im in X]
targets = []
lab = {}
# THIS IS IMPORTANT!!!!!
# get rid of the first dimension (batch)
# IF you have >1 images, make another loop
# REPEAT: DO NOT USE BATCH DIMENSION
lab['boxes'] = y['boxes'].squeeze_(0)
lab['labels'] = y['labels'].squeeze_(0)
lab['masks'] = y['masks'].squeeze_(0)
if len(lab['boxes']) > 0 and len(lab['labels']) > 0 and len(lab['masks']) > 0:
targets.append(lab)
else:
pass
# avoid empty objects
if len(targets) > 0:
loss = model(images, targets)
total_loss = 0
for k in loss.keys():
total_loss += loss[k]
if stage == "train":
total_loss.backward()
optimizer.step()
else:
pass
epoch_loss += total_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_name':model_name, 'model_weights': model.state_dict(),
'optimizer_state': optimizer.state_dict(), 'lrate': lrate, 'anchor_generator':anchors}
if model_name is None:
print(save_dir, "mrcnn_covid_segmentation_model_ckpt_" + str(e+1) + ".pth")
torch.save(state, os.path.join(save_dir, "mrcnn_covid_segmentation_model_ckpt_" + str(e+1) + ".pth"))
else:
torch.save(state, os.path.join(save_dir, model_name + "_ckpt_" + str(e+1) + ".pth"))
model.train()
return epoch_loss
# run the training of the segmentation algoithm
if __name__ == '__main__':
config_train = config.get_config_pars("trainval")
main(config_train, step)
Datasets
Models
