# 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
# 2020
import argparse
import os
from collections import OrderedDict
import config_segmentation as config
import torch
import torch.utils.data as data
import torchvision
# implementation of the mAP
import utils
from datasets import dataset_segmentation as dataset
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
def main(config, step):
devices = ['cpu', 'cuda']
mask_classes = ['both', 'ggo', 'merge']
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
if config.device == 'cuda' and torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
#
model_name = None
ckpt = torch.load(config.ckpt, map_location=device)
if 'model_name' in ckpt.keys():
model_name = ckpt['model_name']
device = torch.device('cpu')
if torch.cuda.is_available():
device = torch.device('cuda')
# get the thresholds
confidence_threshold, mask_threshold, save_dir, data_dir, img_dir, gt_dir, mask_type, rpn_nms, roi_nms, backbone_name, truncation \
= config.confidence_th, config.mask_logits_th, config.save_dir, config.test_data_dir, config.test_imgs_dir, \
config.gt_dir, config.mask_type, config.rpn_nms_th, config.roi_nms_th, config.backbone_name, config.truncation
if model_name is None:
model_name = "maskrcnn_segmentation"
elif model_name is not None and config.model_name != model_name:
print("Using model name from the config.")
model_name = config.model_name
# either 2+1 or 1+1 classes
assert mask_type in mask_classes
if mask_type == "both":
n_c = 3
else:
n_c = 2
# dataset interface
dataset_covid_eval_pars = {'stage': 'eval', 'gt': os.path.join(data_dir, gt_dir),
'data': os.path.join(data_dir, img_dir), 'mask_type': mask_type, 'ignore_small':True}
datapoint_eval_covid = dataset.CovidCTData(**dataset_covid_eval_pars)
dataloader_covid_eval_pars = {'shuffle': False, 'batch_size': 1}
dataloader_eval_covid = data.DataLoader(datapoint_eval_covid, **dataloader_covid_eval_pars)
# MASK R-CNN model
# Alex: these settings could also be added to the config
ckpt = torch.load(config.ckpt, map_location=device)
sizes = ckpt['anchor_generator'].sizes
aspect_ratios = ckpt['anchor_generator'].aspect_ratios
anchor_generator = AnchorGenerator(sizes, aspect_ratios)
print("Anchors: ", anchor_generator.sizes, anchor_generator.aspect_ratios)
# create modules
# this assumes FPN with 256 channels
box_head = TwoMLPHead(in_channels=7 * 7 * 256, 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)
# keyword arguments
maskrcnn_args = {'num_classes': None, 'min_size': 512, 'max_size': 1024, 'box_detections_per_img': 128,
'box_nms_thresh': roi_nms, 'box_score_thresh': confidence_threshold, 'rpn_nms_thresh': rpn_nms,
'box_head': box_head, 'rpn_anchor_generator': anchor_generator, 'mask_head':maskrcnn_heads,
'mask_predictor': mask_predictor, 'box_predictor': box_predictor}
# Instantiate the segmentation model
maskrcnn_model = mask_net.maskrcnn_resnet_fpn(backbone_name, truncation, pretrained_backbone=False, **maskrcnn_args)
# Load weights
maskrcnn_model.load_state_dict(ckpt['model_weights'])
# Set to the evaluation mode
print(maskrcnn_model)
maskrcnn_model.eval().to(device)
# IoU thresholds. By default the model computes AP for each threshold between 0.5 and 0.95 with the step of 0.05
thresholds = torch.arange(0.5, 1, 0.05).to(device)
mean_aps_all_th = torch.zeros(thresholds.size()[0]).to(device)
ap_th = OrderedDict()
# run the loop for all thresholds
for t, th in enumerate(thresholds):
# main method
ap = step(maskrcnn_model, th, dataloader_eval_covid, device, mask_threshold)
mean_aps_all_th[t] = ap
th_name = 'AP@{0:.2f}'.format(th)
ap_th[th_name] = ap
print("Done evaluation for {}".format(model_name))
print("mAP:{0:.4f}".format(mean_aps_all_th.mean().item()))
for k, aps in ap_th.items():
print("{0:}:{1:.4f}".format(k, aps))
def compute_map(model, iou_th, dl, device, mask_th):
mean_aps_this_th = torch.zeros(len(dl), dtype=torch.float)
for v, b in enumerate(dl):
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)
lab = {'boxes': y['boxes'].squeeze_(0), 'labels': y['labels'].squeeze_(0), 'masks': y['masks'].squeeze_(0)}
image = [x.squeeze_(0)] # remove the batch dimension
out = model(image)
# scores + bounding boxes + labels + masks
scores = out[0]['scores']
bboxes = out[0]['boxes']
classes = out[0]['labels']
# remove the empty dimension,
# output_size x 512 x 512
predict_mask = out[0]['masks'].squeeze_(1) > mask_th
if len(scores) > 0 and len(lab['labels']) > 0:
# threshold for the masks:
ap, _, _, _ = utils.compute_ap(lab['boxes'], lab['labels'], lab['masks'], bboxes, classes, scores,
predict_mask, iou_threshold=iou_th)
mean_aps_this_th[v] = ap
elif not len(scores) and not len(lab['labels']):
mean_aps_this_th[v] = 1
elif not len(scores) and len(lab['labels']) > 0:
continue
elif len(scores) > 0 and not len(y['labels']):
continue
return mean_aps_this_th.mean().item()
if __name__ == "__main__":
config_mean_ap = config.get_config_pars("precision")
main(config_mean_ap, compute_map)
Datasets
Models
