# 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 copy
import torch
import torchvision
import numpy as np
import os
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 torchvision import transforms
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, 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
assert config.mask_type in mask_classes
if config.device == 'cuda' and torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
# get the configuration
# 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 mask_type == "both":
n_c = 3
else:
n_c = 2
ckpt = torch.load(config.ckpt, map_location=device)
model_name = None
if 'model_name' in ckpt.keys():
model_name = ckpt['model_name']
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': 100,
'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 evaluation mode
print(maskrcnn_model)
maskrcnn_model.eval().to(device)
start_time = time.time()
# get the correct masks and mask colors
if mask_type == "ggo":
ct_classes = {0: '__bgr', 1: 'GGO'}
ct_colors = {1: 'red', 'mask_cols': np.array([[255, 0, 0]])}
elif mask_type == "merge":
ct_classes = {0: '__bgr', 1: 'Lesion'}
ct_colors = {1: 'red', 'mask_cols': np.array([[255, 0, 0]])}
elif mask_type == "both":
ct_classes = {0: '__bgr', 1: 'GGO', 2: 'CL'}
ct_colors = {1: 'red', 2: 'blue', 'mask_cols': np.array([[255, 0, 0], [0, 0, 255]])}
if not save_dir in os.listdir('.'):
os.mkdir(save_dir)
# model name from config, not checkpoint
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
# run the inference with provided hyperparameters
test_ims = os.listdir(os.path.join(data_dir, img_dir))
for j, ims in enumerate(test_ims):
step(os.path.join(os.path.join(data_dir, img_dir), ims), device, maskrcnn_model, model_name,
confidence_threshold, mask_threshold, save_dir, ct_classes, ct_colors, j)
end_time = time.time()
print("Inference took {0:.1f} seconds".format(end_time - start_time))
def test_step(image, device, model, model_name, theta_conf, theta_mask, save_dir, cls, cols, num):
im = PILImage.open(image)
# convert image to RGB, remove the alpha channel
if im.mode != 'RGB':
im = im.convert(mode='RGB')
img = np.array(im)
# copy image to make background for plotting
bgr_img = copy.deepcopy(img)
if img.shape[2] > 3:
img = img[:, :, :3]
# torchvision transforms, the rest Mask R-CNN does internally
t_ = transforms.Compose([
transforms.ToPILImage(),
transforms.ToTensor()])
img = t_(img).to(device)
out = model([img])
# scores + bounding boxes + labels + masks
scores = out[0]['scores']
bboxes = out[0]['boxes']
classes = out[0]['labels']
mask = out[0]['masks']
# this is the array for all masks
best_scores = scores[scores > theta_conf]
# Are there any detections with confidence above the threshold?
if len(best_scores):
best_idx = np.where(scores > theta_conf)
best_bboxes = bboxes[best_idx]
best_classes = classes[best_idx]
best_masks = mask[best_idx]
print('bm', best_masks.shape)
mask_array = np.zeros([best_masks[0].shape[1], best_masks[0].shape[2], 3], dtype=np.uint8)
fig, ax = plt.subplots(1, 1)
fig.set_size_inches(12, 6)
ax.axis("off")
# plot predictions
for idx, dets in enumerate(best_bboxes):
found_masks = best_masks[idx][0].detach().clone().to(device).numpy()
pred_class = best_classes[idx].item()
pred_col_n = cols[pred_class]
pred_class_txt = cls[pred_class]
pred_col = cols['mask_cols'][pred_class - 1]
mask_array[found_masks > theta_mask] = pred_col
rect = Rectangle((dets[0], dets[1]), dets[2] - dets[0], dets[3] - dets[1], linewidth=1,
edgecolor=pred_col_n, facecolor='none', linestyle="--")
ax.text(dets[0] + 40, dets[1], '{0:}'.format(pred_class_txt), fontsize=10, color=pred_col_n)
ax.text(dets[0], dets[1], '{0:.2f}'.format(best_scores[idx]), fontsize=10, color=pred_col_n)
ax.add_patch(rect)
added_image = cv2.addWeighted(bgr_img, 0.5, mask_array, 0.75, gamma=0)
ax.imshow(added_image)
fig.savefig(os.path.join(save_dir, model_name + "_" + str(num) + ".png"),
bbox_inches='tight', pad_inches=0.0)
else:
print("No detections")
# run the inference
if __name__ == '__main__':
config_test = config.get_config_pars("test")
main(config_test, test_step)
Datasets
Models
