--- a
+++ b/inference_segmentation.py
@@ -0,0 +1,184 @@
+
+# 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)