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