diff --git a/exec.py b/exec.py
index 2eb4337..3b56352 100644
--- a/exec.py
+++ b/exec.py
@@ -1,264 +1,273 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""execution script."""
import argparse
import os, warnings
import time
import torch
import utils.exp_utils as utils
from evaluator import Evaluator
from predictor import Predictor
from plotting import plot_batch_prediction
for msg in ["Attempting to set identical bottom==top results",
"This figure includes Axes that are not compatible with tight_layout",
"Data has no positive values, and therefore cannot be log-scaled.",
".*invalid value encountered in double_scalars.*",
".*Mean of empty slice.*"]:
warnings.filterwarnings("ignore", msg)
def train(logger):
"""
perform the training routine for a given fold. saves plots and selected parameters to the experiment dir
specified in the configs.
"""
logger.info('performing training in {}D over fold {} on experiment {} with model {}'.format(
cf.dim, cf.fold, cf.exp_dir, cf.model))
net = model.net(cf, logger).cuda()
- optimizer = torch.optim.Adam(net.parameters(), lr=cf.learning_rate[0], weight_decay=cf.weight_decay)
+ optimizer = torch.optim.AdamW(net.parameters(), lr=cf.learning_rate[0], weight_decay=cf.weight_decay)
if cf.dynamic_lr_scheduling:
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode=cf.scheduling_mode, factor=cf.lr_decay_factor,
patience=cf.scheduling_patience)
model_selector = utils.ModelSelector(cf, logger)
train_evaluator = Evaluator(cf, logger, mode='train')
val_evaluator = Evaluator(cf, logger, mode=cf.val_mode)
starting_epoch = 1
# prepare monitoring
monitor_metrics = utils.prepare_monitoring(cf)
if cf.resume_to_checkpoint:
starting_epoch, monitor_metrics = utils.load_checkpoint(cf.resume_to_checkpoint, net, optimizer)
logger.info('resumed to checkpoint {} at epoch {}'.format(cf.resume_to_checkpoint, starting_epoch))
logger.info('loading dataset and initializing batch generators...')
batch_gen = data_loader.get_train_generators(cf, logger)
for epoch in range(starting_epoch, cf.num_epochs + 1):
logger.info('starting training epoch {}'.format(epoch))
start_time = time.time()
net.train()
train_results_list = []
-
for bix in range(cf.num_train_batches):
batch = next(batch_gen['train'])
tic_fw = time.time()
results_dict = net.train_forward(batch)
tic_bw = time.time()
optimizer.zero_grad()
results_dict['torch_loss'].backward()
optimizer.step()
- logger.info('tr. batch {0}/{1} (ep. {2}) fw {3:.2f}s / bw {4:.2f} s / total {5:.2f} s || '
- .format(bix + 1, cf.num_train_batches, epoch, tic_bw - tic_fw,
- time.time() - tic_bw, time.time() - tic_fw) + results_dict['logger_string'])
- #train_results_list.append([results_dict['boxes'], batch['pid']])
+ print('\rtr. batch {0}/{1} (ep. {2}) fw {3:.2f}s / bw {4:.2f} s / total {5:.2f} s || '.format(
+ bix + 1, cf.num_train_batches, epoch, tic_bw - tic_fw, time.time() - tic_bw,
+ time.time() - tic_fw) + results_dict['logger_string'], flush=True, end="")
train_results_list.append(({k:v for k,v in results_dict.items() if k != "seg_preds"}, batch["pid"]))
+ print()
_, monitor_metrics['train'] = train_evaluator.evaluate_predictions(train_results_list, monitor_metrics['train'])
+ logger.info('generating training example plot.')
+ plot_batch_prediction(batch, results_dict, cf, outfile=os.path.join(
+ cf.plot_dir, 'pred_example_{}_train.png'.format(cf.fold)))
+
train_time = time.time() - start_time
logger.info('starting validation in mode {}.'.format(cf.val_mode))
with torch.no_grad():
net.eval()
if cf.do_validation:
val_results_list = []
val_predictor = Predictor(cf, net, logger, mode='val')
for _ in range(batch_gen['n_val']):
batch = next(batch_gen[cf.val_mode])
if cf.val_mode == 'val_patient':
results_dict = val_predictor.predict_patient(batch)
elif cf.val_mode == 'val_sampling':
results_dict = net.train_forward(batch, is_validation=True)
#val_results_list.append([results_dict['boxes'], batch['pid']])
val_results_list.append(({k:v for k,v in results_dict.items() if k != "seg_preds"}, batch["pid"]))
_, monitor_metrics['val'] = val_evaluator.evaluate_predictions(val_results_list, monitor_metrics['val'])
model_selector.run_model_selection(net, optimizer, monitor_metrics, epoch)
# update monitoring and prediction plots
monitor_metrics.update({"lr":
{str(g): group['lr'] for (g, group) in enumerate(optimizer.param_groups)}})
logger.metrics2tboard(monitor_metrics, global_step=epoch)
epoch_time = time.time() - start_time
logger.info('trained epoch {}: took {:.2f} s ({:.2f} s train / {:.2f} s val)'.format(
epoch, epoch_time, train_time, epoch_time-train_time))
batch = next(batch_gen['val_sampling'])
results_dict = net.train_forward(batch, is_validation=True)
- logger.info('plotting predictions from validation sampling.')
- plot_batch_prediction(batch, results_dict, cf)
+ logger.info('generating validation-sampling example plot.')
+ plot_batch_prediction(batch, results_dict, cf, outfile=os.path.join(
+ cf.plot_dir, 'pred_example_{}_val.png'.format(cf.fold)))
# -------------- scheduling -----------------
if cf.dynamic_lr_scheduling:
scheduler.step(monitor_metrics["val"][cf.scheduling_criterion][-1])
else:
for param_group in optimizer.param_groups:
param_group['lr'] = cf.learning_rate[epoch-1]
def test(logger):
"""
perform testing for a given fold (or hold out set). save stats in evaluator.
"""
logger.info('starting testing model of fold {} in exp {}'.format(cf.fold, cf.exp_dir))
net = model.net(cf, logger).cuda()
test_predictor = Predictor(cf, net, logger, mode='test')
test_evaluator = Evaluator(cf, logger, mode='test')
batch_gen = data_loader.get_test_generator(cf, logger)
test_results_list = test_predictor.predict_test_set(batch_gen, return_results=True)
test_evaluator.evaluate_predictions(test_results_list)
test_evaluator.score_test_df()
if __name__ == '__main__':
stime = time.time()
parser = argparse.ArgumentParser()
parser.add_argument('-m', '--mode', type=str, default='train_test',
help='one out of: train / test / train_test / analysis / create_exp')
parser.add_argument('-f','--folds', nargs='+', type=int, default=None,
help='None runs over all folds in CV. otherwise specify list of folds.')
parser.add_argument('--exp_dir', type=str, default='/path/to/experiment/directory',
help='path to experiment dir. will be created if non existent.')
parser.add_argument('--server_env', default=False, action='store_true',
help='change IO settings to deploy models on a cluster.')
parser.add_argument('--data_dest', type=str, default=None, help="path to final data folder if different from config.")
parser.add_argument('--use_stored_settings', default=False, action='store_true',
help='load configs from existing exp_dir instead of source dir. always done for testing, '
'but can be set to true to do the same for training. useful in job scheduler environment, '
'where source code might change before the job actually runs.')
parser.add_argument('--resume_to_checkpoint', type=str, default=None,
help='if resuming to checkpoint, the desired fold still needs to be parsed via --folds.')
parser.add_argument('--exp_source', type=str, default='experiments/toy_exp',
help='specifies, from which source experiment to load configs and data_loader.')
+ parser.add_argument('--no_benchmark', action='store_true', help="Do not use cudnn.benchmark.")
parser.add_argument('-d', '--dev', default=False, action='store_true', help="development mode: shorten everything")
args = parser.parse_args()
folds = args.folds
resume_to_checkpoint = args.resume_to_checkpoint
+ torch.backends.cudnn.benchmark = not args.no_benchmark
+
if args.mode == 'train' or args.mode == 'train_test':
cf = utils.prep_exp(args.exp_source, args.exp_dir, args.server_env, args.use_stored_settings)
if args.dev:
folds = [0,1]
cf.batch_size, cf.num_epochs, cf.min_save_thresh, cf.save_n_models = 3 if cf.dim==2 else 1, 1, 0, 1
cf.num_train_batches, cf.num_val_batches, cf.max_val_patients = 5, 1, 1
cf.test_n_epochs = cf.save_n_models
cf.max_test_patients = 1
cf.data_dest = args.data_dest
logger = utils.get_logger(cf.exp_dir, cf.server_env)
+ logger.info("cudnn benchmark: {}, deterministic: {}.".format(torch.backends.cudnn.benchmark,
+ torch.backends.cudnn.deterministic))
data_loader = utils.import_module('dl', os.path.join(args.exp_source, 'data_loader.py'))
model = utils.import_module('model', cf.model_path)
logger.info("loaded model from {}".format(cf.model_path))
if folds is None:
folds = range(cf.n_cv_splits)
for fold in folds:
cf.fold_dir = os.path.join(cf.exp_dir, 'fold_{}'.format(fold))
cf.fold = fold
cf.resume_to_checkpoint = resume_to_checkpoint
if not os.path.exists(cf.fold_dir):
os.mkdir(cf.fold_dir)
logger.set_logfile(fold=fold)
train(logger)
cf.resume_to_checkpoint = None
if args.mode == 'train_test':
test(logger)
elif args.mode == 'test':
cf = utils.prep_exp(args.exp_source, args.exp_dir, args.server_env, is_training=False, use_stored_settings=True)
if args.dev:
folds = [0,1]
cf.test_n_epochs = 1; cf.max_test_patients = 1
cf.data_dest = args.data_dest
logger = utils.get_logger(cf.exp_dir, cf.server_env)
data_loader = utils.import_module('dl', os.path.join(args.exp_source, 'data_loader.py'))
model = utils.import_module('model', cf.model_path)
logger.info("loaded model from {}".format(cf.model_path))
if folds is None:
folds = range(cf.n_cv_splits)
for fold in folds:
cf.fold_dir = os.path.join(cf.exp_dir, 'fold_{}'.format(fold))
cf.fold = fold
logger.set_logfile(fold=fold)
test(logger)
# load raw predictions saved by predictor during testing, run aggregation algorithms and evaluation.
elif args.mode == 'analysis':
cf = utils.prep_exp(args.exp_source, args.exp_dir, args.server_env, is_training=False, use_stored_settings=True)
logger = utils.get_logger(cf.exp_dir, cf.server_env)
if cf.hold_out_test_set:
cf.folds = args.folds
predictor = Predictor(cf, net=None, logger=logger, mode='analysis')
results_list = predictor.load_saved_predictions(apply_wbc=True)
utils.create_csv_output([(res_dict["boxes"], pid) for res_dict, pid in results_list], cf, logger)
else:
if folds is None:
folds = range(cf.n_cv_splits)
for fold in folds:
cf.fold_dir = os.path.join(cf.exp_dir, 'fold_{}'.format(fold))
cf.fold = fold
logger.set_logfile(fold=fold)
predictor = Predictor(cf, net=None, logger=logger, mode='analysis')
results_list = predictor.load_saved_predictions(apply_wbc=True)
logger.info('starting evaluation...')
evaluator = Evaluator(cf, logger, mode='test')
evaluator.evaluate_predictions(results_list)
evaluator.score_test_df()
# create experiment folder and copy scripts without starting job.
# useful for cloud deployment where configs might change before job actually runs.
elif args.mode == 'create_exp':
cf = utils.prep_exp(args.exp_source, args.exp_dir, args.server_env, use_stored_settings=False)
logger = utils.get_logger(cf.exp_dir)
logger.info('created experiment directory at {}'.format(cf.exp_dir))
else:
raise RuntimeError('mode specified in args is not implemented...')
mins, secs = divmod((time.time() - stime), 60)
h, mins = divmod(mins, 60)
t = "{:d}h:{:02d}m:{:02d}s".format(int(h), int(mins), int(secs))
logger.info("{} total runtime: {}".format(os.path.split(__file__)[1], t))
del logger
\ No newline at end of file
diff --git a/experiments/lidc_exp/configs.py b/experiments/lidc_exp/configs.py
index 97bdb23..6db1598 100644
--- a/experiments/lidc_exp/configs.py
+++ b/experiments/lidc_exp/configs.py
@@ -1,341 +1,341 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import sys
import os
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
import numpy as np
from default_configs import DefaultConfigs
class configs(DefaultConfigs):
def __init__(self, server_env=None):
#########################
# Preprocessing #
#########################
self.root_dir = '/home/gregor/networkdrives/E130-Personal/Goetz/Datenkollektive/Lungendaten/Nodules_LIDC_IDRI'
self.raw_data_dir = '{}/new_nrrd'.format(self.root_dir)
self.pp_dir = '/media/gregor/HDD2TB/data/lidc/lidc_mdt'
self.target_spacing = (0.7, 0.7, 1.25)
#########################
# I/O #
#########################
# one out of [2, 3]. dimension the model operates in.
self.dim = 2
# one out of ['mrcnn', 'retina_net', 'retina_unet', 'detection_unet', 'ufrcnn'].
self.model = 'retina_unet'
DefaultConfigs.__init__(self, self.model, server_env, self.dim)
# int [0 < dataset_size]. select n patients from dataset for prototyping. If None, all data is used.
self.select_prototype_subset = None
# path to preprocessed data.
self.pp_name = 'lidc_mdt'
self.input_df_name = 'info_df.pickle'
self.pp_data_path = '/media/gregor/HDD2TB/data/lidc/{}'.format(self.pp_name)
self.pp_test_data_path = self.pp_data_path #change if test_data in separate folder.
# settings for deployment in cloud.
if server_env:
# path to preprocessed data.
self.pp_name = 'lidc_mdt_npz'
self.crop_name = 'pp_fg_slices_packed'
self.pp_data_path = '/datasets/datasets_ramien/lidc_exp/data/{}'.format(self.pp_name)
self.pp_test_data_path = self.pp_data_path
self.select_prototype_subset = None
#########################
# Data Loader #
#########################
# select modalities from preprocessed data
self.channels = [0]
self.n_channels = len(self.channels)
# patch_size to be used for training. pre_crop_size is the patch_size before data augmentation.
self.pre_crop_size_2D = [300, 300]
self.patch_size_2D = [288, 288]
self.pre_crop_size_3D = [156, 156, 96]
self.patch_size_3D = [128, 128, 64]
self.patch_size = self.patch_size_2D if self.dim == 2 else self.patch_size_3D
self.pre_crop_size = self.pre_crop_size_2D if self.dim == 2 else self.pre_crop_size_3D
# ratio of free sampled batch elements before class balancing is triggered
# (>0 to include "empty"/background patches.)
self.batch_sample_slack = 0.2
# set 2D network to operate in 3D images.
self.merge_2D_to_3D_preds = self.dim == 2
# feed +/- n neighbouring slices into channel dimension. set to None for no context.
self.n_3D_context = None
if self.n_3D_context is not None and self.dim == 2:
self.n_channels *= (self.n_3D_context * 2 + 1)
#########################
# Architecture #
#########################
self.start_filts = 48 if self.dim == 2 else 18
self.end_filts = self.start_filts * 4 if self.dim == 2 else self.start_filts * 2
self.res_architecture = 'resnet50' # 'resnet101' , 'resnet50'
self.norm = None # one of None, 'instance_norm', 'batch_norm'
self.weight_decay = 0
# one of 'xavier_uniform', 'xavier_normal', or 'kaiming_normal', None (=default = 'kaiming_uniform')
self.weight_init = None
#########################
# Schedule / Selection #
#########################
self.num_epochs = 100
self.num_train_batches = 200 if self.dim == 2 else 200
self.batch_size = 20 if self.dim == 2 else 8
self.do_validation = True
# decide whether to validate on entire patient volumes (like testing) or sampled patches (like training)
# the former is morge accurate, while the latter is faster (depending on volume size)
self.val_mode = 'val_sampling' # one of 'val_sampling' , 'val_patient'
if self.val_mode == 'val_patient':
self.max_val_patients = 50 # if 'None' iterates over entire val_set once.
if self.val_mode == 'val_sampling':
self.num_val_batches = 50
# set dynamic_lr_scheduling to True to apply LR scheduling with below settings.
self.dynamic_lr_scheduling = True
self.lr_decay_factor = 0.5
- self.scheduling_patience = int(self.num_train_batches * self.batch_size / 6000)
+ self.scheduling_patience = np.ceil(6000 / (self.num_train_batches * self.batch_size))
self.scheduling_criterion = 'malignant_ap'
self.scheduling_mode = 'min' if "loss" in self.scheduling_criterion else 'max'
#########################
# Testing / Plotting #
#########################
# set the top-n-epochs to be saved for temporal averaging in testing.
self.save_n_models = 5
self.test_n_epochs = 5
# set a minimum epoch number for saving in case of instabilities in the first phase of training.
self.min_save_thresh = 0 if self.dim == 2 else 0
self.report_score_level = ['patient', 'rois'] # choose list from 'patient', 'rois'
self.class_dict = {1: 'benign', 2: 'malignant'} # 0 is background.
self.patient_class_of_interest = 2 # patient metrics are only plotted for one class.
self.ap_match_ious = [0.1] # list of ious to be evaluated for ap-scoring.
self.model_selection_criteria = ['malignant_ap', 'benign_ap'] # criteria to average over for saving epochs.
self.min_det_thresh = 0.1 # minimum confidence value to select predictions for evaluation.
# threshold for clustering predictions together (wcs = weighted cluster scoring).
# needs to be >= the expected overlap of predictions coming from one model (typically NMS threshold).
# if too high, preds of the same object are separate clusters.
self.wcs_iou = 1e-5
self.plot_prediction_histograms = True
self.plot_stat_curves = False
#########################
# Data Augmentation #
#########################
self.da_kwargs={
'do_elastic_deform': True,
'alpha':(0., 1500.),
'sigma':(30., 50.),
'do_rotation':True,
'angle_x': (0., 2 * np.pi),
'angle_y': (0., 0),
'angle_z': (0., 0),
'do_scale': True,
'scale':(0.8, 1.1),
'random_crop':False,
'rand_crop_dist': (self.patch_size[0] / 2. - 3, self.patch_size[1] / 2. - 3),
'border_mode_data': 'constant',
'border_cval_data': 0,
'order_data': 1
}
if self.dim == 3:
self.da_kwargs['do_elastic_deform'] = False
self.da_kwargs['angle_x'] = (0, 0.0)
self.da_kwargs['angle_y'] = (0, 0.0) #must be 0!!
self.da_kwargs['angle_z'] = (0., 2 * np.pi)
#########################
# Add model specifics #
#########################
{'detection_unet': self.add_det_unet_configs,
'mrcnn': self.add_mrcnn_configs,
'ufrcnn': self.add_mrcnn_configs,
'retina_net': self.add_mrcnn_configs,
'retina_unet': self.add_mrcnn_configs,
}[self.model]()
def add_det_unet_configs(self):
self.learning_rate = [1e-4] * self.num_epochs
# aggregation from pixel perdiction to object scores (connected component). One of ['max', 'median']
self.aggregation_operation = 'max'
# max number of roi candidates to identify per batch element and class.
self.n_roi_candidates = 10 if self.dim == 2 else 30
# loss mode: either weighted cross entropy ('wce'), batch-wise dice loss ('dice), or the sum of both ('dice_wce')
self.seg_loss_mode = 'dice_wce'
# if <1, false positive predictions in foreground are penalized less.
self.fp_dice_weight = 1 if self.dim == 2 else 1
self.wce_weights = [1, 1, 1]
self.detection_min_confidence = self.min_det_thresh
# if 'True', loss distinguishes all classes, else only foreground vs. background (class agnostic).
self.class_specific_seg_flag = True
self.num_seg_classes = 3 if self.class_specific_seg_flag else 2
self.head_classes = self.num_seg_classes
def add_mrcnn_configs(self):
# learning rate is a list with one entry per epoch.
self.learning_rate = [1e-4] * self.num_epochs
# disable the re-sampling of mask proposals to original size for speed-up.
# since evaluation is detection-driven (box-matching) and not instance segmentation-driven (iou-matching),
# mask-outputs are optional.
self.return_masks_in_val = True
self.return_masks_in_test = False
# set number of proposal boxes to plot after each epoch.
self.n_plot_rpn_props = 5 if self.dim == 2 else 30
# number of classes for head networks: n_foreground_classes + 1 (background)
self.head_classes = 3
# seg_classes hier refers to the first stage classifier (RPN)
self.num_seg_classes = 2 # foreground vs. background
# feature map strides per pyramid level are inferred from architecture.
self.backbone_strides = {'xy': [4, 8, 16, 32], 'z': [1, 2, 4, 8]}
# anchor scales are chosen according to expected object sizes in data set. Default uses only one anchor scale
# per pyramid level. (outer list are pyramid levels (corresponding to BACKBONE_STRIDES), inner list are scales per level.)
self.rpn_anchor_scales = {'xy': [[8], [16], [32], [64]], 'z': [[2], [4], [8], [16]]}
# choose which pyramid levels to extract features from: P2: 0, P3: 1, P4: 2, P5: 3.
self.pyramid_levels = [0, 1, 2, 3]
# number of feature maps in rpn. typically lowered in 3D to save gpu-memory.
self.n_rpn_features = 512 if self.dim == 2 else 128
# anchor ratios and strides per position in feature maps.
self.rpn_anchor_ratios = [0.5, 1, 2]
self.rpn_anchor_stride = 1
# Threshold for first stage (RPN) non-maximum suppression (NMS): LOWER == HARDER SELECTION
self.rpn_nms_threshold = 0.7 if self.dim == 2 else 0.7
# loss sampling settings.
self.rpn_train_anchors_per_image = 6 #per batch element
self.train_rois_per_image = 6 #per batch element
self.roi_positive_ratio = 0.5
self.anchor_matching_iou = 0.7
# factor of top-k candidates to draw from per negative sample (stochastic-hard-example-mining).
# poolsize to draw top-k candidates from will be shem_poolsize * n_negative_samples.
self.shem_poolsize = 10
self.pool_size = (7, 7) if self.dim == 2 else (7, 7, 3)
self.mask_pool_size = (14, 14) if self.dim == 2 else (14, 14, 5)
self.mask_shape = (28, 28) if self.dim == 2 else (28, 28, 10)
self.rpn_bbox_std_dev = np.array([0.1, 0.1, 0.1, 0.2, 0.2, 0.2])
self.bbox_std_dev = np.array([0.1, 0.1, 0.1, 0.2, 0.2, 0.2])
self.window = np.array([0, 0, self.patch_size[0], self.patch_size[1], 0, self.patch_size_3D[2]])
self.scale = np.array([self.patch_size[0], self.patch_size[1], self.patch_size[0], self.patch_size[1],
self.patch_size_3D[2], self.patch_size_3D[2]])
if self.dim == 2:
self.rpn_bbox_std_dev = self.rpn_bbox_std_dev[:4]
self.bbox_std_dev = self.bbox_std_dev[:4]
self.window = self.window[:4]
self.scale = self.scale[:4]
# pre-selection in proposal-layer (stage 1) for NMS-speedup. applied per batch element.
self.pre_nms_limit = 3000 if self.dim == 2 else 6000
# n_proposals to be selected after NMS per batch element. too high numbers blow up memory if "detect_while_training" is True,
# since proposals of the entire batch are forwarded through second stage in as one "batch".
self.roi_chunk_size = 2500 if self.dim == 2 else 600
self.post_nms_rois_training = 500 if self.dim == 2 else 75
self.post_nms_rois_inference = 500
# Final selection of detections (refine_detections)
self.model_max_instances_per_batch_element = 10 if self.dim == 2 else 30 # per batch element and class.
self.detection_nms_threshold = 1e-5 # needs to be > 0, otherwise all predictions are one cluster.
self.model_min_confidence = 0.1
if self.dim == 2:
self.backbone_shapes = np.array(
[[int(np.ceil(self.patch_size[0] / stride)),
int(np.ceil(self.patch_size[1] / stride))]
for stride in self.backbone_strides['xy']])
else:
self.backbone_shapes = np.array(
[[int(np.ceil(self.patch_size[0] / stride)),
int(np.ceil(self.patch_size[1] / stride)),
int(np.ceil(self.patch_size[2] / stride_z))]
for stride, stride_z in zip(self.backbone_strides['xy'], self.backbone_strides['z']
)])
if self.model == 'ufrcnn':
self.operate_stride1 = True
self.class_specific_seg_flag = True
self.num_seg_classes = 3 if self.class_specific_seg_flag else 2
self.frcnn_mode = True
if self.model == 'retina_net' or self.model == 'retina_unet' or self.model == 'prob_detector':
# implement extra anchor-scales according to retina-net publication.
self.rpn_anchor_scales['xy'] = [[ii[0], ii[0] * (2 ** (1 / 3)), ii[0] * (2 ** (2 / 3))] for ii in
self.rpn_anchor_scales['xy']]
self.rpn_anchor_scales['z'] = [[ii[0], ii[0] * (2 ** (1 / 3)), ii[0] * (2 ** (2 / 3))] for ii in
self.rpn_anchor_scales['z']]
self.n_anchors_per_pos = len(self.rpn_anchor_ratios) * 3
self.n_rpn_features = 256 if self.dim == 2 else 64
# pre-selection of detections for NMS-speedup. per entire batch.
self.pre_nms_limit = 10000 if self.dim == 2 else 50000
# anchor matching iou is lower than in Mask R-CNN according to https://arxiv.org/abs/1708.02002
self.anchor_matching_iou = 0.5
# if 'True', seg loss distinguishes all classes, else only foreground vs. background (class agnostic).
self.num_seg_classes = 3 if self.class_specific_seg_flag else 2
if self.model == 'retina_unet':
self.operate_stride1 = True
diff --git a/experiments/toy_exp/configs.py b/experiments/toy_exp/configs.py
index f37f262..8863f70 100644
--- a/experiments/toy_exp/configs.py
+++ b/experiments/toy_exp/configs.py
@@ -1,350 +1,351 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import sys
import os
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
import numpy as np
from default_configs import DefaultConfigs
class configs(DefaultConfigs):
def __init__(self, server_env=None):
#########################
# Preprocessing #
#########################
self.root_dir = '/home/gregor/datasets/toy_mdt'
#########################
# I/O #
#########################
# one out of [2, 3]. dimension the model operates in.
self.dim = 2
# one out of ['mrcnn', 'retina_net', 'retina_unet', 'detection_unet', 'ufrcnn'].
self.model = 'retina_net'
DefaultConfigs.__init__(self, self.model, server_env, self.dim)
# int [0 < dataset_size]. select n patients from dataset for prototyping.
self.select_prototype_subset = None
self.hold_out_test_set = True
- self.n_train_data = 2500
+ # including val set. will be 3/4 train, 1/4 val.
+ self.n_train_val_data = 2500
# choose one of the 3 toy experiments described in https://arxiv.org/pdf/1811.08661.pdf
# one of ['donuts_shape', 'donuts_pattern', 'circles_scale'].
toy_mode = 'donuts_shape_noise'
# path to preprocessed data.
self.input_df_name = 'info_df.pickle'
self.pp_name = os.path.join(toy_mode, 'train')
self.pp_data_path = os.path.join(self.root_dir, self.pp_name)
self.pp_test_name = os.path.join(toy_mode, 'test')
self.pp_test_data_path = os.path.join(self.root_dir, self.pp_test_name)
# settings for deployment in cloud.
if server_env:
# path to preprocessed data.
pp_root_dir = '/datasets/datasets_ramien/toy_exp/data'
self.pp_name = os.path.join(toy_mode, 'train')
self.pp_data_path = os.path.join(pp_root_dir, self.pp_name)
self.pp_test_name = os.path.join(toy_mode, 'test')
self.pp_test_data_path = os.path.join(pp_root_dir, self.pp_test_name)
self.select_prototype_subset = None
#########################
# Data Loader #
#########################
# select modalities from preprocessed data
self.channels = [0]
self.n_channels = len(self.channels)
# patch_size to be used for training. pre_crop_size is the patch_size before data augmentation.
self.pre_crop_size_2D = [320, 320]
self.patch_size_2D = [320, 320]
self.patch_size = self.patch_size_2D if self.dim == 2 else self.patch_size_3D
self.pre_crop_size = self.pre_crop_size_2D if self.dim == 2 else self.pre_crop_size_3D
# ratio of free sampled batch elements before class balancing is triggered
# (>0 to include "empty"/background patches.)
self.batch_sample_slack = 0.2
# set 2D network to operate in 3D images.
self.merge_2D_to_3D_preds = False
# feed +/- n neighbouring slices into channel dimension. set to None for no context.
self.n_3D_context = None
if self.n_3D_context is not None and self.dim == 2:
self.n_channels *= (self.n_3D_context * 2 + 1)
#########################
# Architecture #
#########################
self.start_filts = 48 if self.dim == 2 else 18
self.end_filts = self.start_filts * 4 if self.dim == 2 else self.start_filts * 2
self.res_architecture = 'resnet50' # 'resnet101' , 'resnet50'
self.norm = None # one of None, 'instance_norm', 'batch_norm'
- self.weight_decay = 0
+ self.weight_decay = 3e-5
# one of 'xavier_uniform', 'xavier_normal', or 'kaiming_normal', None (=default = 'kaiming_uniform')
self.weight_init = None
#########################
# Schedule / Selection #
#########################
- self.num_epochs = 22
+ self.num_epochs = 24
self.num_train_batches = 100 if self.dim == 2 else 200
self.batch_size = 20 if self.dim == 2 else 8
self.do_validation = True
# decide whether to validate on entire patient volumes (like testing) or sampled patches (like training)
# the former is morge accurate, while the latter is faster (depending on volume size)
self.val_mode = 'val_patient' # one of 'val_sampling' , 'val_patient'
if self.val_mode == 'val_patient':
self.max_val_patients = None # if 'None' iterates over entire val_set once.
if self.val_mode == 'val_sampling':
self.num_val_batches = 50
# set dynamic_lr_scheduling to True to apply LR scheduling with below settings.
self.dynamic_lr_scheduling = True
self.lr_decay_factor = 0.5
- self.scheduling_patience = int(self.num_train_batches * self.batch_size / 2400)
+ self.scheduling_patience = np.ceil(3600 / (self.num_train_batches * self.batch_size))
self.scheduling_criterion = 'malignant_ap'
self.scheduling_mode = 'min' if "loss" in self.scheduling_criterion else 'max'
#########################
# Testing / Plotting #
#########################
# set the top-n-epochs to be saved for temporal averaging in testing.
self.save_n_models = 5
self.test_n_epochs = 5
# set a minimum epoch number for saving in case of instabilities in the first phase of training.
self.min_save_thresh = 0 if self.dim == 2 else 0
self.report_score_level = ['patient', 'rois'] # choose list from 'patient', 'rois'
self.class_dict = {1: 'benign', 2: 'malignant'} # 0 is background.
self.patient_class_of_interest = 2 # patient metrics are only plotted for one class.
self.ap_match_ious = [0.1] # list of ious to be evaluated for ap-scoring.
self.model_selection_criteria = ['benign_ap', 'malignant_ap'] # criteria to average over for saving epochs.
self.min_det_thresh = 0.1 # minimum confidence value to select predictions for evaluation.
# threshold for clustering predictions together (wcs = weighted cluster scoring).
# needs to be >= the expected overlap of predictions coming from one model (typically NMS threshold).
# if too high, preds of the same object are separate clusters.
self.wcs_iou = 1e-5
self.plot_prediction_histograms = True
self.plot_stat_curves = False
#########################
# Data Augmentation #
#########################
self.da_kwargs={
'do_elastic_deform': True,
'alpha':(0., 1500.),
'sigma':(30., 50.),
'do_rotation':True,
'angle_x': (0., 2 * np.pi),
'angle_y': (0., 0),
'angle_z': (0., 0),
'do_scale': True,
'scale':(0.8, 1.1),
'random_crop':False,
'rand_crop_dist': (self.patch_size[0] / 2. - 3, self.patch_size[1] / 2. - 3),
'border_mode_data': 'constant',
'border_cval_data': 0,
'order_data': 1
}
if self.dim == 3:
self.da_kwargs['do_elastic_deform'] = False
self.da_kwargs['angle_x'] = (0, 0.0)
self.da_kwargs['angle_y'] = (0, 0.0) #must be 0!!
self.da_kwargs['angle_z'] = (0., 2 * np.pi)
#########################
# Add model specifics #
#########################
{'detection_unet': self.add_det_unet_configs,
'mrcnn': self.add_mrcnn_configs,
'ufrcnn': self.add_mrcnn_configs,
'ufrcnn_surrounding': self.add_mrcnn_configs,
'retina_net': self.add_mrcnn_configs,
'retina_unet': self.add_mrcnn_configs,
'prob_detector': self.add_mrcnn_configs,
}[self.model]()
def add_det_unet_configs(self):
self.learning_rate = [1e-4] * self.num_epochs
# aggregation from pixel perdiction to object scores (connected component). One of ['max', 'median']
self.aggregation_operation = 'max'
# max number of roi candidates to identify per image (slice in 2D, volume in 3D)
self.n_roi_candidates = 3 if self.dim == 2 else 8
# loss mode: either weighted cross entropy ('wce'), batch-wise dice loss ('dice), or the sum of both ('dice_wce')
- self.seg_loss_mode = 'dice_wce'
+ self.seg_loss_mode = 'wce'
# if <1, false positive predictions in foreground are penalized less.
self.fp_dice_weight = 1 if self.dim == 2 else 1
self.wce_weights = [1, 1, 1]
self.detection_min_confidence = self.min_det_thresh
# if 'True', loss distinguishes all classes, else only foreground vs. background (class agnostic).
self.class_specific_seg_flag = True
self.num_seg_classes = 3 if self.class_specific_seg_flag else 2
self.head_classes = self.num_seg_classes
def add_mrcnn_configs(self):
# learning rate is a list with one entry per epoch.
- self.learning_rate = [1e-4] * self.num_epochs
+ self.learning_rate = [3e-4] * self.num_epochs
# disable mask head loss. (e.g. if no pixelwise annotations available)
self.frcnn_mode = False
# disable the re-sampling of mask proposals to original size for speed-up.
# since evaluation is detection-driven (box-matching) and not instance segmentation-driven (iou-matching),
# mask-outputs are optional.
self.return_masks_in_val = True
self.return_masks_in_test = False
# set number of proposal boxes to plot after each epoch.
self.n_plot_rpn_props = 0 if self.dim == 2 else 0
# number of classes for head networks: n_foreground_classes + 1 (background)
self.head_classes = 3
# seg_classes hier refers to the first stage classifier (RPN)
self.num_seg_classes = 2 # foreground vs. background
# feature map strides per pyramid level are inferred from architecture.
self.backbone_strides = {'xy': [4, 8, 16, 32], 'z': [1, 2, 4, 8]}
# anchor scales are chosen according to expected object sizes in data set. Default uses only one anchor scale
# per pyramid level. (outer list are pyramid levels (corresponding to BACKBONE_STRIDES), inner list are scales per level.)
self.rpn_anchor_scales = {'xy': [[8], [16], [32], [64]], 'z': [[2], [4], [8], [16]]}
# choose which pyramid levels to extract features from: P2: 0, P3: 1, P4: 2, P5: 3.
self.pyramid_levels = [0, 1, 2, 3]
# number of feature maps in rpn. typically lowered in 3D to save gpu-memory.
self.n_rpn_features = 512 if self.dim == 2 else 128
# anchor ratios and strides per position in feature maps.
self.rpn_anchor_ratios = [0.5, 1., 2.]
self.rpn_anchor_stride = 1
# Threshold for first stage (RPN) non-maximum suppression (NMS): LOWER == HARDER SELECTION
self.rpn_nms_threshold = 0.7 if self.dim == 2 else 0.7
# loss sampling settings.
- self.rpn_train_anchors_per_image = 2 #per batch element
+ self.rpn_train_anchors_per_image = 64 #per batch element
self.train_rois_per_image = 2 #per batch element
self.roi_positive_ratio = 0.5
self.anchor_matching_iou = 0.7
# factor of top-k candidates to draw from per negative sample (stochastic-hard-example-mining).
# poolsize to draw top-k candidates from will be shem_poolsize * n_negative_samples.
self.shem_poolsize = 10
self.pool_size = (7, 7) if self.dim == 2 else (7, 7, 3)
self.mask_pool_size = (14, 14) if self.dim == 2 else (14, 14, 5)
self.mask_shape = (28, 28) if self.dim == 2 else (28, 28, 10)
self.rpn_bbox_std_dev = np.array([0.1, 0.1, 0.1, 0.2, 0.2, 0.2])
self.bbox_std_dev = np.array([0.1, 0.1, 0.1, 0.2, 0.2, 0.2])
self.window = np.array([0, 0, self.patch_size[0], self.patch_size[1]])
self.scale = np.array([self.patch_size[0], self.patch_size[1], self.patch_size[0], self.patch_size[1]])
if self.dim == 2:
self.rpn_bbox_std_dev = self.rpn_bbox_std_dev[:4]
self.bbox_std_dev = self.bbox_std_dev[:4]
self.window = self.window[:4]
self.scale = self.scale[:4]
# pre-selection in proposal-layer (stage 1) for NMS-speedup. applied per batch element.
self.pre_nms_limit = 3000 if self.dim == 2 else 6000
# n_proposals to be selected after NMS per batch element. too high numbers blow up memory if "detect_while_training" is True,
# since proposals of the entire batch are forwarded through second stage in as one "batch".
self.roi_chunk_size = 800 if self.dim == 2 else 600
self.post_nms_rois_training = 500 if self.dim == 2 else 75
self.post_nms_rois_inference = 500
# Final selection of detections (refine_detections)
self.model_max_instances_per_batch_element = 10 if self.dim == 2 else 30 # per batch element and class.
self.detection_nms_threshold = 1e-5 # needs to be > 0, otherwise all predictions are one cluster.
self.model_min_confidence = 0.1
if self.dim == 2:
self.backbone_shapes = np.array(
[[int(np.ceil(self.patch_size[0] / stride)),
int(np.ceil(self.patch_size[1] / stride))]
for stride in self.backbone_strides['xy']])
else:
self.backbone_shapes = np.array(
[[int(np.ceil(self.patch_size[0] / stride)),
int(np.ceil(self.patch_size[1] / stride)),
int(np.ceil(self.patch_size[2] / stride_z))]
for stride, stride_z in zip(self.backbone_strides['xy'], self.backbone_strides['z']
)])
if self.model == 'ufrcnn':
self.operate_stride1 = True
self.class_specific_seg_flag = True
self.num_seg_classes = 3 if self.class_specific_seg_flag else 2
self.frcnn_mode = True
if self.model == 'retina_net' or self.model == 'retina_unet' or self.model == 'prob_detector':
# implement extra anchor-scales according to retina-net publication.
self.rpn_anchor_scales['xy'] = [[ii[0], ii[0] * (2 ** (1 / 3)), ii[0] * (2 ** (2 / 3))] for ii in
self.rpn_anchor_scales['xy']]
self.rpn_anchor_scales['z'] = [[ii[0], ii[0] * (2 ** (1 / 3)), ii[0] * (2 ** (2 / 3))] for ii in
self.rpn_anchor_scales['z']]
self.n_anchors_per_pos = len(self.rpn_anchor_ratios) * 3
self.n_rpn_features = 256 if self.dim == 2 else 64
# pre-selection of detections for NMS-speedup. per entire batch.
self.pre_nms_limit = 10000 if self.dim == 2 else 50000
# anchor matching iou is lower than in Mask R-CNN according to https://arxiv.org/abs/1708.02002
self.anchor_matching_iou = 0.5
# if 'True', seg loss distinguishes all classes, else only foreground vs. background (class agnostic).
self.num_seg_classes = 3 if self.class_specific_seg_flag else 2
if self.model == 'retina_unet':
self.operate_stride1 = True
diff --git a/experiments/toy_exp/data_loader.py b/experiments/toy_exp/data_loader.py
index e1c84b4..29490e3 100644
--- a/experiments/toy_exp/data_loader.py
+++ b/experiments/toy_exp/data_loader.py
@@ -1,309 +1,309 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import numpy as np
import os
from collections import OrderedDict
import pandas as pd
import pickle
import time
import subprocess
import utils.dataloader_utils as dutils
# batch generator tools from https://github.com/MIC-DKFZ/batchgenerators
from batchgenerators.dataloading.data_loader import SlimDataLoaderBase
from batchgenerators.transforms.spatial_transforms import MirrorTransform as Mirror
from batchgenerators.transforms.abstract_transforms import Compose
from batchgenerators.dataloading.multi_threaded_augmenter import MultiThreadedAugmenter
from batchgenerators.dataloading import SingleThreadedAugmenter
from batchgenerators.transforms.spatial_transforms import SpatialTransform
from batchgenerators.transforms.crop_and_pad_transforms import CenterCropTransform
from batchgenerators.transforms.utility_transforms import ConvertSegToBoundingBoxCoordinates
def get_train_generators(cf, logger):
"""
wrapper function for creating the training batch generator pipeline. returns the train/val generators.
selects patients according to cv folds (generated by first run/fold of experiment):
splits the data into n-folds, where 1 split is used for val, 1 split for testing and the rest for training. (inner loop test set)
If cf.hold_out_test_set is True, adds the test split to the training data.
"""
all_data = load_dataset(cf, logger)
all_pids_list = np.unique([v['pid'] for (k, v) in all_data.items()])
- train_pids = all_pids_list[:cf.n_train_data]
- val_pids = all_pids_list[1000:1500]
+ train_pids = all_pids_list[:int(2*cf.n_train_val_data//3)]
+ val_pids = all_pids_list[int(np.ceil(2*cf.n_train_val_data//3)):cf.n_train_val_data]
train_data = {k: v for (k, v) in all_data.items() if any(p == v['pid'] for p in train_pids)}
val_data = {k: v for (k, v) in all_data.items() if any(p == v['pid'] for p in val_pids)}
logger.info("data set loaded with: {} train / {} val patients".format(len(train_pids), len(val_pids)))
batch_gen = {}
batch_gen['train'] = create_data_gen_pipeline(train_data, cf=cf, do_aug=False)
batch_gen['val_sampling'] = create_data_gen_pipeline(val_data, cf=cf, do_aug=False)
if cf.val_mode == 'val_patient':
batch_gen['val_patient'] = PatientBatchIterator(val_data, cf=cf)
batch_gen['n_val'] = len(val_pids) if cf.max_val_patients is None else min(len(val_pids), cf.max_val_patients)
else:
batch_gen['n_val'] = cf.num_val_batches
return batch_gen
def get_test_generator(cf, logger):
"""
wrapper function for creating the test batch generator pipeline.
selects patients according to cv folds (generated by first run/fold of experiment)
If cf.hold_out_test_set is True, gets the data from an external folder instead.
"""
if cf.hold_out_test_set:
pp_name = cf.pp_test_name
test_ix = None
else:
pp_name = None
with open(os.path.join(cf.exp_dir, 'fold_ids.pickle'), 'rb') as handle:
fold_list = pickle.load(handle)
_, _, test_ix, _ = fold_list[cf.fold]
# warnings.warn('WARNING: using validation set for testing!!!')
test_data = load_dataset(cf, logger, test_ix, pp_data_path=cf.pp_test_data_path, pp_name=pp_name)
logger.info("data set loaded with: {} test patients from {}".format(len(test_data.keys()), cf.pp_test_data_path))
batch_gen = {}
batch_gen['test'] = PatientBatchIterator(test_data, cf=cf)
batch_gen['n_test'] = len(test_data.keys()) if cf.max_test_patients=="all" else \
min(cf.max_test_patients, len(test_data.keys()))
return batch_gen
def load_dataset(cf, logger, subset_ixs=None, pp_data_path=None, pp_name=None):
"""
loads the dataset. if deployed in cloud also copies and unpacks the data to the working directory.
:param subset_ixs: subset indices to be loaded from the dataset. used e.g. for testing to only load the test folds.
:return: data: dictionary with one entry per patient (in this case per patient-breast, since they are treated as
individual images for training) each entry is a dictionary containing respective meta-info as well as paths to the preprocessed
numpy arrays to be loaded during batch-generation
"""
if pp_data_path is None:
pp_data_path = cf.pp_data_path
if pp_name is None:
pp_name = cf.pp_name
if cf.server_env:
copy_data = True
target_dir = os.path.join(cf.data_dest, pp_name)
if not os.path.exists(target_dir):
cf.data_source_dir = pp_data_path
os.makedirs(target_dir)
subprocess.call('rsync -av {} {}'.format(
os.path.join(cf.data_source_dir, cf.input_df_name), os.path.join(target_dir, cf.input_df_name)), shell=True)
logger.info('created target dir and info df at {}'.format(os.path.join(target_dir, cf.input_df_name)))
elif subset_ixs is None:
copy_data = False
pp_data_path = target_dir
p_df = pd.read_pickle(os.path.join(pp_data_path, cf.input_df_name))
if subset_ixs is not None:
subset_pids = [np.unique(p_df.pid.tolist())[ix] for ix in subset_ixs]
p_df = p_df[p_df.pid.isin(subset_pids)]
logger.info('subset: selected {} instances from df'.format(len(p_df)))
if cf.server_env:
if copy_data:
copy_and_unpack_data(logger, p_df.pid.tolist(), cf.fold_dir, cf.data_source_dir, target_dir)
class_targets = p_df['class_id'].tolist()
pids = p_df.pid.tolist()
imgs = [os.path.join(pp_data_path, '{}.npy'.format(pid)) for pid in pids]
segs = [os.path.join(pp_data_path,'{}.npy'.format(pid)) for pid in pids]
data = OrderedDict()
for ix, pid in enumerate(pids):
data[pid] = {'data': imgs[ix], 'seg': segs[ix], 'pid': pid, 'class_target': [class_targets[ix]]}
return data
def create_data_gen_pipeline(patient_data, cf, do_aug=True):
"""
create mutli-threaded train/val/test batch generation and augmentation pipeline.
:param patient_data: dictionary containing one dictionary per patient in the train/test subset.
:param is_training: (optional) whether to perform data augmentation (training) or not (validation/testing)
:return: multithreaded_generator
"""
# create instance of batch generator as first element in pipeline.
data_gen = BatchGenerator(patient_data, batch_size=cf.batch_size, cf=cf)
# add transformations to pipeline.
my_transforms = []
if do_aug:
mirror_transform = Mirror(axes=np.arange(2, cf.dim+2, 1))
my_transforms.append(mirror_transform)
spatial_transform = SpatialTransform(patch_size=cf.patch_size[:cf.dim],
patch_center_dist_from_border=cf.da_kwargs['rand_crop_dist'],
do_elastic_deform=cf.da_kwargs['do_elastic_deform'],
alpha=cf.da_kwargs['alpha'], sigma=cf.da_kwargs['sigma'],
do_rotation=cf.da_kwargs['do_rotation'], angle_x=cf.da_kwargs['angle_x'],
angle_y=cf.da_kwargs['angle_y'], angle_z=cf.da_kwargs['angle_z'],
do_scale=cf.da_kwargs['do_scale'], scale=cf.da_kwargs['scale'],
random_crop=cf.da_kwargs['random_crop'])
my_transforms.append(spatial_transform)
else:
my_transforms.append(CenterCropTransform(crop_size=cf.patch_size[:cf.dim]))
my_transforms.append(ConvertSegToBoundingBoxCoordinates(cf.dim, get_rois_from_seg_flag=False, class_specific_seg_flag=cf.class_specific_seg_flag))
all_transforms = Compose(my_transforms)
# multithreaded_generator = SingleThreadedAugmenter(data_gen, all_transforms)
multithreaded_generator = MultiThreadedAugmenter(data_gen, all_transforms, num_processes=cf.n_workers, seeds=range(cf.n_workers))
return multithreaded_generator
class BatchGenerator(SlimDataLoaderBase):
"""
creates the training/validation batch generator. Samples n_batch_size patients (draws a slice from each patient if 2D)
from the data set while maintaining foreground-class balance. Returned patches are cropped/padded to pre_crop_size.
Actual patch_size is obtained after data augmentation.
:param data: data dictionary as provided by 'load_dataset'.
:param batch_size: number of patients to sample for the batch
:return dictionary containing the batch data (b, c, x, y, (z)) / seg (b, 1, x, y, (z)) / pids / class_target
"""
def __init__(self, data, batch_size, cf):
super(BatchGenerator, self).__init__(data, batch_size)
self.cf = cf
def generate_train_batch(self):
batch_data, batch_segs, batch_pids, batch_targets = [], [], [], []
class_targets_list = [v['class_target'] for (k, v) in self._data.items()]
#samples patients towards equilibrium of foreground classes on a roi-level (after randomly sampling the ratio "batch_sample_slack).
batch_ixs = dutils.get_class_balanced_patients(
class_targets_list, self.batch_size, self.cf.head_classes - 1, slack_factor=self.cf.batch_sample_slack)
patients = list(self._data.items())
for b in batch_ixs:
patient = patients[b][1]
all_data = np.load(patient['data'], mmap_mode='r')
data = all_data[0]
seg = all_data[1].astype('uint8')
batch_pids.append(patient['pid'])
batch_targets.append(patient['class_target'])
batch_data.append(data[np.newaxis])
batch_segs.append(seg[np.newaxis])
data = np.array(batch_data)
seg = np.array(batch_segs).astype(np.uint8)
class_target = np.array(batch_targets)
return {'data': data, 'seg': seg, 'pid': batch_pids, 'class_target': class_target}
class PatientBatchIterator(SlimDataLoaderBase):
"""
creates a test generator that iterates over entire given dataset returning 1 patient per batch.
Can be used for monitoring if cf.val_mode = 'patient_val' for a monitoring closer to actualy evaluation (done in 3D),
if willing to accept speed-loss during training.
:return: out_batch: dictionary containing one patient with batch_size = n_3D_patches in 3D or
batch_size = n_2D_patches in 2D .
"""
def __init__(self, data, cf): #threads in augmenter
super(PatientBatchIterator, self).__init__(data, 0)
self.cf = cf
self.patient_ix = 0
self.dataset_pids = [v['pid'] for (k, v) in data.items()]
self.patch_size = cf.patch_size
if len(self.patch_size) == 2:
self.patch_size = self.patch_size + [1]
def generate_train_batch(self):
pid = self.dataset_pids[self.patient_ix]
patient = self._data[pid]
all_data = np.load(patient['data'], mmap_mode='r')
data = all_data[0]
seg = all_data[1].astype('uint8')
batch_class_targets = np.array([patient['class_target']])
out_data = data[None, None]
out_seg = seg[None, None]
- print('check patient data loader', out_data.shape, out_seg.shape)
+ #print('check patient data loader', out_data.shape, out_seg.shape)
batch_2D = {'data': out_data, 'seg': out_seg, 'class_target': batch_class_targets, 'pid': pid}
converter = ConvertSegToBoundingBoxCoordinates(dim=2, get_rois_from_seg_flag=False, class_specific_seg_flag=self.cf.class_specific_seg_flag)
batch_2D = converter(**batch_2D)
batch_2D.update({'patient_bb_target': batch_2D['bb_target'],
'patient_roi_labels': batch_2D['roi_labels'],
'original_img_shape': out_data.shape})
self.patient_ix += 1
if self.patient_ix == len(self.dataset_pids):
self.patient_ix = 0
return batch_2D
def copy_and_unpack_data(logger, pids, fold_dir, source_dir, target_dir):
start_time = time.time()
with open(os.path.join(fold_dir, 'file_list.txt'), 'w') as handle:
for pid in pids:
handle.write('{}.npy\n'.format(pid))
subprocess.call('rsync -ahv --files-from {} {} {}'.format(os.path.join(fold_dir, 'file_list.txt'),
source_dir, target_dir), shell=True)
# dutils.unpack_dataset(target_dir)
copied_files = os.listdir(target_dir)
logger.info("copying and unpacking data set finished : {} files in target dir: {}. took {} sec".format(
len(copied_files), target_dir, np.round(time.time() - start_time, 0)))
if __name__=="__main__":
import utils.exp_utils as utils
total_stime = time.time()
cf_file = utils.import_module("cf", "configs.py")
cf = cf_file.configs()
logger = utils.get_logger("dev")
batch_gen = get_train_generators(cf, logger)
train_batch = next(batch_gen["train"])
pids = []
total = 100
for i in range(total):
print("\r producing batch {}/{}.".format(i, total), end="", flush=True)
train_batch = next(batch_gen["train"])
pids.append(train_batch["pid"])
print()
mins, secs = divmod((time.time() - total_stime), 60)
h, mins = divmod(mins, 60)
t = "{:d}h:{:02d}m:{:02d}s".format(int(h), int(mins), int(secs))
print("{} total runtime: {}".format(os.path.split(__file__)[1], t))
\ No newline at end of file
diff --git a/experiments/toy_exp/generate_toys.py b/experiments/toy_exp/generate_toys.py
index 8f0c35b..2cdbaa6 100644
--- a/experiments/toy_exp/generate_toys.py
+++ b/experiments/toy_exp/generate_toys.py
@@ -1,130 +1,130 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import os, time
import numpy as np
import pandas as pd
import pickle
import argparse
from multiprocessing import Pool
def multi_processing_create_image(inputs):
out_dir, six, foreground_margin, class_diameters, mode, noisy_bg = inputs
print('processing {} {}'.format(out_dir, six))
img = np.random.rand(320, 320) if noisy_bg else np.zeros((320, 320))
seg = np.zeros((320, 320)).astype('uint8')
center_x = np.random.randint(foreground_margin, img.shape[0] - foreground_margin)
center_y = np.random.randint(foreground_margin, img.shape[1] - foreground_margin)
class_id = np.random.randint(0, 2)
for y in range(img.shape[0]):
for x in range(img.shape[0]):
if ((x - center_x) ** 2 + (y - center_y) ** 2 - class_diameters[class_id] ** 2) < 0:
img[y][x] += 0.2
seg[y][x] = 1
if 'donuts' in mode:
hole_diameter = 4
if class_id == 1:
for y in range(img.shape[0]):
for x in range(img.shape[0]):
if ((x - center_x) ** 2 + (y - center_y) ** 2 - hole_diameter ** 2) < 0:
img[y][x] -= 0.2
if mode == 'donuts_shape':
seg[y][x] = 0
out = np.concatenate((img[None], seg[None]))
out_path = os.path.join(out_dir, '{}.npy'.format(six))
np.save(out_path, out)
with open(os.path.join(out_dir, 'meta_info_{}.pickle'.format(six)), 'wb') as handle:
pickle.dump([out_path, class_id, str(six)], handle)
def generate_experiment(cf, exp_name, n_train_images, n_test_images, mode, class_diameters=(20, 20), noisy_bg=False):
train_dir = os.path.join(cf.root_dir, exp_name, 'train')
test_dir = os.path.join(cf.root_dir, exp_name, 'test')
os.makedirs(train_dir, exist_ok=True)
os.makedirs(test_dir, exist_ok=True)
# enforced distance between object center and image edge.
foreground_margin = int(np.ceil(np.max(class_diameters) / 1.25))
info = []
info += [[train_dir, six, foreground_margin, class_diameters, mode, noisy_bg] for six in range(n_train_images)]
info += [[test_dir, six, foreground_margin, class_diameters, mode, noisy_bg] for six in range(n_test_images)]
print('starting creation of {} images'.format(len(info)))
pool = Pool(processes=os.cpu_count()-1)
- pool.map(multi_processing_create_image, info)
+ pool.map(multi_processing_create_image, info, chunksize=1)
pool.close()
pool.join()
aggregate_meta_info(train_dir)
aggregate_meta_info(test_dir)
def aggregate_meta_info(exp_dir):
files = [os.path.join(exp_dir, f) for f in os.listdir(exp_dir) if 'meta_info' in f]
df = pd.DataFrame(columns=['path', 'class_id', 'pid'])
for f in files:
with open(f, 'rb') as handle:
df.loc[len(df)] = pickle.load(handle)
df.to_pickle(os.path.join(exp_dir, 'info_df.pickle'))
print ("aggregated meta info to df with length", len(df))
if __name__ == '__main__':
stime = time.time()
import sys
sys.path.append("../..")
import utils.exp_utils as utils
parser = argparse.ArgumentParser()
mode_choices = ['donuts_shape', 'donuts_pattern', 'circles_scale']
parser.add_argument('-m', '--modes', nargs='+', type=str, default=mode_choices, choices=mode_choices)
parser.add_argument('--noise', action='store_true', help="if given, add noise to the sample bg.")
parser.add_argument('--n_train', type=int, default=1500, help="Nr. of train images to generate.")
parser.add_argument('--n_test', type=int, default=1000, help="Nr. of test images to generate.")
args = parser.parse_args()
cf_file = utils.import_module("cf", "configs.py")
cf = cf_file.configs()
class_diameters = {
'donuts_shape': (20, 20),
'donuts_pattern': (20, 20),
'circles_scale': (19, 20)
}
for mode in args.modes:
generate_experiment(cf, mode + ("_noise" if args.noise else ""), n_train_images=args.n_train, n_test_images=args.n_test, mode=mode,
class_diameters=class_diameters[mode], noisy_bg=args.noise)
mins, secs = divmod((time.time() - stime), 60)
h, mins = divmod(mins, 60)
t = "{:d}h:{:02d}m:{:02d}s".format(int(h), int(mins), int(secs))
print("{} total runtime: {}".format(os.path.split(__file__)[1], t))
diff --git a/models/retina_net.py b/models/retina_net.py
index 52435f5..a514a91 100644
--- a/models/retina_net.py
+++ b/models/retina_net.py
@@ -1,513 +1,504 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""
Retina Net. According to https://arxiv.org/abs/1708.02002
Retina U-Net. According to https://arxiv.org/abs/1811.08661
"""
import utils.model_utils as mutils
import utils.exp_utils as utils
import sys
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils
sys.path.append('..')
from custom_extensions.nms import nms
############################################################
# Network Heads
############################################################
class Classifier(nn.Module):
def __init__(self, cf, conv):
"""
Builds the classifier sub-network.
"""
super(Classifier, self).__init__()
self.dim = conv.dim
self.n_classes = cf.head_classes
n_input_channels = cf.end_filts
n_features = cf.n_rpn_features
n_output_channels = cf.n_anchors_per_pos * cf.head_classes
anchor_stride = cf.rpn_anchor_stride
self.conv_1 = conv(n_input_channels, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_2 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_3 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_4 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_final = conv(n_features, n_output_channels, ks=3, stride=anchor_stride, pad=1, relu=None)
def forward(self, x):
"""
:param x: input feature map (b, in_c, y, x, (z))
:return: class_logits (b, n_anchors, n_classes)
"""
x = self.conv_1(x)
x = self.conv_2(x)
x = self.conv_3(x)
x = self.conv_4(x)
class_logits = self.conv_final(x)
axes = (0, 2, 3, 1) if self.dim == 2 else (0, 2, 3, 4, 1)
class_logits = class_logits.permute(*axes)
class_logits = class_logits.contiguous()
class_logits = class_logits.view(x.size()[0], -1, self.n_classes)
return [class_logits]
class BBRegressor(nn.Module):
def __init__(self, cf, conv):
"""
Builds the bb-regression sub-network.
"""
super(BBRegressor, self).__init__()
self.dim = conv.dim
n_input_channels = cf.end_filts
n_features = cf.n_rpn_features
n_output_channels = cf.n_anchors_per_pos * self.dim * 2
anchor_stride = cf.rpn_anchor_stride
self.conv_1 = conv(n_input_channels, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_2 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_3 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_4 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_final = conv(n_features, n_output_channels, ks=3, stride=anchor_stride,
pad=1, relu=None)
def forward(self, x):
"""
:param x: input feature map (b, in_c, y, x, (z))
:return: bb_logits (b, n_anchors, dim * 2)
"""
x = self.conv_1(x)
x = self.conv_2(x)
x = self.conv_3(x)
x = self.conv_4(x)
bb_logits = self.conv_final(x)
axes = (0, 2, 3, 1) if self.dim == 2 else (0, 2, 3, 4, 1)
bb_logits = bb_logits.permute(*axes)
bb_logits = bb_logits.contiguous()
bb_logits = bb_logits.view(x.size()[0], -1, self.dim * 2)
return [bb_logits]
############################################################
# Loss Functions
############################################################
def compute_class_loss(anchor_matches, class_pred_logits, shem_poolsize=20):
"""
:param anchor_matches: (n_anchors). [-1, 0, class_id] for negative, neutral, and positive matched anchors.
:param class_pred_logits: (n_anchors, n_classes). logits from classifier sub-network.
:param shem_poolsize: int. factor of top-k candidates to draw from per negative sample (online-hard-example-mining).
:return: loss: torch tensor.
:return: np_neg_ix: 1D array containing indices of the neg_roi_logits, which have been sampled for training.
"""
# Positive and Negative anchors contribute to the loss,
# but neutral anchors (match value = 0) don't.
pos_indices = torch.nonzero(anchor_matches > 0)
neg_indices = torch.nonzero(anchor_matches == -1)
- # get positive samples and calucalte loss.
- if 0 not in pos_indices.size():
+ # get positive samples and calculate loss.
+ if 0 not in pos_indices.shape:
pos_indices = pos_indices.squeeze(1)
roi_logits_pos = class_pred_logits[pos_indices]
targets_pos = anchor_matches[pos_indices]
pos_loss = F.cross_entropy(roi_logits_pos, targets_pos.long())
else:
pos_loss = torch.FloatTensor([0]).cuda()
# get negative samples, such that the amount matches the number of positive samples, but at least 1.
# get high scoring negatives by applying online-hard-example-mining.
- if 0 not in neg_indices.size():
+ if 0 not in neg_indices.shape:
neg_indices = neg_indices.squeeze(1)
roi_logits_neg = class_pred_logits[neg_indices]
- negative_count = np.max((1, pos_indices.size()[0]))
+ negative_count = np.max((1, pos_indices.shape[0]))
roi_probs_neg = F.softmax(roi_logits_neg, dim=1)
neg_ix = mutils.shem(roi_probs_neg, negative_count, shem_poolsize)
neg_loss = F.cross_entropy(roi_logits_neg[neg_ix], torch.LongTensor([0] * neg_ix.shape[0]).cuda())
# return the indices of negative samples, which contributed to the loss (for monitoring plots).
np_neg_ix = neg_ix.cpu().data.numpy()
else:
neg_loss = torch.FloatTensor([0]).cuda()
np_neg_ix = np.array([]).astype('int32')
loss = (pos_loss + neg_loss) / 2
return loss, np_neg_ix
def compute_bbox_loss(target_deltas, pred_deltas, anchor_matches):
"""
:param target_deltas: (b, n_positive_anchors, (dy, dx, (dz), log(dh), log(dw), (log(dd)))).
Uses 0 padding to fill in unsed bbox deltas.
:param pred_deltas: predicted deltas from bbox regression head. (b, n_anchors, (dy, dx, (dz), log(dh), log(dw), (log(dd))))
:param anchor_matches: (n_anchors). [-1, 0, class_id] for negative, neutral, and positive matched anchors.
:return: loss: torch 1D tensor.
"""
if 0 not in torch.nonzero(anchor_matches > 0).size():
indices = torch.nonzero(anchor_matches > 0).squeeze(1)
# Pick bbox deltas that contribute to the loss
pred_deltas = pred_deltas[indices]
# Trim target bounding box deltas to the same length as pred_deltas.
target_deltas = target_deltas[:pred_deltas.size()[0], :]
# Smooth L1 loss
loss = F.smooth_l1_loss(pred_deltas, target_deltas)
else:
loss = torch.FloatTensor([0]).cuda()
return loss
############################################################
# Output Handler
############################################################
def refine_detections(anchors, probs, deltas, batch_ixs, cf):
"""
Refine classified proposals, filter overlaps and return final
detections. n_proposals here is typically a very large number: batch_size * n_anchors.
This function is hence optimized on trimming down n_proposals.
:param anchors: (n_anchors, 2 * dim)
:param probs: (n_proposals, n_classes) softmax probabilities for all rois as predicted by classifier head.
:param deltas: (n_proposals, n_classes, 2 * dim) box refinement deltas as predicted by bbox regressor head.
:param batch_ixs: (n_proposals) batch element assignemnt info for re-allocation.
:return: result: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score))
"""
anchors = anchors.repeat(batch_ixs.unique().shape[0], 1)
# flatten foreground probabilities, sort and trim down to highest confidences by pre_nms limit.
fg_probs = probs[:, 1:].contiguous()
flat_probs, flat_probs_order = fg_probs.view(-1).sort(descending=True)
keep_ix = flat_probs_order[:cf.pre_nms_limit]
# reshape indices to 2D index array with shape like fg_probs.
keep_arr = torch.cat(((keep_ix / fg_probs.shape[1]).unsqueeze(1), (keep_ix % fg_probs.shape[1]).unsqueeze(1)), 1)
pre_nms_scores = flat_probs[:cf.pre_nms_limit]
pre_nms_class_ids = keep_arr[:, 1] + 1 # add background again.
pre_nms_batch_ixs = batch_ixs[keep_arr[:, 0]]
pre_nms_anchors = anchors[keep_arr[:, 0]]
pre_nms_deltas = deltas[keep_arr[:, 0]]
- keep = torch.arange(pre_nms_scores.size()[0]).long().cuda()
+ keep = torch.arange(pre_nms_scores.shape[0]).long().cuda()
# apply bounding box deltas. re-scale to image coordinates.
std_dev = torch.from_numpy(np.reshape(cf.rpn_bbox_std_dev, [1, cf.dim * 2])).float().cuda()
scale = torch.from_numpy(cf.scale).float().cuda()
refined_rois = mutils.apply_box_deltas_2D(pre_nms_anchors / scale, pre_nms_deltas * std_dev) * scale \
if cf.dim == 2 else mutils.apply_box_deltas_3D(pre_nms_anchors / scale, pre_nms_deltas * std_dev) * scale
# round and cast to int since we're deadling with pixels now
refined_rois = mutils.clip_to_window(cf.window, refined_rois)
pre_nms_rois = torch.round(refined_rois)
for j, b in enumerate(mutils.unique1d(pre_nms_batch_ixs)):
bixs = torch.nonzero(pre_nms_batch_ixs == b)[:, 0]
bix_class_ids = pre_nms_class_ids[bixs]
bix_rois = pre_nms_rois[bixs]
bix_scores = pre_nms_scores[bixs]
for i, class_id in enumerate(mutils.unique1d(bix_class_ids)):
ixs = torch.nonzero(bix_class_ids == class_id)[:, 0]
# nms expects boxes sorted by score.
ix_rois = bix_rois[ixs]
ix_scores = bix_scores[ixs]
ix_scores, order = ix_scores.sort(descending=True)
ix_rois = ix_rois[order, :]
- ix_scores = ix_scores
class_keep = nms.nms(ix_rois, ix_scores, cf.detection_nms_threshold)
# map indices back.
class_keep = keep[bixs[ixs[order[class_keep]]]]
# merge indices over classes for current batch element
b_keep = class_keep if i == 0 else mutils.unique1d(torch.cat((b_keep, class_keep)))
# only keep top-k boxes of current batch-element.
top_ids = pre_nms_scores[b_keep].sort(descending=True)[1][:cf.model_max_instances_per_batch_element]
b_keep = b_keep[top_ids]
# merge indices over batch elements.
batch_keep = b_keep if j == 0 else mutils.unique1d(torch.cat((batch_keep, b_keep)))
keep = batch_keep
# arrange output.
result = torch.cat((pre_nms_rois[keep],
pre_nms_batch_ixs[keep].unsqueeze(1).float(),
pre_nms_class_ids[keep].unsqueeze(1).float(),
pre_nms_scores[keep].unsqueeze(1)), dim=1)
return result
def get_results(cf, img_shape, detections, seg_logits, box_results_list=None):
"""
Restores batch dimension of merged detections, unmolds detections, creates and fills results dict.
:param img_shape:
:param detections: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score)
:param box_results_list: None or list of output boxes for monitoring/plotting.
each element is a list of boxes per batch element.
:return: results_dict: dictionary with keys:
'boxes': list over batch elements. each batch element is a list of boxes. each box is a dictionary:
[[{box_0}, ... {box_n}], [{box_0}, ... {box_n}], ...]
'seg_preds': pixel-wise class predictions (b, 1, y, x, (z)) with values [0, ..., n_classes] for
retina_unet and dummy array for retina_net.
"""
detections = detections.cpu().data.numpy()
batch_ixs = detections[:, cf.dim*2]
detections = [detections[batch_ixs == ix] for ix in range(img_shape[0])]
# for test_forward, where no previous list exists.
if box_results_list is None:
box_results_list = [[] for _ in range(img_shape[0])]
for ix in range(img_shape[0]):
if 0 not in detections[ix].shape:
boxes = detections[ix][:, :2 * cf.dim].astype(np.int32)
class_ids = detections[ix][:, 2 * cf.dim + 1].astype(np.int32)
scores = detections[ix][:, 2 * cf.dim + 2]
# Filter out detections with zero area. Often only happens in early
# stages of training when the network weights are still a bit random.
if cf.dim == 2:
exclude_ix = np.where((boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) <= 0)[0]
else:
exclude_ix = np.where(
(boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 5] - boxes[:, 4]) <= 0)[0]
if exclude_ix.shape[0] > 0:
boxes = np.delete(boxes, exclude_ix, axis=0)
class_ids = np.delete(class_ids, exclude_ix, axis=0)
scores = np.delete(scores, exclude_ix, axis=0)
if 0 not in boxes.shape:
for ix2, score in enumerate(scores):
if score >= cf.model_min_confidence:
box_results_list[ix].append({'box_coords': boxes[ix2],
'box_score': score,
'box_type': 'det',
'box_pred_class_id': class_ids[ix2]})
results_dict = {'boxes': box_results_list}
if seg_logits is None:
# output dummy segmentation for retina_net.
results_dict['seg_preds'] = np.zeros(img_shape)[:, 0][:, np.newaxis]
else:
# output label maps for retina_unet.
results_dict['seg_preds'] = F.softmax(seg_logits, 1).argmax(1).cpu().data.numpy()[:, np.newaxis].astype('uint8')
return results_dict
############################################################
# Retina (U-)Net Class
############################################################
class net(nn.Module):
def __init__(self, cf, logger):
super(net, self).__init__()
self.cf = cf
self.logger = logger
self.build()
if self.cf.weight_init is not None:
logger.info("using pytorch weight init of type {}".format(self.cf.weight_init))
mutils.initialize_weights(self)
else:
logger.info("using default pytorch weight init")
def build(self):
"""
Build Retina Net architecture.
"""
# Image size must be dividable by 2 multiple times.
h, w = self.cf.patch_size[:2]
if h / 2 ** 5 != int(h / 2 ** 5) or w / 2 ** 5 != int(w / 2 ** 5):
raise Exception("Image size must be dividable by 2 at least 5 times "
"to avoid fractions when downscaling and upscaling."
"For example, use 256, 320, 384, 448, 512, ... etc. ")
# instanciate abstract multi dimensional conv class and backbone model.
conv = mutils.NDConvGenerator(self.cf.dim)
backbone = utils.import_module('bbone', self.cf.backbone_path)
# build Anchors, FPN, Classifier / Bbox-Regressor -head
self.np_anchors = mutils.generate_pyramid_anchors(self.logger, self.cf)
self.anchors = torch.from_numpy(self.np_anchors).float().cuda()
self.Fpn = backbone.FPN(self.cf, conv, operate_stride1=self.cf.operate_stride1)
self.Classifier = Classifier(self.cf, conv)
self.BBRegressor = BBRegressor(self.cf, conv)
def train_forward(self, batch, **kwargs):
"""
train method (also used for validation monitoring). wrapper around forward pass of network. prepares input data
for processing, computes losses, and stores outputs in a dictionary.
:param batch: dictionary containing 'data', 'seg', etc.
:return: results_dict: dictionary with keys:
'boxes': list over batch elements. each batch element is a list of boxes. each box is a dictionary:
[[{box_0}, ... {box_n}], [{box_0}, ... {box_n}], ...]
'seg_preds': pixelwise segmentation output (b, c, y, x, (z)) with values [0, .., n_classes].
'monitor_values': dict of values to be monitored.
"""
img = batch['data']
gt_class_ids = batch['roi_labels']
gt_boxes = batch['bb_target']
img = torch.from_numpy(img).float().cuda()
batch_class_loss = torch.FloatTensor([0]).cuda()
batch_bbox_loss = torch.FloatTensor([0]).cuda()
# list of output boxes for monitoring/plotting. each element is a list of boxes per batch element.
box_results_list = [[] for _ in range(img.shape[0])]
detections, class_logits, pred_deltas, seg_logits = self.forward(img)
-
-
# loop over batch
for b in range(img.shape[0]):
# add gt boxes to results dict for monitoring.
if len(gt_boxes[b]) > 0:
for ix in range(len(gt_boxes[b])):
box_results_list[b].append({'box_coords': batch['bb_target'][b][ix],
'box_label': batch['roi_labels'][b][ix], 'box_type': 'gt'})
# match gt boxes with anchors to generate targets.
anchor_class_match, anchor_target_deltas = mutils.gt_anchor_matching(
self.cf, self.np_anchors, gt_boxes[b], gt_class_ids[b])
# add positive anchors used for loss to results_dict for monitoring.
pos_anchors = mutils.clip_boxes_numpy(
self.np_anchors[np.argwhere(anchor_class_match > 0)][:, 0], img.shape[2:])
for p in pos_anchors:
box_results_list[b].append({'box_coords': p, 'box_type': 'pos_anchor'})
else:
anchor_class_match = np.array([-1]*self.np_anchors.shape[0])
anchor_target_deltas = np.array([0])
anchor_class_match = torch.from_numpy(anchor_class_match).cuda()
anchor_target_deltas = torch.from_numpy(anchor_target_deltas).float().cuda()
- # todo debug print
- pos_indices = torch.nonzero(anchor_class_match > 0).squeeze(0)
- neg_indices = torch.nonzero(anchor_class_match == -1).squeeze(0)
- softmax = F.softmax(class_logits[b][pos_indices].detach(), 1)
- #ics = np.random.choice(range(softmax.shape[0]), size=min(softmax.shape[0], 6))
- comb_view = torch.cat((anchor_class_match[pos_indices].detach().unsqueeze(1).float(), softmax), dim=1)
- print(comb_view)
# compute losses.
class_loss, neg_anchor_ix = compute_class_loss(anchor_class_match, class_logits[b])
bbox_loss = compute_bbox_loss(anchor_target_deltas, pred_deltas[b], anchor_class_match)
# add negative anchors used for loss to results_dict for monitoring.
neg_anchors = mutils.clip_boxes_numpy(
self.np_anchors[np.argwhere(anchor_class_match.cpu().numpy() == -1)][neg_anchor_ix, 0],
img.shape[2:])
for n in neg_anchors:
box_results_list[b].append({'box_coords': n, 'box_type': 'neg_anchor'})
batch_class_loss += class_loss / img.shape[0]
batch_bbox_loss += bbox_loss / img.shape[0]
results_dict = get_results(self.cf, img.shape, detections, seg_logits, box_results_list)
loss = batch_class_loss + batch_bbox_loss
results_dict['torch_loss'] = loss
results_dict['class_loss'] = batch_class_loss.item()
results_dict['logger_string'] = "loss: {0:.2f}, class: {1:.2f}, bbox: {2:.2f}"\
.format(loss.item(), batch_class_loss.item(), batch_bbox_loss.item())
+
return results_dict
def test_forward(self, batch, **kwargs):
"""
test method. wrapper around forward pass of network without usage of any ground truth information.
prepares input data for processing and stores outputs in a dictionary.
:param batch: dictionary containing 'data'
:return: results_dict: dictionary with keys:
'boxes': list over batch elements. each batch element is a list of boxes. each box is a dictionary:
[[{box_0}, ... {box_n}], [{box_0}, ... {box_n}], ...]
'seg_preds': pixel-wise class predictions (b, 1, y, x, (z)) with values [0, ..., n_classes] for
retina_unet and dummy array for retina_net.
"""
img = batch['data']
img = torch.from_numpy(img).float().cuda()
detections, _, _, seg_logits = self.forward(img)
results_dict = get_results(self.cf, img.shape, detections, seg_logits)
return results_dict
def forward(self, img):
"""
forward pass of the model.
:param img: input img (b, c, y, x, (z)).
:return: rpn_pred_logits: (b, n_anchors, 2)
:return: rpn_pred_deltas: (b, n_anchors, (y, x, (z), log(h), log(w), (log(d))))
:return: batch_proposal_boxes: (b, n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ix)) only for monitoring/plotting.
:return: detections: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score)
:return: detection_masks: (n_final_detections, n_classes, y, x, (z)) raw molded masks as returned by mask-head.
"""
# Feature extraction
fpn_outs = self.Fpn(img)
seg_logits = None
selected_fmaps = [fpn_outs[i] for i in self.cf.pyramid_levels]
# Loop through pyramid layers
class_layer_outputs, bb_reg_layer_outputs = [], [] # list of lists
for p in selected_fmaps:
class_layer_outputs.append(self.Classifier(p))
bb_reg_layer_outputs.append(self.BBRegressor(p))
# Concatenate layer outputs
# Convert from list of lists of level outputs to list of lists
# of outputs across levels.
# e.g. [[a1, b1, c1], [a2, b2, c2]] => [[a1, a2], [b1, b2], [c1, c2]]
class_logits = list(zip(*class_layer_outputs))
class_logits = [torch.cat(list(o), dim=1) for o in class_logits][0]
bb_outputs = list(zip(*bb_reg_layer_outputs))
bb_outputs = [torch.cat(list(o), dim=1) for o in bb_outputs][0]
# merge batch_dimension and store info in batch_ixs for re-allocation.
batch_ixs = torch.arange(class_logits.shape[0]).unsqueeze(1).repeat(1, class_logits.shape[1]).view(-1).cuda()
flat_class_softmax = F.softmax(class_logits.view(-1, class_logits.shape[-1]), 1)
flat_bb_outputs = bb_outputs.view(-1, bb_outputs.shape[-1])
detections = refine_detections(self.anchors, flat_class_softmax, flat_bb_outputs, batch_ixs, self.cf)
return detections, class_logits, bb_outputs, seg_logits
diff --git a/plotting.py b/plotting.py
index 34b5246..023e739 100644
--- a/plotting.py
+++ b/plotting.py
@@ -1,272 +1,275 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import numpy as np
import os
from copy import deepcopy
-def plot_batch_prediction(batch, results_dict, cf, outfile= None):
+def plot_batch_prediction(batch, results_dict, cf, outfile=None, suptitle=None):
"""
plot the input images, ground truth annotations, and output predictions of a batch. If 3D batch, plots a 2D projection
of one randomly sampled element (patient) in the batch. Since plotting all slices of patient volume blows up costs of
time and space, only a section containing a randomly sampled ground truth annotation is plotted.
:param batch: dict with keys: 'data' (input image), 'seg' (pixelwise annotations), 'pid'
:param results_dict: list over batch element. Each element is a list of boxes (prediction and ground truth),
where every box is a dictionary containing box_coords, box_score and box_type.
"""
if outfile is None:
outfile = os.path.join(cf.plot_dir, 'pred_example_{}.png'.format(cf.fold))
data = batch['data']
segs = batch['seg']
pids = batch['pid']
# for 3D, repeat pid over batch elements.
if len(set(pids)) == 1:
pids = [pids] * data.shape[0]
seg_preds = results_dict['seg_preds']
roi_results = deepcopy(results_dict['boxes'])
# Randomly sampled one patient of batch and project data into 2D slices for plotting.
if cf.dim == 3:
patient_ix = np.random.choice(data.shape[0])
data = np.transpose(data[patient_ix], axes=(3, 0, 1, 2))
# select interesting foreground section to plot.
gt_boxes = [box['box_coords'] for box in roi_results[patient_ix] if box['box_type'] == 'gt']
if len(gt_boxes) > 0:
z_cuts = [np.max((int(gt_boxes[0][4]) - 5, 0)), np.min((int(gt_boxes[0][5]) + 5, data.shape[0]))]
else:
z_cuts = [data.shape[0]//2 - 5, int(data.shape[0]//2 + np.min([10, data.shape[0]//2]))]
p_roi_results = roi_results[patient_ix]
roi_results = [[] for _ in range(data.shape[0])]
# iterate over cubes and spread across slices.
for box in p_roi_results:
b = box['box_coords']
# dismiss negative anchor slices.
slices = np.round(np.unique(np.clip(np.arange(b[4], b[5] + 1), 0, data.shape[0]-1)))
for s in slices:
roi_results[int(s)].append(box)
roi_results[int(s)][-1]['box_coords'] = b[:4]
roi_results = roi_results[z_cuts[0]: z_cuts[1]]
data = data[z_cuts[0]: z_cuts[1]]
segs = np.transpose(segs[patient_ix], axes=(3, 0, 1, 2))[z_cuts[0]: z_cuts[1]]
seg_preds = np.transpose(seg_preds[patient_ix], axes=(3, 0, 1, 2))[z_cuts[0]: z_cuts[1]]
pids = [pids[patient_ix]] * data.shape[0]
try:
# all dimensions except for the 'channel-dimension' are required to match
for i in [0, 2, 3]:
assert data.shape[i] == segs.shape[i] == seg_preds.shape[i]
except:
raise Warning('Shapes of arrays to plot not in agreement!'
'Shapes {} vs. {} vs {}'.format(data.shape, segs.shape, seg_preds.shape))
show_arrays = np.concatenate([data, segs, seg_preds, data[:, 0][:, None]], axis=1).astype(float)
approx_figshape = (4 * show_arrays.shape[0], 4 * show_arrays.shape[1])
fig = plt.figure(figsize=approx_figshape)
gs = gridspec.GridSpec(show_arrays.shape[1] + 1, show_arrays.shape[0])
gs.update(wspace=0.1, hspace=0.1)
for b in range(show_arrays.shape[0]):
for m in range(show_arrays.shape[1]):
ax = plt.subplot(gs[m, b])
ax.axis('off')
if m < show_arrays.shape[1]:
arr = show_arrays[b, m]
if m < data.shape[1] or m == show_arrays.shape[1] - 1:
cmap = 'gray'
vmin = None
vmax = None
else:
cmap = None
vmin = 0
vmax = cf.num_seg_classes - 1
if m == 0:
plt.title('{}'.format(pids[b][:10]), fontsize=20)
plt.imshow(arr, cmap=cmap, vmin=vmin, vmax=vmax)
if m >= (data.shape[1]):
for box in roi_results[b]:
if box['box_type'] != 'patient_tn_box': # don't plot true negative dummy boxes.
coords = box['box_coords']
if box['box_type'] == 'det':
# dont plot background preds or low confidence boxes.
if box['box_pred_class_id'] > 0 and box['box_score'] > 0.1:
plot_text = True
score = np.max(box['box_score'])
score_text = '{}|{:.0f}'.format(box['box_pred_class_id'], score*100)
# if prob detection: plot only boxes from correct sampling instance.
if 'sample_id' in box.keys() and int(box['sample_id']) != m - data.shape[1] - 2:
continue
# if prob detection: plot reconstructed boxes only in corresponding line.
if not 'sample_id' in box.keys() and m != data.shape[1] + 1:
continue
score_font_size = 7
text_color = 'w'
text_x = coords[1] + 10*(box['box_pred_class_id'] -1) #avoid overlap of scores in plot.
text_y = coords[2] + 5
else:
continue
elif box['box_type'] == 'gt':
plot_text = True
score_text = int(box['box_label'])
score_font_size = 7
text_color = 'r'
text_x = coords[1]
text_y = coords[0] - 1
else:
plot_text = False
color_var = 'extra_usage' if 'extra_usage' in list(box.keys()) else 'box_type'
color = cf.box_color_palette[box[color_var]]
plt.plot([coords[1], coords[3]], [coords[0], coords[0]], color=color, linewidth=1, alpha=1) # up
plt.plot([coords[1], coords[3]], [coords[2], coords[2]], color=color, linewidth=1, alpha=1) # down
plt.plot([coords[1], coords[1]], [coords[0], coords[2]], color=color, linewidth=1, alpha=1) # left
plt.plot([coords[3], coords[3]], [coords[0], coords[2]], color=color, linewidth=1, alpha=1) # right
if plot_text:
plt.text(text_x, text_y, score_text, fontsize=score_font_size, color=text_color)
+ if suptitle is not None:
+ plt.suptitle(suptitle, fontsize=22)
+
try:
plt.savefig(outfile)
except:
raise Warning('failed to save plot.')
plt.close(fig)
class TrainingPlot_2Panel():
# todo remove since replaced by tensorboard?
def __init__(self, cf):
self.file_name = cf.plot_dir + '/monitor_{}'.format(cf.fold)
self.exp_name = cf.fold_dir
self.do_validation = cf.do_validation
self.separate_values_dict = cf.assign_values_to_extra_figure
self.figure_list = []
for n in range(cf.n_monitoring_figures):
self.figure_list.append(plt.figure(figsize=(10, 6)))
self.figure_list[-1].ax1 = plt.subplot(111)
self.figure_list[-1].ax1.set_xlabel('epochs')
self.figure_list[-1].ax1.set_ylabel('loss / metrics')
self.figure_list[-1].ax1.set_xlim(0, cf.num_epochs)
self.figure_list[-1].ax1.grid()
self.figure_list[0].ax1.set_ylim(0, 1.5)
self.color_palette = ['b', 'c', 'r', 'purple', 'm', 'y', 'k', 'tab:gray']
def update_and_save(self, metrics, epoch):
for figure_ix in range(len(self.figure_list)):
fig = self.figure_list[figure_ix]
detection_monitoring_plot(fig.ax1, metrics, self.exp_name, self.color_palette, epoch, figure_ix,
self.separate_values_dict,
self.do_validation)
fig.savefig(self.file_name + '_{}'.format(figure_ix))
def detection_monitoring_plot(ax1, metrics, exp_name, color_palette, epoch, figure_ix, separate_values_dict, do_validation):
# todo remove since replaced by tensorboard?
monitor_values_keys = metrics['train']['monitor_values'][1][0].keys()
separate_values = [v for fig_ix in separate_values_dict.values() for v in fig_ix]
if figure_ix == 0:
plot_keys = [ii for ii in monitor_values_keys if ii not in separate_values]
plot_keys += [k for k in metrics['train'].keys() if k != 'monitor_values']
else:
plot_keys = separate_values_dict[figure_ix]
x = np.arange(1, epoch + 1)
for kix, pk in enumerate(plot_keys):
if pk in metrics['train'].keys():
y_train = metrics['train'][pk][1:]
if do_validation:
y_val = metrics['val'][pk][1:]
else:
y_train = [np.mean([er[pk] for er in metrics['train']['monitor_values'][e]]) for e in x]
if do_validation:
y_val = [np.mean([er[pk] for er in metrics['val']['monitor_values'][e]]) for e in x]
ax1.plot(x, y_train, label='train_{}'.format(pk), linestyle='--', color=color_palette[kix])
if do_validation:
ax1.plot(x, y_val, label='val_{}'.format(pk), linestyle='-', color=color_palette[kix])
if epoch == 1:
box = ax1.get_position()
ax1.set_position([box.x0, box.y0, box.width * 0.8, box.height])
ax1.legend(loc='center left', bbox_to_anchor=(1, 0.5))
ax1.set_title(exp_name)
def plot_prediction_hist(label_list, pred_list, type_list, outfile):
"""
plot histogram of predictions for a specific class.
:param label_list: list of 1s and 0s specifying whether prediction is a true positive match (1) or a false positive (0).
False negatives (missed ground truth objects) are artificially added predictions with score 0 and label 1.
:param pred_list: list of prediction-scores.
:param type_list: list of prediction-types for stastic-info in title.
"""
preds = np.array(pred_list)
labels = np.array(label_list)
title = outfile.split('/')[-1] + ' count:{}'.format(len(label_list))
plt.figure()
plt.yscale('log')
if 0 in labels:
plt.hist(preds[labels == 0], alpha=0.3, color='g', range=(0, 1), bins=50, label='false pos.')
if 1 in labels:
plt.hist(preds[labels == 1], alpha=0.3, color='b', range=(0, 1), bins=50, label='true pos. (false neg. @ score=0)')
if type_list is not None:
fp_count = type_list.count('det_fp')
fn_count = type_list.count('det_fn')
tp_count = type_list.count('det_tp')
pos_count = fn_count + tp_count
title += ' tp:{} fp:{} fn:{} pos:{}'. format(tp_count, fp_count, fn_count, pos_count)
plt.legend()
plt.title(title)
plt.xlabel('confidence score')
plt.ylabel('log n')
plt.savefig(outfile)
plt.close()
def plot_stat_curves(stats, outfile):
for c in ['roc', 'prc']:
plt.figure()
for s in stats:
if s[c] is not None:
plt.plot(s[c][0], s[c][1], label=s['name'] + '_' + c)
plt.title(outfile.split('/')[-1] + '_' + c)
plt.legend(loc=3 if c == 'prc' else 4)
plt.xlabel('precision' if c == 'prc' else '1-spec.')
plt.ylabel('recall')
plt.savefig(outfile + '_' + c)
plt.close()
diff --git a/predictor.py b/predictor.py
index 95fa872..a023388 100644
--- a/predictor.py
+++ b/predictor.py
@@ -1,850 +1,851 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import os
import numpy as np
import torch
from scipy.stats import norm
from collections import OrderedDict
from multiprocessing import Pool
import pickle
import pandas as pd
class Predictor:
"""
Prediction pipeline:
- receives a patched patient image (n_patches, c, y, x, (z)) from patient data loader.
- forwards patches through model in chunks of batch_size. (method: batch_tiling_forward)
- unmolds predictions (boxes and segmentations) to original patient coordinates. (method: spatial_tiling_forward)
Ensembling (mode == 'test'):
- for inference, forwards 4 mirrored versions of image to through model and unmolds predictions afterwards
accordingly (method: data_aug_forward)
- for inference, loads multiple parameter-sets of the trained model corresponding to different epochs. for each
parameter-set loops over entire test set, runs prediction pipeline for each patient. (method: predict_test_set)
Consolidation of predictions:
- consolidates a patient's predictions (boxes, segmentations) collected over patches, data_aug- and temporal ensembling,
performs clustering and weighted averaging (external function: apply_wbc_to_patient) to obtain consistent outptus.
- for 2D networks, consolidates box predictions to 3D cubes via clustering (adaption of non-maximum surpression).
(external function: merge_2D_to_3D_preds_per_patient)
Ground truth handling:
- dissmisses any ground truth boxes returned by the model (happens in validation mode, patch-based groundtruth)
- if provided by data loader, adds 3D ground truth to the final predictions to be passed to the evaluator.
"""
def __init__(self, cf, net, logger, mode):
self.cf = cf
self.logger = logger
# mode is 'val' for patient-based validation/monitoring and 'test' for inference.
self.mode = mode
# model instance. In validation mode, contains parameters of current epoch.
self.net = net
# rank of current epoch loaded (for temporal averaging). this info is added to each prediction,
# for correct weighting during consolidation.
self.rank_ix = '0'
# number of ensembled models. used to calculate the number of expected predictions per position
# during consolidation of predictions. Default is 1 (no ensembling, e.g. in validation).
self.n_ens = 1
if self.mode == 'test':
try:
self.epoch_ranking = np.load(os.path.join(self.cf.fold_dir, 'epoch_ranking.npy'))[:cf.test_n_epochs]
except:
raise RuntimeError('no epoch ranking file in fold directory. '
'seems like you are trying to run testing without prior training...')
self.n_ens = cf.test_n_epochs
if self.cf.test_aug:
self.n_ens *= 4
def predict_patient(self, batch):
"""
predicts one patient.
called either directly via loop over validation set in exec.py (mode=='val')
or from self.predict_test_set (mode=='test).
in val mode: adds 3D ground truth info to predictions and runs consolidation and 2Dto3D merging of predictions.
in test mode: returns raw predictions (ground truth addition, consolidation, 2D to 3D merging are
done in self.predict_test_set, because patient predictions across several epochs might be needed
to be collected first, in case of temporal ensembling).
:return. results_dict: stores the results for one patient. dictionary with keys:
- 'boxes': list over batch elements. each element is a list over boxes, where each box is
one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions
(if not merged to 3D), and a dummy batch dimension of 1 for 3D predictions.
- 'seg_preds': pixel-wise predictions. (b, 1, y, x, (z))
- losses (only in validation mode)
"""
- self.logger.info('evaluating patient {} for fold {} '.format(batch['pid'], self.cf.fold))
+ #self.logger.info('\revaluating patient {} for fold {} '.format(batch['pid'], self.cf.fold))
+ print('\revaluating patient {} for fold {} '.format(batch['pid'], self.cf.fold), end="", flush=True)
# True if patient is provided in patches and predictions need to be tiled.
self.patched_patient = True if 'patch_crop_coords' in list(batch.keys()) else False
# forward batch through prediction pipeline.
results_dict = self.data_aug_forward(batch)
if self.mode == 'val':
for b in range(batch['patient_bb_target'].shape[0]):
for t in range(len(batch['patient_bb_target'][b])):
results_dict['boxes'][b].append({'box_coords': batch['patient_bb_target'][b][t],
'box_label': batch['patient_roi_labels'][b][t],
'box_type': 'gt'})
if self.patched_patient:
wcs_input = [results_dict['boxes'], 'dummy_pid', self.cf.class_dict, self.cf.wcs_iou, self.n_ens]
results_dict['boxes'] = apply_wbc_to_patient(wcs_input)[0]
if self.cf.merge_2D_to_3D_preds:
merge_dims_inputs = [results_dict['boxes'], 'dummy_pid', self.cf.class_dict, self.cf.merge_3D_iou]
results_dict['boxes'] = merge_2D_to_3D_preds_per_patient(merge_dims_inputs)[0]
return results_dict
def predict_test_set(self, batch_gen, return_results=True):
"""
wrapper around test method, which loads multiple (or one) epoch parameters (temporal ensembling), loops through
the test set and collects predictions per patient. Also flattens the results per patient and epoch
and adds optional ground truth boxes for evaluation. Saves out the raw result list for later analysis and
optionally consolidates and returns predictions immediately.
:return: (optionally) list_of_results_per_patient: list over patient results. each entry is a dict with keys:
- 'boxes': list over batch elements. each element is a list over boxes, where each box is
one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions
(if not merged to 3D), and a dummy batch dimension of 1 for 3D predictions.
- 'seg_preds': not implemented yet. todo for evaluation of instance/semantic segmentation.
"""
dict_of_patient_results = OrderedDict()
# get paths of all parameter sets to be loaded for temporal ensembling. (or just one for no temp. ensembling).
weight_paths = [os.path.join(self.cf.fold_dir, '{}_best_checkpoint'.format(epoch), 'params.pth') for epoch in
self.epoch_ranking]
for rank_ix, weight_path in enumerate(weight_paths):
self.logger.info(('tmp ensembling over rank_ix:{} epoch:{}'.format(rank_ix, weight_path)))
self.net.load_state_dict(torch.load(weight_path))
self.net.eval()
self.rank_ix = str(rank_ix) # get string of current rank for unique patch ids.
with torch.no_grad():
for _ in range(batch_gen['n_test']):
batch = next(batch_gen['test'])
# store batch info in patient entry of results dict.
if rank_ix == 0:
dict_of_patient_results[batch['pid']] = {}
dict_of_patient_results[batch['pid']]['results_list'] = []
dict_of_patient_results[batch['pid']]['patient_bb_target'] = batch['patient_bb_target']
dict_of_patient_results[batch['pid']]['patient_roi_labels'] = batch['patient_roi_labels']
# call prediction pipeline and store results in dict.
results_dict = self.predict_patient(batch)
dict_of_patient_results[batch['pid']]['results_list'].append({"boxes": results_dict['boxes']})
self.logger.info('finished predicting test set. starting post-processing of predictions.')
results_per_patient = []
# loop over patients again to flatten results across epoch predictions.
# if provided, add ground truth boxes for evaluation.
for pid, p_dict in dict_of_patient_results.items():
tmp_ens_list = p_dict['results_list']
results_dict = {}
# collect all boxes/seg_preds of same batch_instance over temporal instances.
b_size = len(tmp_ens_list[0])
results_dict['boxes'] = [[item for rank_dict in tmp_ens_list for item in rank_dict["boxes"][batch_instance]]
for batch_instance in range(b_size)]
# TODO return for instance segmentation:
# results_dict['seg_preds'] = np.mean(results_dict['seg_preds'], 1)[:, None]
# results_dict['seg_preds'] = np.array([[item for d in tmp_ens_list for item in d['seg_preds'][batch_instance]]
# for batch_instance in range(len(tmp_ens_list[0]['boxes']))])
# add 3D ground truth boxes for evaluation.
for b in range(p_dict['patient_bb_target'].shape[0]):
for t in range(len(p_dict['patient_bb_target'][b])):
results_dict['boxes'][b].append({'box_coords': p_dict['patient_bb_target'][b][t],
'box_label': p_dict['patient_roi_labels'][b][t],
'box_type': 'gt'})
results_per_patient.append([results_dict, pid])
# save out raw predictions.
out_string = 'raw_pred_boxes_hold_out_list' if self.cf.hold_out_test_set else 'raw_pred_boxes_list'
with open(os.path.join(self.cf.fold_dir, '{}.pickle'.format(out_string)), 'wb') as handle:
pickle.dump(results_per_patient, handle)
if return_results:
final_patient_box_results = [(res_dict["boxes"], pid) for res_dict, pid in results_per_patient]
# consolidate predictions.
self.logger.info('applying wcs to test set predictions with iou = {} and n_ens = {}.'.format(
self.cf.wcs_iou, self.n_ens))
pool = Pool(processes=6)
mp_inputs = [[ii[0], ii[1], self.cf.class_dict, self.cf.wcs_iou, self.n_ens] for ii in final_patient_box_results]
final_patient_box_results = pool.map(apply_wbc_to_patient, mp_inputs, chunksize=1)
pool.close()
pool.join()
# merge 2D boxes to 3D cubes. (if model predicts 2D but evaluation is run in 3D)
if self.cf.merge_2D_to_3D_preds:
self.logger.info('applying 2Dto3D merging to test set predictions with iou = {}.'.format(self.cf.merge_3D_iou))
pool = Pool(processes=6)
mp_inputs = [[ii[0], ii[1], self.cf.class_dict, self.cf.merge_3D_iou] for ii in final_patient_box_results]
final_patient_box_results = pool.map(merge_2D_to_3D_preds_per_patient, mp_inputs, chunksize=1)
pool.close()
pool.join()
# final_patient_box_results holds [avg_boxes, pid] if wbc
for ix in range(len(results_per_patient)):
assert results_per_patient[ix][1] == final_patient_box_results[ix][1], "should be same pid"
results_per_patient[ix][0]["boxes"] = final_patient_box_results[ix][0]
return results_per_patient
def load_saved_predictions(self, apply_wbc=False):
"""
loads raw predictions saved by self.predict_test_set. consolidates and merges 2D boxes to 3D cubes for evaluation.
(if model predicts 2D but evaluation is run in 3D)
:return: (optionally) results_list: list over patient results. each entry is a dict with keys:
- 'boxes': list over batch elements. each element is a list over boxes, where each box is
one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions
(if not merged to 3D), and a dummy batch dimension of 1 for 3D predictions.
- 'seg_preds': not implemented yet. todo for evaluation of instance/semantic segmentation.
"""
# load predictions for a single test-set fold.
if not self.cf.hold_out_test_set:
with open(os.path.join(self.cf.fold_dir, 'raw_pred_boxes_list.pickle'), 'rb') as handle:
results_list = pickle.load(handle)
box_results_list = [(res_dict["boxes"], pid) for res_dict, pid in results_list]
da_factor = 4 if self.cf.test_aug else 1
n_ens = self.cf.test_n_epochs * da_factor
self.logger.info('loaded raw test set predictions with n_patients = {} and n_ens = {}'.format(
len(results_list), n_ens))
# if hold out test set was perdicted, aggregate predictions of all trained models
# corresponding to all CV-folds and flatten them.
else:
self.logger.info("loading saved predictions of hold-out test set")
fold_dirs = sorted([os.path.join(self.cf.exp_dir, f) for f in os.listdir(self.cf.exp_dir) if
os.path.isdir(os.path.join(self.cf.exp_dir, f)) and f.startswith("fold")])
results_list = []
folds_loaded = 0
for fold in range(self.cf.n_cv_splits):
fold_dir = os.path.join(self.cf.exp_dir, 'fold_{}'.format(fold))
if fold_dir in fold_dirs:
with open(os.path.join(fold_dir, 'raw_pred_boxes_hold_out_list.pickle'), 'rb') as handle:
fold_list = pickle.load(handle)
results_list += fold_list
folds_loaded += 1
else:
self.logger.info("Skipping fold {} since no saved predictions found.".format(fold))
box_results_list = []
for res_dict, pid in results_list: #without filtering gt out:
box_results_list.append((res_dict['boxes'], pid))
da_factor = 4 if self.cf.test_aug else 1
n_ens = self.cf.test_n_epochs * da_factor * folds_loaded
# consolidate predictions.
if apply_wbc:
self.logger.info('applying wcs to test set predictions with iou = {} and n_ens = {}.'.format(
self.cf.wcs_iou, n_ens))
pool = Pool(processes=6)
mp_inputs = [[ii[0], ii[1], self.cf.class_dict, self.cf.wcs_iou, n_ens] for ii in box_results_list]
box_results_list = pool.map(apply_wbc_to_patient, mp_inputs, chunksize=1)
pool.close()
pool.join()
# merge 2D box predictions to 3D cubes (if model predicts 2D but evaluation is run in 3D)
if self.cf.merge_2D_to_3D_preds:
self.logger.info(
'applying 2Dto3D merging to test set predictions with iou = {}.'.format(self.cf.merge_3D_iou))
pool = Pool(processes=6)
mp_inputs = [[ii[0], ii[1], self.cf.class_dict, self.cf.merge_3D_iou] for ii in box_results_list]
box_results_list = pool.map(merge_2D_to_3D_preds_per_patient, mp_inputs, chunksize=1)
pool.close()
pool.join()
for ix in range(len(results_list)):
assert np.all(
results_list[ix][1] == box_results_list[ix][1]), "pid mismatch between loaded and aggregated results"
results_list[ix][0]["boxes"] = box_results_list[ix][0]
return results_list # holds (results_dict, pid)
def data_aug_forward(self, batch):
"""
in val_mode: passes batch through to spatial_tiling method without data_aug.
in test_mode: if cf.test_aug is set in configs, createst 4 mirrored versions of the input image,
passes all of them to the next processing step (spatial_tiling method) and re-transforms returned predictions
to original image version.
:return. results_dict: stores the results for one patient. dictionary with keys:
- 'boxes': list over batch elements. each element is a list over boxes, where each box is
one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions,
and a dummy batch dimension of 1 for 3D predictions.
- 'seg_preds': pixel-wise predictions. (b, 1, y, x, (z))
- losses (only in validation mode)
"""
patch_crops = batch['patch_crop_coords'] if self.patched_patient else None
results_list = [self.spatial_tiling_forward(batch, patch_crops)]
org_img_shape = batch['original_img_shape']
if self.mode == 'test' and self.cf.test_aug:
if self.patched_patient:
# apply mirror transformations to patch-crop coordinates, for correct tiling in spatial_tiling method.
mirrored_patch_crops = get_mirrored_patch_crops(patch_crops, batch['original_img_shape'])
else:
mirrored_patch_crops = [None] * 3
img = np.copy(batch['data'])
# first mirroring: y-axis.
batch['data'] = np.flip(img, axis=2).copy()
chunk_dict = self.spatial_tiling_forward(batch, mirrored_patch_crops[0], n_aug='1')
# re-transform coordinates.
for ix in range(len(chunk_dict['boxes'])):
for boxix in range(len(chunk_dict['boxes'][ix])):
coords = chunk_dict['boxes'][ix][boxix]['box_coords'].copy()
coords[0] = org_img_shape[2] - chunk_dict['boxes'][ix][boxix]['box_coords'][2]
coords[2] = org_img_shape[2] - chunk_dict['boxes'][ix][boxix]['box_coords'][0]
assert coords[2] >= coords[0], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]
assert coords[3] >= coords[1], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]
chunk_dict['boxes'][ix][boxix]['box_coords'] = coords
# re-transform segmentation predictions.
chunk_dict['seg_preds'] = np.flip(chunk_dict['seg_preds'], axis=2)
results_list.append(chunk_dict)
# second mirroring: x-axis.
batch['data'] = np.flip(img, axis=3).copy()
chunk_dict = self.spatial_tiling_forward(batch, mirrored_patch_crops[1], n_aug='2')
# re-transform coordinates.
for ix in range(len(chunk_dict['boxes'])):
for boxix in range(len(chunk_dict['boxes'][ix])):
coords = chunk_dict['boxes'][ix][boxix]['box_coords'].copy()
coords[1] = org_img_shape[3] - chunk_dict['boxes'][ix][boxix]['box_coords'][3]
coords[3] = org_img_shape[3] - chunk_dict['boxes'][ix][boxix]['box_coords'][1]
assert coords[2] >= coords[0], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]
assert coords[3] >= coords[1], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]
chunk_dict['boxes'][ix][boxix]['box_coords'] = coords
# re-transform segmentation predictions.
chunk_dict['seg_preds'] = np.flip(chunk_dict['seg_preds'], axis=3)
results_list.append(chunk_dict)
# third mirroring: y-axis and x-axis.
batch['data'] = np.flip(np.flip(img, axis=2), axis=3).copy()
chunk_dict = self.spatial_tiling_forward(batch, mirrored_patch_crops[2], n_aug='3')
# re-transform coordinates.
for ix in range(len(chunk_dict['boxes'])):
for boxix in range(len(chunk_dict['boxes'][ix])):
coords = chunk_dict['boxes'][ix][boxix]['box_coords'].copy()
coords[0] = org_img_shape[2] - chunk_dict['boxes'][ix][boxix]['box_coords'][2]
coords[2] = org_img_shape[2] - chunk_dict['boxes'][ix][boxix]['box_coords'][0]
coords[1] = org_img_shape[3] - chunk_dict['boxes'][ix][boxix]['box_coords'][3]
coords[3] = org_img_shape[3] - chunk_dict['boxes'][ix][boxix]['box_coords'][1]
assert coords[2] >= coords[0], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]
assert coords[3] >= coords[1], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]
chunk_dict['boxes'][ix][boxix]['box_coords'] = coords
# re-transform segmentation predictions.
chunk_dict['seg_preds'] = np.flip(np.flip(chunk_dict['seg_preds'], axis=2), axis=3).copy()
results_list.append(chunk_dict)
batch['data'] = img
# aggregate all boxes/seg_preds per batch element from data_aug predictions.
results_dict = {}
results_dict['boxes'] = [[item for d in results_list for item in d['boxes'][batch_instance]]
for batch_instance in range(org_img_shape[0])]
results_dict['seg_preds'] = np.array([[item for d in results_list for item in d['seg_preds'][batch_instance]]
for batch_instance in range(org_img_shape[0])])
if self.mode == 'val':
try:
results_dict['torch_loss'] = results_list[0]['torch_loss']
results_dict['class_loss'] = results_list[0]['class_loss']
except KeyError:
pass
return results_dict
def spatial_tiling_forward(self, batch, patch_crops=None, n_aug='0'):
"""
forwards batch to batch_tiling_forward method and receives and returns a dictionary with results.
if patch-based prediction, the results received from batch_tiling_forward will be on a per-patch-basis.
this method uses the provided patch_crops to re-transform all predictions to whole-image coordinates.
Patch-origin information of all box-predictions will be needed for consolidation, hence it is stored as
'patch_id', which is a unique string for each patch (also takes current data aug and temporal epoch instances
into account). all box predictions get additional information about the amount overlapping patches at the
respective position (used for consolidation).
:return. results_dict: stores the results for one patient. dictionary with keys:
- 'boxes': list over batch elements. each element is a list over boxes, where each box is
one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions,
and a dummy batch dimension of 1 for 3D predictions.
- 'seg_preds': pixel-wise predictions. (b, 1, y, x, (z))
- losses (only in validation mode)
"""
if patch_crops is not None:
patches_dict = self.batch_tiling_forward(batch)
results_dict = {'boxes': [[] for _ in range(batch['original_img_shape'][0])]}
# instanciate segemntation output array. Will contain averages over patch predictions.
out_seg_preds = np.zeros(batch['original_img_shape'], dtype=np.float16)[:, 0][:, None]
# counts patch instances per pixel-position.
patch_overlap_map = np.zeros_like(out_seg_preds, dtype='uint8')
#unmold segmentation outputs. loop over patches.
for pix, pc in enumerate(patch_crops):
if self.cf.dim == 3:
out_seg_preds[:, :, pc[0]:pc[1], pc[2]:pc[3], pc[4]:pc[5]] += patches_dict['seg_preds'][pix][None]
patch_overlap_map[:, :, pc[0]:pc[1], pc[2]:pc[3], pc[4]:pc[5]] += 1
else:
out_seg_preds[pc[4]:pc[5], :, pc[0]:pc[1], pc[2]:pc[3], ] += patches_dict['seg_preds'][pix]
patch_overlap_map[pc[4]:pc[5], :, pc[0]:pc[1], pc[2]:pc[3], ] += 1
# take mean in overlapping areas.
out_seg_preds[patch_overlap_map > 0] /= patch_overlap_map[patch_overlap_map > 0]
results_dict['seg_preds'] = out_seg_preds
# unmold box outputs. loop over patches.
for pix, pc in enumerate(patch_crops):
patch_boxes = patches_dict['boxes'][pix]
for box in patch_boxes:
# add unique patch id for consolidation of predictions.
box['patch_id'] = self.rank_ix + '_' + n_aug + '_' + str(pix)
# boxes from the edges of a patch have a lower prediction quality, than the ones at patch-centers.
# hence they will be downweighted for consolidation, using the 'box_patch_center_factor', which is
# obtained by a normal distribution over positions in the patch and average over spatial dimensions.
# Also the info 'box_n_overlaps' is stored for consolidation, which depicts the amount over
# overlapping patches at the box's position.
c = box['box_coords']
box_centers = [(c[ii] + c[ii + 2]) / 2 for ii in range(2)]
if self.cf.dim == 3:
box_centers.append((c[4] + c[5]) / 2)
box['box_patch_center_factor'] = np.mean(
[norm.pdf(bc, loc=pc, scale=pc * 0.8) * np.sqrt(2 * np.pi) * pc * 0.8 for bc, pc in
zip(box_centers, np.array(self.cf.patch_size) / 2)])
if self.cf.dim == 3:
c += np.array([pc[0], pc[2], pc[0], pc[2], pc[4], pc[4]])
int_c = [int(np.floor(ii)) if ix%2 == 0 else int(np.ceil(ii)) for ix, ii in enumerate(c)]
box['box_n_overlaps'] = np.mean(patch_overlap_map[:, :, int_c[1]:int_c[3], int_c[0]:int_c[2], int_c[4]:int_c[5]])
results_dict['boxes'][0].append(box)
else:
c += np.array([pc[0], pc[2], pc[0], pc[2]])
int_c = [int(np.floor(ii)) if ix % 2 == 0 else int(np.ceil(ii)) for ix, ii in enumerate(c)]
box['box_n_overlaps'] = np.mean(patch_overlap_map[pc[4], :, int_c[1]:int_c[3], int_c[0]:int_c[2]])
results_dict['boxes'][pc[4]].append(box)
if self.mode == 'val':
try:
results_dict['torch_loss'] = patches_dict['torch_loss']
results_dict['class_loss'] = patches_dict['class_loss']
except KeyError:
pass
# if predictions are not patch-based:
# add patch-origin info to boxes (entire image is the same patch with overlap=1) and return results.
else:
results_dict = self.batch_tiling_forward(batch)
for b in results_dict['boxes']:
for box in b:
box['box_patch_center_factor'] = 1
box['box_n_overlaps'] = 1
box['patch_id'] = self.rank_ix + '_' + n_aug
return results_dict
def batch_tiling_forward(self, batch):
"""
calls the actual network forward method. in patch-based prediction, the batch dimension might be overladed
with n_patches >> batch_size, which would exceed gpu memory. In this case, batches are processed in chunks of
batch_size. validation mode calls the train method to monitor losses (returned ground truth objects are discarded).
test mode calls the test forward method, no ground truth required / involved.
:return. results_dict: stores the results for one patient. dictionary with keys:
- 'boxes': list over batch elements. each element is a list over boxes, where each box is
one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions,
and a dummy batch dimension of 1 for 3D predictions.
- 'seg_preds': pixel-wise predictions. (b, 1, y, x, (z))
- losses (only in validation mode)
"""
- self.logger.info('forwarding (patched) patient with shape: {}'.format(batch['data'].shape))
+ #self.logger.info('forwarding (patched) patient with shape: {}'.format(batch['data'].shape))
img = batch['data']
if img.shape[0] <= self.cf.batch_size:
if self.mode == 'val':
# call training method to monitor losses
results_dict = self.net.train_forward(batch, is_validation=True)
# discard returned ground-truth boxes (also training info boxes).
results_dict['boxes'] = [[box for box in b if box['box_type'] == 'det'] for b in results_dict['boxes']]
else:
results_dict = self.net.test_forward(batch, return_masks=self.cf.return_masks_in_test)
else:
split_ixs = np.split(np.arange(img.shape[0]), np.arange(img.shape[0])[::self.cf.batch_size])
chunk_dicts = []
for chunk_ixs in split_ixs[1:]: # first split is elements before 0, so empty
b = {k: batch[k][chunk_ixs] for k in batch.keys()
if (isinstance(batch[k], np.ndarray) and batch[k].shape[0] == img.shape[0])}
if self.mode == 'val':
chunk_dicts += [self.net.train_forward(b, is_validation=True)]
else:
chunk_dicts += [self.net.test_forward(b, return_masks=self.cf.return_masks_in_test)]
results_dict = {}
# flatten out batch elements from chunks ([chunk, chunk] -> [b, b, b, b, ...])
results_dict['boxes'] = [item for d in chunk_dicts for item in d['boxes']]
results_dict['seg_preds'] = np.array([item for d in chunk_dicts for item in d['seg_preds']])
if self.mode == 'val':
try:
# estimate metrics by mean over batch_chunks. Most similar to training metrics.
results_dict['torch_loss'] = torch.mean(torch.cat([d['torch_loss'] for d in chunk_dicts]))
results_dict['class_loss'] = np.mean([d['class_loss'] for d in chunk_dicts])
except KeyError:
# losses are not necessarily monitored
pass
# discard returned ground-truth boxes (also training info boxes).
results_dict['boxes'] = [[box for box in b if box['box_type'] == 'det'] for b in results_dict['boxes']]
return results_dict
def apply_wbc_to_patient(inputs):
"""
wrapper around prediction box consolidation: weighted cluster scoring (wcs). processes a single patient.
loops over batch elements in patient results (1 in 3D, slices in 2D) and foreground classes,
aggregates and stores results in new list.
:return. patient_results_list: list over batch elements. each element is a list over boxes, where each box is
one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D
predictions, and a dummy batch dimension of 1 for 3D predictions.
:return. pid: string. patient id.
"""
in_patient_results_list, pid, class_dict, wcs_iou, n_ens = inputs
out_patient_results_list = [[] for _ in range(len(in_patient_results_list))]
for bix, b in enumerate(in_patient_results_list):
for cl in list(class_dict.keys()):
boxes = [(ix, box) for ix, box in enumerate(b) if (box['box_type'] == 'det' and box['box_pred_class_id'] == cl)]
box_coords = np.array([b[1]['box_coords'] for b in boxes])
box_scores = np.array([b[1]['box_score'] for b in boxes])
box_center_factor = np.array([b[1]['box_patch_center_factor'] for b in boxes])
box_n_overlaps = np.array([b[1]['box_n_overlaps'] for b in boxes])
box_patch_id = np.array([b[1]['patch_id'] for b in boxes])
if 0 not in box_scores.shape:
keep_scores, keep_coords = weighted_box_clustering(
np.concatenate((box_coords, box_scores[:, None], box_center_factor[:, None],
box_n_overlaps[:, None]), axis=1), box_patch_id, wcs_iou, n_ens)
for boxix in range(len(keep_scores)):
out_patient_results_list[bix].append({'box_type': 'det', 'box_coords': keep_coords[boxix],
'box_score': keep_scores[boxix], 'box_pred_class_id': cl})
# add gt boxes back to new output list.
out_patient_results_list[bix].extend([box for box in b if box['box_type'] == 'gt'])
return [out_patient_results_list, pid]
def merge_2D_to_3D_preds_per_patient(inputs):
"""
wrapper around 2Dto3D merging operation. Processes a single patient. Takes 2D patient results (slices in batch dimension)
and returns 3D patient results (dummy batch dimension of 1). Applies an adaption of Non-Maximum Surpression
(Detailed methodology is described in nms_2to3D).
:return. results_dict_boxes: list over batch elements (1 in 3D). each element is a list over boxes, where each box is
one dictionary: [[box_0, ...], [box_n,...]].
:return. pid: string. patient id.
"""
in_patient_results_list, pid, class_dict, merge_3D_iou = inputs
out_patient_results_list = []
for cl in list(class_dict.keys()):
boxes, slice_ids = [], []
# collect box predictions over batch dimension (slices) and store slice info as slice_ids.
for bix, b in enumerate(in_patient_results_list):
det_boxes = [(ix, box) for ix, box in enumerate(b) if
(box['box_type'] == 'det' and box['box_pred_class_id'] == cl)]
boxes += det_boxes
slice_ids += [bix] * len(det_boxes)
box_coords = np.array([b[1]['box_coords'] for b in boxes])
box_scores = np.array([b[1]['box_score'] for b in boxes])
slice_ids = np.array(slice_ids)
if 0 not in box_scores.shape:
keep_ix, keep_z = nms_2to3D(
np.concatenate((box_coords, box_scores[:, None], slice_ids[:, None]), axis=1), merge_3D_iou)
else:
keep_ix, keep_z = [], []
# store kept predictions in new results list and add corresponding z-dimension info to coordinates.
for kix, kz in zip(keep_ix, keep_z):
out_patient_results_list.append({'box_type': 'det', 'box_coords': list(box_coords[kix]) + kz,
'box_score': box_scores[kix], 'box_pred_class_id': cl})
out_patient_results_list += [box for b in in_patient_results_list for box in b if box['box_type'] == 'gt']
out_patient_results_list = [out_patient_results_list] # add dummy batch dimension 1 for 3D.
return [out_patient_results_list, pid]
def weighted_box_clustering(dets, box_patch_id, thresh, n_ens):
"""
consolidates overlapping predictions resulting from patch overlaps, test data augmentations and temporal ensembling.
clusters predictions together with iou > thresh (like in NMS). Output score and coordinate for one cluster are the
average weighted by individual patch center factors (how trustworthy is this candidate measured by how centered
its position the patch is) and the size of the corresponding box.
The number of expected predictions at a position is n_data_aug * n_temp_ens * n_overlaps_at_position
(1 prediction per unique patch). Missing predictions at a cluster position are defined as the number of unique
patches in the cluster, which did not contribute any predict any boxes.
:param dets: (n_dets, (y1, x1, y2, x2, (z1), (z2), scores, box_pc_facts, box_n_ovs)
:param thresh: threshold for iou_matching.
:param n_ens: number of models, that are ensembled. (-> number of expected predicitions per position)
:return: keep_scores: (n_keep) new scores of boxes to be kept.
:return: keep_coords: (n_keep, (y1, x1, y2, x2, (z1), (z2)) new coordinates of boxes to be kept.
"""
dim = 2 if dets.shape[1] == 7 else 3
y1 = dets[:, 0]
x1 = dets[:, 1]
y2 = dets[:, 2]
x2 = dets[:, 3]
scores = dets[:, -3]
box_pc_facts = dets[:, -2]
box_n_ovs = dets[:, -1]
areas = (y2 - y1 + 1) * (x2 - x1 + 1)
if dim == 3:
z1 = dets[:, 4]
z2 = dets[:, 5]
areas *= (z2 - z1 + 1)
# order is the sorted index. maps order to index o[1] = 24 (rank1, ix 24)
order = scores.argsort()[::-1]
keep = []
keep_scores = []
keep_coords = []
while order.size > 0:
i = order[0] # higehst scoring element
xx1 = np.maximum(x1[i], x1[order])
yy1 = np.maximum(y1[i], y1[order])
xx2 = np.minimum(x2[i], x2[order])
yy2 = np.minimum(y2[i], y2[order])
w = np.maximum(0.0, xx2 - xx1 + 1)
h = np.maximum(0.0, yy2 - yy1 + 1)
inter = w * h
if dim == 3:
zz1 = np.maximum(z1[i], z1[order])
zz2 = np.minimum(z2[i], z2[order])
d = np.maximum(0.0, zz2 - zz1 + 1)
inter *= d
# overall between currently highest scoring box and all boxes.
ovr = inter / (areas[i] + areas[order] - inter)
# get all the predictions that match the current box to build one cluster.
matches = np.argwhere(ovr > thresh)
match_n_ovs = box_n_ovs[order[matches]]
match_pc_facts = box_pc_facts[order[matches]]
match_patch_id = box_patch_id[order[matches]]
match_ov_facts = ovr[matches]
match_areas = areas[order[matches]]
match_scores = scores[order[matches]]
# weight all socres in cluster by patch factors, and size.
match_score_weights = match_ov_facts * match_areas * match_pc_facts
match_scores *= match_score_weights
# for the weigted average, scores have to be divided by the number of total expected preds at the position
# of the current cluster. 1 Prediction per patch is expected. therefore, the number of ensembled models is
# multiplied by the mean overlaps of patches at this position (boxes of the cluster might partly be
# in areas of different overlaps).
n_expected_preds = n_ens * np.mean(match_n_ovs)
# the number of missing predictions is obtained as the number of patches,
# which did not contribute any prediction to the current cluster.
n_missing_preds = np.max((0, n_expected_preds - np.unique(match_patch_id).shape[0]))
# missing preds are given the mean weighting
# (expected prediction is the mean over all predictions in cluster).
denom = np.sum(match_score_weights) + n_missing_preds * np.mean(match_score_weights)
# compute weighted average score for the cluster
avg_score = np.sum(match_scores) / denom
# compute weighted average of coordinates for the cluster. now only take existing
# predictions into account.
avg_coords = [np.sum(y1[order[matches]] * match_scores) / np.sum(match_scores),
np.sum(x1[order[matches]] * match_scores) / np.sum(match_scores),
np.sum(y2[order[matches]] * match_scores) / np.sum(match_scores),
np.sum(x2[order[matches]] * match_scores) / np.sum(match_scores)]
if dim == 3:
avg_coords.append(np.sum(z1[order[matches]] * match_scores) / np.sum(match_scores))
avg_coords.append(np.sum(z2[order[matches]] * match_scores) / np.sum(match_scores))
# some clusters might have very low scores due to high amounts of missing predictions.
# filter out the with a conservative threshold, to speed up evaluation.
if avg_score > 0.01:
keep_scores.append(avg_score)
keep_coords.append(avg_coords)
# get index of all elements that were not matched and discard all others.
inds = np.where(ovr <= thresh)[0]
order = order[inds]
return keep_scores, keep_coords
def nms_2to3D(dets, thresh):
"""
Merges 2D boxes to 3D cubes. Therefore, boxes of all slices are projected into one slices. An adaptation of Non-maximum surpression
is applied, where clusters are found (like in NMS) with an extra constrained, that surpressed boxes have to have 'connected'
z-coordinates w.r.t the core slice (cluster center, highest scoring box). 'connected' z-coordinates are determined
as the z-coordinates with predictions until the first coordinate, where no prediction was found.
example: a cluster of predictions was found overlap > iou thresh in xy (like NMS). The z-coordinate of the highest
scoring box is 50. Other predictions have 23, 46, 48, 49, 51, 52, 53, 56, 57.
Only the coordinates connected with 50 are clustered to one cube: 48, 49, 51, 52, 53. (46 not because nothing was
found in 47, so 47 is a 'hole', which interrupts the connection). Only the boxes corresponding to these coordinates
are surpressed. All others are kept for building of further clusters.
This algorithm works better with a certain min_confidence of predictions, because low confidence (e.g. noisy/cluttery)
predictions can break the relatively strong assumption of defining cubes' z-boundaries at the first 'hole' in the cluster.
:param dets: (n_detections, (y1, x1, y2, x2, scores, slice_id)
:param thresh: iou matchin threshold (like in NMS).
:return: keep: (n_keep) 1D tensor of indices to be kept.
:return: keep_z: (n_keep, [z1, z2]) z-coordinates to be added to boxes, which are kept in order to form cubes.
"""
y1 = dets[:, 0]
x1 = dets[:, 1]
y2 = dets[:, 2]
x2 = dets[:, 3]
scores = dets[:, -2]
slice_id = dets[:, -1]
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
order = scores.argsort()[::-1]
keep = []
keep_z = []
while order.size > 0: # order is the sorted index. maps order to index o[1] = 24 (rank1, ix 24)
i = order[0] # pop higehst scoring element
xx1 = np.maximum(x1[i], x1[order])
yy1 = np.maximum(y1[i], y1[order])
xx2 = np.minimum(x2[i], x2[order])
yy2 = np.minimum(y2[i], y2[order])
w = np.maximum(0.0, xx2 - xx1 + 1)
h = np.maximum(0.0, yy2 - yy1 + 1)
inter = w * h
ovr = inter / (areas[i] + areas[order] - inter)
matches = np.argwhere(ovr > thresh) # get all the elements that match the current box and have a lower score
slice_ids = slice_id[order[matches]]
core_slice = slice_id[int(i)]
upper_wholes = [ii for ii in np.arange(core_slice, np.max(slice_ids)) if ii not in slice_ids]
lower_wholes = [ii for ii in np.arange(np.min(slice_ids), core_slice) if ii not in slice_ids]
max_valid_slice_id = np.min(upper_wholes) if len(upper_wholes) > 0 else np.max(slice_ids)
min_valid_slice_id = np.max(lower_wholes) if len(lower_wholes) > 0 else np.min(slice_ids)
z_matches = matches[(slice_ids <= max_valid_slice_id) & (slice_ids >= min_valid_slice_id)]
z1 = np.min(slice_id[order[z_matches]]) - 1
z2 = np.max(slice_id[order[z_matches]]) + 1
keep.append(i)
keep_z.append([z1, z2])
order = np.delete(order, z_matches, axis=0)
return keep, keep_z
def get_mirrored_patch_crops(patch_crops, org_img_shape):
"""
apply 3 mirrror transformations (x-axis, y-axis, x&y-axis)
to given patch crop coordinates and return the transformed coordinates.
Handles 2D and 3D coordinates.
:param patch_crops: list of crops: each element is a list of coordinates for given crop [[y1, x1, ...], [y1, x1, ..]]
:param org_img_shape: shape of patient volume used as world coordinates.
:return: list of mirrored patch crops: lenght=3. each element is a list of transformed patch crops.
"""
mirrored_patch_crops = []
# y-axis transform.
mirrored_patch_crops.append([[org_img_shape[2] - ii[1],
org_img_shape[2] - ii[0],
ii[2], ii[3]] if len(ii) == 4 else
[org_img_shape[2] - ii[1],
org_img_shape[2] - ii[0],
ii[2], ii[3], ii[4], ii[5]] for ii in patch_crops])
# x-axis transform.
mirrored_patch_crops.append([[ii[0], ii[1],
org_img_shape[3] - ii[3],
org_img_shape[3] - ii[2]] if len(ii) == 4 else
[ii[0], ii[1],
org_img_shape[3] - ii[3],
org_img_shape[3] - ii[2],
ii[4], ii[5]] for ii in patch_crops])
# y-axis and x-axis transform.
mirrored_patch_crops.append([[org_img_shape[2] - ii[1],
org_img_shape[2] - ii[0],
org_img_shape[3] - ii[3],
org_img_shape[3] - ii[2]] if len(ii) == 4 else
[org_img_shape[2] - ii[1],
org_img_shape[2] - ii[0],
org_img_shape[3] - ii[3],
org_img_shape[3] - ii[2],
ii[4], ii[5]] for ii in patch_crops])
return mirrored_patch_crops
diff --git a/utils/dataloader_utils.py b/utils/dataloader_utils.py
index 530f311..062af62 100644
--- a/utils/dataloader_utils.py
+++ b/utils/dataloader_utils.py
@@ -1,277 +1,278 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import numpy as np
import os
from multiprocessing import Pool
def get_class_balanced_patients(class_targets, batch_size, num_classes, slack_factor=0.1):
'''
samples patients towards equilibrium of classes on a roi-level. For highly imbalanced datasets, this might be a too strong requirement.
Hence a slack factor determines the ratio of the batch, that is randomly sampled, before class-balance is triggered.
:param class_targets: list of patient targets. where each patient target is a list of class labels of respective rois.
:param batch_size:
:param num_classes:
:param slack_factor:
:return: batch_ixs: list of indices referring to a subset in class_targets-list, sampled to build one batch.
'''
batch_ixs = []
class_count = {k: 0 for k in range(num_classes)}
weakest_class = 0
for ix in range(batch_size):
keep_looking = True
while keep_looking:
#choose a random patient.
cand = np.random.choice(len(class_targets), 1)[0]
# check the least occuring class among this patient's rois.
tmp_weakest_class = np.argmin([class_targets[cand].count(ii) for ii in range(num_classes)])
# if current batch already bigger than the slack_factor ratio, then
# check that weakest class in this patient is not the weakest in current batch (since needs to be boosted)
# also that at least one roi of this patient belongs to weakest class. If True, keep patient, else keep looking.
if (tmp_weakest_class != weakest_class and class_targets[cand].count(weakest_class) > 0) or ix < int(batch_size * slack_factor):
keep_looking = False
for c in range(num_classes):
class_count[c] += class_targets[cand].count(c)
weakest_class = np.argmin(([class_count[c] for c in range(num_classes)]))
batch_ixs.append(cand)
return batch_ixs
class fold_generator:
"""
generates splits of indices for a given length of a dataset to perform n-fold cross-validation.
splits each fold into 3 subsets for training, validation and testing.
This form of cross validation uses an inner loop test set, which is useful if test scores shall be reported on a
statistically reliable amount of patients, despite limited size of a dataset.
If hold out test set is provided and hence no inner loop test set needed, just add test_idxs to the training data in the dataloader.
This creates straight-forward train-val splits.
:returns names list: list of len n_splits. each element is a list of len 3 for train_ix, val_ix, test_ix.
"""
def __init__(self, seed, n_splits, len_data):
"""
:param seed: Random seed for splits.
:param n_splits: number of splits, e.g. 5 splits for 5-fold cross-validation
:param len_data: number of elements in the dataset.
"""
self.tr_ix = []
self.val_ix = []
self.te_ix = []
self.slicer = None
self.missing = 0
self.fold = 0
self.len_data = len_data
self.n_splits = n_splits
self.myseed = seed
self.boost_val = 0
def init_indices(self):
t = list(np.arange(self.l))
# round up to next splittable data amount.
split_length = int(np.ceil(len(t) / float(self.n_splits)))
self.slicer = split_length
self.mod = len(t) % self.n_splits
if self.mod > 0:
# missing is the number of folds, in which the new splits are reduced to account for missing data.
self.missing = self.n_splits - self.mod
self.te_ix = t[:self.slicer]
self.tr_ix = t[self.slicer:]
self.val_ix = self.tr_ix[:self.slicer]
self.tr_ix = self.tr_ix[self.slicer:]
def new_fold(self):
slicer = self.slicer
if self.fold < self.missing :
slicer = self.slicer - 1
temp = self.te_ix
# catch exception mod == 1: test set collects 1+ data since walk through both roudned up splits.
# account for by reducing last fold split by 1.
if self.fold == self.n_splits-2 and self.mod ==1:
temp += self.val_ix[-1:]
self.val_ix = self.val_ix[:-1]
self.te_ix = self.val_ix
self.val_ix = self.tr_ix[:slicer]
self.tr_ix = self.tr_ix[slicer:] + temp
def get_fold_names(self):
names_list = []
rgen = np.random.RandomState(self.myseed)
cv_names = np.arange(self.len_data)
rgen.shuffle(cv_names)
self.l = len(cv_names)
self.init_indices()
for split in range(self.n_splits):
train_names, val_names, test_names = cv_names[self.tr_ix], cv_names[self.val_ix], cv_names[self.te_ix]
names_list.append([train_names, val_names, test_names, self.fold])
self.new_fold()
self.fold += 1
return names_list
def get_patch_crop_coords(img, patch_size, min_overlap=30):
"""
_:param img (y, x, (z))
_:param patch_size: list of len 2 (2D) or 3 (3D).
_:param min_overlap: minimum required overlap of patches.
If too small, some areas are poorly represented only at edges of single patches.
_:return ndarray: shape (n_patches, 2*dim). crop coordinates for each patch.
"""
crop_coords = []
for dim in range(len(img.shape)):
n_patches = int(np.ceil(img.shape[dim] / patch_size[dim]))
# no crops required in this dimension, add image shape as coordinates.
if n_patches == 1:
crop_coords.append([(0, img.shape[dim])])
continue
# fix the two outside patches to coords patchsize/2 and interpolate.
center_dists = (img.shape[dim] - patch_size[dim]) / (n_patches - 1)
if (patch_size[dim] - center_dists) < min_overlap:
n_patches += 1
center_dists = (img.shape[dim] - patch_size[dim]) / (n_patches - 1)
patch_centers = np.round([(patch_size[dim] / 2 + (center_dists * ii)) for ii in range(n_patches)])
dim_crop_coords = [(center - patch_size[dim] / 2, center + patch_size[dim] / 2) for center in patch_centers]
crop_coords.append(dim_crop_coords)
coords_mesh_grid = []
for ymin, ymax in crop_coords[0]:
for xmin, xmax in crop_coords[1]:
if len(crop_coords) == 3 and patch_size[2] > 1:
for zmin, zmax in crop_coords[2]:
coords_mesh_grid.append([ymin, ymax, xmin, xmax, zmin, zmax])
elif len(crop_coords) == 3 and patch_size[2] == 1:
for zmin in range(img.shape[2]):
coords_mesh_grid.append([ymin, ymax, xmin, xmax, zmin, zmin + 1])
else:
coords_mesh_grid.append([ymin, ymax, xmin, xmax])
return np.array(coords_mesh_grid).astype(int)
def pad_nd_image(image, new_shape=None, mode="edge", kwargs=None, return_slicer=False, shape_must_be_divisible_by=None):
"""
one padder to pad them all. Documentation? Well okay. A little bit. by Fabian Isensee
:param image: nd image. can be anything
:param new_shape: what shape do you want? new_shape does not have to have the same dimensionality as image. If
len(new_shape) < len(image.shape) then the last axes of image will be padded. If new_shape < image.shape in any of
the axes then we will not pad that axis, but also not crop! (interpret new_shape as new_min_shape)
Example:
image.shape = (10, 1, 512, 512); new_shape = (768, 768) -> result: (10, 1, 768, 768). Cool, huh?
image.shape = (10, 1, 512, 512); new_shape = (364, 768) -> result: (10, 1, 512, 768).
:param mode: see np.pad for documentation
:param return_slicer: if True then this function will also return what coords you will need to use when cropping back
to original shape
:param shape_must_be_divisible_by: for network prediction. After applying new_shape, make sure the new shape is
divisibly by that number (can also be a list with an entry for each axis). Whatever is missing to match that will
be padded (so the result may be larger than new_shape if shape_must_be_divisible_by is not None)
:param kwargs: see np.pad for documentation
"""
if kwargs is None:
kwargs = {}
if new_shape is not None:
old_shape = np.array(image.shape[-len(new_shape):])
else:
assert shape_must_be_divisible_by is not None
assert isinstance(shape_must_be_divisible_by, (list, tuple, np.ndarray))
new_shape = image.shape[-len(shape_must_be_divisible_by):]
old_shape = new_shape
num_axes_nopad = len(image.shape) - len(new_shape)
new_shape = [max(new_shape[i], old_shape[i]) for i in range(len(new_shape))]
if not isinstance(new_shape, np.ndarray):
new_shape = np.array(new_shape)
if shape_must_be_divisible_by is not None:
if not isinstance(shape_must_be_divisible_by, (list, tuple, np.ndarray)):
shape_must_be_divisible_by = [shape_must_be_divisible_by] * len(new_shape)
else:
assert len(shape_must_be_divisible_by) == len(new_shape)
for i in range(len(new_shape)):
if new_shape[i] % shape_must_be_divisible_by[i] == 0:
new_shape[i] -= shape_must_be_divisible_by[i]
new_shape = np.array([new_shape[i] + shape_must_be_divisible_by[i] - new_shape[i] % shape_must_be_divisible_by[i] for i in range(len(new_shape))])
difference = new_shape - old_shape
pad_below = difference // 2
pad_above = difference // 2 + difference % 2
pad_list = [[0, 0]]*num_axes_nopad + list([list(i) for i in zip(pad_below, pad_above)])
res = np.pad(image, pad_list, mode, **kwargs)
if not return_slicer:
return res
else:
pad_list = np.array(pad_list)
pad_list[:, 1] = np.array(res.shape) - pad_list[:, 1]
slicer = list(slice(*i) for i in pad_list)
return res, slicer
#############################
# data packing / unpacking #
#############################
def get_case_identifiers(folder):
case_identifiers = [i[:-4] for i in os.listdir(folder) if i.endswith("npz")]
return case_identifiers
def convert_to_npy(npz_file):
- if not os.path.isfile(npz_file[:-3] + "npy"):
- a = np.load(npz_file)['data']
- np.save(npz_file[:-3] + "npy", a)
+ identifier = os.path.split(npz_file)[1][:-4]
+ if not os.path.isfile(npz_file[:-4] + ".npy"):
+ a = np.load(npz_file)[identifier]
+ np.save(npz_file[:-4] + ".npy", a)
def unpack_dataset(folder, threads=8):
case_identifiers = get_case_identifiers(folder)
p = Pool(threads)
npz_files = [os.path.join(folder, i + ".npz") for i in case_identifiers]
p.map(convert_to_npy, npz_files)
p.close()
p.join()
def delete_npy(folder):
case_identifiers = get_case_identifiers(folder)
npy_files = [os.path.join(folder, i + ".npy") for i in case_identifiers]
npy_files = [i for i in npy_files if os.path.isfile(i)]
for n in npy_files:
os.remove(n)
\ No newline at end of file
diff --git a/utils/model_utils.py b/utils/model_utils.py
index 3251577..70c1fae 100644
--- a/utils/model_utils.py
+++ b/utils/model_utils.py
@@ -1,1011 +1,1012 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""
Parts are based on https://github.com/multimodallearning/pytorch-mask-rcnn
published under MIT license.
"""
import numpy as np
import scipy.misc
import scipy.ndimage
import scipy.interpolate
import torch
from torch.autograd import Variable
import torch.nn as nn
import tqdm
############################################################
# Bounding Boxes
############################################################
def compute_iou_2D(box, boxes, box_area, boxes_area):
"""Calculates IoU of the given box with the array of the given boxes.
box: 1D vector [y1, x1, y2, x2] THIS IS THE GT BOX
boxes: [boxes_count, (y1, x1, y2, x2)]
box_area: float. the area of 'box'
boxes_area: array of length boxes_count.
Note: the areas are passed in rather than calculated here for
efficency. Calculate once in the caller to avoid duplicate work.
"""
# Calculate intersection areas
y1 = np.maximum(box[0], boxes[:, 0])
y2 = np.minimum(box[2], boxes[:, 2])
x1 = np.maximum(box[1], boxes[:, 1])
x2 = np.minimum(box[3], boxes[:, 3])
intersection = np.maximum(x2 - x1, 0) * np.maximum(y2 - y1, 0)
union = box_area + boxes_area[:] - intersection[:]
iou = intersection / union
return iou
def compute_iou_3D(box, boxes, box_volume, boxes_volume):
"""Calculates IoU of the given box with the array of the given boxes.
box: 1D vector [y1, x1, y2, x2, z1, z2] (typically gt box)
boxes: [boxes_count, (y1, x1, y2, x2, z1, z2)]
box_area: float. the area of 'box'
boxes_area: array of length boxes_count.
Note: the areas are passed in rather than calculated here for
efficency. Calculate once in the caller to avoid duplicate work.
"""
# Calculate intersection areas
y1 = np.maximum(box[0], boxes[:, 0])
y2 = np.minimum(box[2], boxes[:, 2])
x1 = np.maximum(box[1], boxes[:, 1])
x2 = np.minimum(box[3], boxes[:, 3])
z1 = np.maximum(box[4], boxes[:, 4])
z2 = np.minimum(box[5], boxes[:, 5])
intersection = np.maximum(x2 - x1, 0) * np.maximum(y2 - y1, 0) * np.maximum(z2 - z1, 0)
union = box_volume + boxes_volume[:] - intersection[:]
iou = intersection / union
return iou
def compute_overlaps(boxes1, boxes2):
"""Computes IoU overlaps between two sets of boxes.
boxes1, boxes2: [N, (y1, x1, y2, x2)]. / 3D: (z1, z2))
For better performance, pass the largest set first and the smaller second.
"""
# Areas of anchors and GT boxes
if boxes1.shape[1] == 4:
area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])
# Compute overlaps to generate matrix [boxes1 count, boxes2 count]
# Each cell contains the IoU value.
overlaps = np.zeros((boxes1.shape[0], boxes2.shape[0]))
for i in range(overlaps.shape[1]):
box2 = boxes2[i] #this is the gt box
overlaps[:, i] = compute_iou_2D(box2, boxes1, area2[i], area1)
return overlaps
else:
# Areas of anchors and GT boxes
volume1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1]) * (boxes1[:, 5] - boxes1[:, 4])
volume2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1]) * (boxes2[:, 5] - boxes2[:, 4])
# Compute overlaps to generate matrix [boxes1 count, boxes2 count]
# Each cell contains the IoU value.
overlaps = np.zeros((boxes1.shape[0], boxes2.shape[0]))
for i in range(overlaps.shape[1]):
box2 = boxes2[i] # this is the gt box
overlaps[:, i] = compute_iou_3D(box2, boxes1, volume2[i], volume1)
return overlaps
def box_refinement(box, gt_box):
"""Compute refinement needed to transform box to gt_box.
box and gt_box are [N, (y1, x1, y2, x2)] / 3D: (z1, z2))
"""
height = box[:, 2] - box[:, 0]
width = box[:, 3] - box[:, 1]
center_y = box[:, 0] + 0.5 * height
center_x = box[:, 1] + 0.5 * width
gt_height = gt_box[:, 2] - gt_box[:, 0]
gt_width = gt_box[:, 3] - gt_box[:, 1]
gt_center_y = gt_box[:, 0] + 0.5 * gt_height
gt_center_x = gt_box[:, 1] + 0.5 * gt_width
dy = (gt_center_y - center_y) / height
dx = (gt_center_x - center_x) / width
dh = torch.log(gt_height / height)
dw = torch.log(gt_width / width)
result = torch.stack([dy, dx, dh, dw], dim=1)
if box.shape[1] > 4:
depth = box[:, 5] - box[:, 4]
center_z = box[:, 4] + 0.5 * depth
gt_depth = gt_box[:, 5] - gt_box[:, 4]
gt_center_z = gt_box[:, 4] + 0.5 * gt_depth
dz = (gt_center_z - center_z) / depth
dd = torch.log(gt_depth / depth)
result = torch.stack([dy, dx, dz, dh, dw, dd], dim=1)
return result
def unmold_mask_2D(mask, bbox, image_shape):
"""Converts a mask generated by the neural network into a format similar
to it's original shape.
mask: [height, width] of type float. A small, typically 28x28 mask.
bbox: [y1, x1, y2, x2]. The box to fit the mask in.
Returns a binary mask with the same size as the original image.
"""
y1, x1, y2, x2 = bbox
out_zoom = [y2 - y1, x2 - x1]
zoom_factor = [i / j for i, j in zip(out_zoom, mask.shape)]
mask = scipy.ndimage.zoom(mask, zoom_factor, order=1).astype(np.float32)
# Put the mask in the right location.
full_mask = np.zeros(image_shape[:2])
full_mask[y1:y2, x1:x2] = mask
return full_mask
def unmold_mask_3D(mask, bbox, image_shape):
"""Converts a mask generated by the neural network into a format similar
to it's original shape.
mask: [height, width] of type float. A small, typically 28x28 mask.
bbox: [y1, x1, y2, x2, z1, z2]. The box to fit the mask in.
Returns a binary mask with the same size as the original image.
"""
y1, x1, y2, x2, z1, z2 = bbox
out_zoom = [y2 - y1, x2 - x1, z2 - z1]
zoom_factor = [i/j for i,j in zip(out_zoom, mask.shape)]
mask = scipy.ndimage.zoom(mask, zoom_factor, order=1).astype(np.float32)
# Put the mask in the right location.
full_mask = np.zeros(image_shape[:3])
full_mask[y1:y2, x1:x2, z1:z2] = mask
return full_mask
############################################################
# Anchors
############################################################
def generate_anchors(scales, ratios, shape, feature_stride, anchor_stride):
"""
scales: 1D array of anchor sizes in pixels. Example: [32, 64, 128]
ratios: 1D array of anchor ratios of width/height. Example: [0.5, 1, 2]
shape: [height, width] spatial shape of the feature map over which
to generate anchors.
feature_stride: Stride of the feature map relative to the image in pixels.
anchor_stride: Stride of anchors on the feature map. For example, if the
value is 2 then generate anchors for every other feature map pixel.
"""
# Get all combinations of scales and ratios
scales, ratios = np.meshgrid(np.array(scales), np.array(ratios))
scales = scales.flatten()
ratios = ratios.flatten()
# Enumerate heights and widths from scales and ratios
heights = scales / np.sqrt(ratios)
widths = scales * np.sqrt(ratios)
# Enumerate shifts in feature space
shifts_y = np.arange(0, shape[0], anchor_stride) * feature_stride
shifts_x = np.arange(0, shape[1], anchor_stride) * feature_stride
shifts_x, shifts_y = np.meshgrid(shifts_x, shifts_y)
# Enumerate combinations of shifts, widths, and heights
box_widths, box_centers_x = np.meshgrid(widths, shifts_x)
box_heights, box_centers_y = np.meshgrid(heights, shifts_y)
# Reshape to get a list of (y, x) and a list of (h, w)
box_centers = np.stack(
[box_centers_y, box_centers_x], axis=2).reshape([-1, 2])
box_sizes = np.stack([box_heights, box_widths], axis=2).reshape([-1, 2])
# Convert to corner coordinates (y1, x1, y2, x2)
boxes = np.concatenate([box_centers - 0.5 * box_sizes,
box_centers + 0.5 * box_sizes], axis=1)
return boxes
def generate_anchors_3D(scales_xy, scales_z, ratios, shape, feature_stride_xy, feature_stride_z, anchor_stride):
"""
scales: 1D array of anchor sizes in pixels. Example: [32, 64, 128]
ratios: 1D array of anchor ratios of width/height. Example: [0.5, 1, 2]
shape: [height, width] spatial shape of the feature map over which
to generate anchors.
feature_stride: Stride of the feature map relative to the image in pixels.
anchor_stride: Stride of anchors on the feature map. For example, if the
value is 2 then generate anchors for every other feature map pixel.
"""
# Get all combinations of scales and ratios
scales_xy, ratios_meshed = np.meshgrid(np.array(scales_xy), np.array(ratios))
scales_xy = scales_xy.flatten()
ratios_meshed = ratios_meshed.flatten()
# Enumerate heights and widths from scales and ratios
heights = scales_xy / np.sqrt(ratios_meshed)
widths = scales_xy * np.sqrt(ratios_meshed)
depths = np.tile(np.array(scales_z), len(ratios_meshed)//np.array(scales_z)[..., None].shape[0])
# Enumerate shifts in feature space
shifts_y = np.arange(0, shape[0], anchor_stride) * feature_stride_xy #translate from fm positions to input coords.
shifts_x = np.arange(0, shape[1], anchor_stride) * feature_stride_xy
shifts_z = np.arange(0, shape[2], anchor_stride) * (feature_stride_z)
shifts_x, shifts_y, shifts_z = np.meshgrid(shifts_x, shifts_y, shifts_z)
# Enumerate combinations of shifts, widths, and heights
box_widths, box_centers_x = np.meshgrid(widths, shifts_x)
box_heights, box_centers_y = np.meshgrid(heights, shifts_y)
box_depths, box_centers_z = np.meshgrid(depths, shifts_z)
# Reshape to get a list of (y, x, z) and a list of (h, w, d)
box_centers = np.stack(
[box_centers_y, box_centers_x, box_centers_z], axis=2).reshape([-1, 3])
box_sizes = np.stack([box_heights, box_widths, box_depths], axis=2).reshape([-1, 3])
# Convert to corner coordinates (y1, x1, y2, x2, z1, z2)
boxes = np.concatenate([box_centers - 0.5 * box_sizes,
box_centers + 0.5 * box_sizes], axis=1)
boxes = np.transpose(np.array([boxes[:, 0], boxes[:, 1], boxes[:, 3], boxes[:, 4], boxes[:, 2], boxes[:, 5]]), axes=(1, 0))
return boxes
def generate_pyramid_anchors(logger, cf):
"""Generate anchors at different levels of a feature pyramid. Each scale
is associated with a level of the pyramid, but each ratio is used in
all levels of the pyramid.
from configs:
:param scales: cf.RPN_ANCHOR_SCALES , e.g. [4, 8, 16, 32]
:param ratios: cf.RPN_ANCHOR_RATIOS , e.g. [0.5, 1, 2]
:param feature_shapes: cf.BACKBONE_SHAPES , e.g. [array of shapes per feature map] [80, 40, 20, 10, 5]
:param feature_strides: cf.BACKBONE_STRIDES , e.g. [2, 4, 8, 16, 32, 64]
:param anchors_stride: cf.RPN_ANCHOR_STRIDE , e.g. 1
:return anchors: (N, (y1, x1, y2, x2, (z1), (z2)). All generated anchors in one array. Sorted
with the same order of the given scales. So, anchors of scale[0] come first, then anchors of scale[1], and so on.
"""
scales = cf.rpn_anchor_scales
ratios = cf.rpn_anchor_ratios
feature_shapes = cf.backbone_shapes
anchor_stride = cf.rpn_anchor_stride
pyramid_levels = cf.pyramid_levels
feature_strides = cf.backbone_strides
anchors = []
logger.info("feature map shapes: {}".format(feature_shapes))
logger.info("anchor scales: {}".format(scales))
expected_anchors = [np.prod(feature_shapes[ii]) * len(ratios) * len(scales['xy'][ii]) for ii in pyramid_levels]
for lix, level in enumerate(pyramid_levels):
if len(feature_shapes[level]) == 2:
anchors.append(generate_anchors(scales['xy'][level], ratios, feature_shapes[level],
feature_strides['xy'][level], anchor_stride))
else:
anchors.append(generate_anchors_3D(scales['xy'][level], scales['z'][level], ratios, feature_shapes[level],
feature_strides['xy'][level], feature_strides['z'][level], anchor_stride))
logger.info("level {}: built anchors {} / expected anchors {} ||| total build {} / total expected {}".format(
level, anchors[-1].shape, expected_anchors[lix], np.concatenate(anchors).shape, np.sum(expected_anchors)))
out_anchors = np.concatenate(anchors, axis=0)
return out_anchors
def apply_box_deltas_2D(boxes, deltas):
"""Applies the given deltas to the given boxes.
boxes: [N, 4] where each row is y1, x1, y2, x2
deltas: [N, 4] where each row is [dy, dx, log(dh), log(dw)]
"""
# Convert to y, x, h, w
height = boxes[:, 2] - boxes[:, 0]
width = boxes[:, 3] - boxes[:, 1]
center_y = boxes[:, 0] + 0.5 * height
center_x = boxes[:, 1] + 0.5 * width
# Apply deltas
center_y += deltas[:, 0] * height
center_x += deltas[:, 1] * width
height *= torch.exp(deltas[:, 2])
width *= torch.exp(deltas[:, 3])
# Convert back to y1, x1, y2, x2
y1 = center_y - 0.5 * height
x1 = center_x - 0.5 * width
y2 = y1 + height
x2 = x1 + width
result = torch.stack([y1, x1, y2, x2], dim=1)
return result
def apply_box_deltas_3D(boxes, deltas):
"""Applies the given deltas to the given boxes.
boxes: [N, 6] where each row is y1, x1, y2, x2, z1, z2
deltas: [N, 6] where each row is [dy, dx, dz, log(dh), log(dw), log(dd)]
"""
# Convert to y, x, h, w
height = boxes[:, 2] - boxes[:, 0]
width = boxes[:, 3] - boxes[:, 1]
depth = boxes[:, 5] - boxes[:, 4]
center_y = boxes[:, 0] + 0.5 * height
center_x = boxes[:, 1] + 0.5 * width
center_z = boxes[:, 4] + 0.5 * depth
# Apply deltas
center_y += deltas[:, 0] * height
center_x += deltas[:, 1] * width
center_z += deltas[:, 2] * depth
height *= torch.exp(deltas[:, 3])
width *= torch.exp(deltas[:, 4])
depth *= torch.exp(deltas[:, 5])
# Convert back to y1, x1, y2, x2
y1 = center_y - 0.5 * height
x1 = center_x - 0.5 * width
z1 = center_z - 0.5 * depth
y2 = y1 + height
x2 = x1 + width
z2 = z1 + depth
result = torch.stack([y1, x1, y2, x2, z1, z2], dim=1)
return result
def clip_boxes_2D(boxes, window):
"""
boxes: [N, 4] each col is y1, x1, y2, x2
window: [4] in the form y1, x1, y2, x2
"""
boxes = torch.stack( \
[boxes[:, 0].clamp(float(window[0]), float(window[2])),
boxes[:, 1].clamp(float(window[1]), float(window[3])),
boxes[:, 2].clamp(float(window[0]), float(window[2])),
boxes[:, 3].clamp(float(window[1]), float(window[3]))], 1)
return boxes
def clip_boxes_3D(boxes, window):
"""
boxes: [N, 6] each col is y1, x1, y2, x2, z1, z2
window: [6] in the form y1, x1, y2, x2, z1, z2
"""
boxes = torch.stack( \
[boxes[:, 0].clamp(float(window[0]), float(window[2])),
boxes[:, 1].clamp(float(window[1]), float(window[3])),
boxes[:, 2].clamp(float(window[0]), float(window[2])),
boxes[:, 3].clamp(float(window[1]), float(window[3])),
boxes[:, 4].clamp(float(window[4]), float(window[5])),
boxes[:, 5].clamp(float(window[4]), float(window[5]))], 1)
return boxes
def clip_boxes_numpy(boxes, window):
"""
boxes: [N, 4] each col is y1, x1, y2, x2 / [N, 6] in 3D.
window: iamge shape (y, x, (z))
"""
if boxes.shape[1] == 4:
boxes = np.concatenate(
(np.clip(boxes[:, 0], 0, window[0])[:, None],
np.clip(boxes[:, 1], 0, window[0])[:, None],
np.clip(boxes[:, 2], 0, window[1])[:, None],
np.clip(boxes[:, 3], 0, window[1])[:, None]), 1
)
else:
boxes = np.concatenate(
(np.clip(boxes[:, 0], 0, window[0])[:, None],
np.clip(boxes[:, 1], 0, window[0])[:, None],
np.clip(boxes[:, 2], 0, window[1])[:, None],
np.clip(boxes[:, 3], 0, window[1])[:, None],
np.clip(boxes[:, 4], 0, window[2])[:, None],
np.clip(boxes[:, 5], 0, window[2])[:, None]), 1
)
return boxes
def bbox_overlaps_2D(boxes1, boxes2):
"""Computes IoU overlaps between two sets of boxes.
boxes1, boxes2: [N, (y1, x1, y2, x2)].
"""
# 1. Tile boxes2 and repeate boxes1. This allows us to compare
# every boxes1 against every boxes2 without loops.
# TF doesn't have an equivalent to np.repeate() so simulate it
# using tf.tile() and tf.reshape.
boxes1_repeat = boxes2.size()[0]
boxes2_repeat = boxes1.size()[0]
boxes1 = boxes1.repeat(1,boxes1_repeat).view(-1,4)
boxes2 = boxes2.repeat(boxes2_repeat,1)
# 2. Compute intersections
b1_y1, b1_x1, b1_y2, b1_x2 = boxes1.chunk(4, dim=1)
b2_y1, b2_x1, b2_y2, b2_x2 = boxes2.chunk(4, dim=1)
y1 = torch.max(b1_y1, b2_y1)[:, 0]
x1 = torch.max(b1_x1, b2_x1)[:, 0]
y2 = torch.min(b1_y2, b2_y2)[:, 0]
x2 = torch.min(b1_x2, b2_x2)[:, 0]
zeros = Variable(torch.zeros(y1.size()[0]), requires_grad=False)
if y1.is_cuda:
zeros = zeros.cuda()
intersection = torch.max(x2 - x1, zeros) * torch.max(y2 - y1, zeros)
# 3. Compute unions
b1_area = (b1_y2 - b1_y1) * (b1_x2 - b1_x1)
b2_area = (b2_y2 - b2_y1) * (b2_x2 - b2_x1)
union = b1_area[:,0] + b2_area[:,0] - intersection
# 4. Compute IoU and reshape to [boxes1, boxes2]
iou = intersection / union
overlaps = iou.view(boxes2_repeat, boxes1_repeat)
return overlaps
def bbox_overlaps_3D(boxes1, boxes2):
"""Computes IoU overlaps between two sets of boxes.
boxes1, boxes2: [N, (y1, x1, y2, x2, z1, z2)].
"""
# 1. Tile boxes2 and repeate boxes1. This allows us to compare
# every boxes1 against every boxes2 without loops.
# TF doesn't have an equivalent to np.repeate() so simulate it
# using tf.tile() and tf.reshape.
boxes1_repeat = boxes2.size()[0]
boxes2_repeat = boxes1.size()[0]
boxes1 = boxes1.repeat(1,boxes1_repeat).view(-1,6)
boxes2 = boxes2.repeat(boxes2_repeat,1)
# 2. Compute intersections
b1_y1, b1_x1, b1_y2, b1_x2, b1_z1, b1_z2 = boxes1.chunk(6, dim=1)
b2_y1, b2_x1, b2_y2, b2_x2, b2_z1, b2_z2 = boxes2.chunk(6, dim=1)
y1 = torch.max(b1_y1, b2_y1)[:, 0]
x1 = torch.max(b1_x1, b2_x1)[:, 0]
y2 = torch.min(b1_y2, b2_y2)[:, 0]
x2 = torch.min(b1_x2, b2_x2)[:, 0]
z1 = torch.max(b1_z1, b2_z1)[:, 0]
z2 = torch.min(b1_z2, b2_z2)[:, 0]
zeros = Variable(torch.zeros(y1.size()[0]), requires_grad=False)
if y1.is_cuda:
zeros = zeros.cuda()
intersection = torch.max(x2 - x1, zeros) * torch.max(y2 - y1, zeros) * torch.max(z2 - z1, zeros)
# 3. Compute unions
b1_volume = (b1_y2 - b1_y1) * (b1_x2 - b1_x1) * (b1_z2 - b1_z1)
b2_volume = (b2_y2 - b2_y1) * (b2_x2 - b2_x1) * (b2_z2 - b2_z1)
union = b1_volume[:,0] + b2_volume[:,0] - intersection
# 4. Compute IoU and reshape to [boxes1, boxes2]
iou = intersection / union
overlaps = iou.view(boxes2_repeat, boxes1_repeat)
return overlaps
def gt_anchor_matching(cf, anchors, gt_boxes, gt_class_ids=None):
"""Given the anchors and GT boxes, compute overlaps and identify positive
anchors and deltas to refine them to match their corresponding GT boxes.
anchors: [num_anchors, (y1, x1, y2, x2, (z1), (z2))]
gt_boxes: [num_gt_boxes, (y1, x1, y2, x2, (z1), (z2))]
gt_class_ids (optional): [num_gt_boxes] Integer class IDs for one stage detectors. in RPN case of Mask R-CNN,
set all positive matches to 1 (foreground)
Returns:
anchor_class_matches: [N] (int32) matches between anchors and GT boxes.
1 = positive anchor, -1 = negative anchor, 0 = neutral.
In case of one stage detectors like RetinaNet/RetinaUNet this flag takes
class_ids as positive anchor values, i.e. values >= 1!
anchor_delta_targets: [N, (dy, dx, (dz), log(dh), log(dw), (log(dd)))] Anchor bbox deltas.
"""
anchor_class_matches = np.zeros([anchors.shape[0]], dtype=np.int32)
anchor_delta_targets = np.zeros((cf.rpn_train_anchors_per_image, 2*cf.dim))
anchor_matching_iou = cf.anchor_matching_iou
if gt_boxes is None:
anchor_class_matches = np.full(anchor_class_matches.shape, fill_value=-1)
return anchor_class_matches, anchor_delta_targets
# for mrcnn: anchor matching is done for RPN loss, so positive labels are all 1 (foreground)
if gt_class_ids is None:
gt_class_ids = np.array([1] * len(gt_boxes))
# Compute overlaps [num_anchors, num_gt_boxes]
overlaps = compute_overlaps(anchors, gt_boxes)
# Match anchors to GT Boxes
# If an anchor overlaps a GT box with IoU >= anchor_matching_iou then it's positive.
# If an anchor overlaps a GT box with IoU < 0.1 then it's negative.
# Neutral anchors are those that don't match the conditions above,
# and they don't influence the loss function.
# However, don't keep any GT box unmatched (rare, but happens). Instead,
# match it to the closest anchor (even if its max IoU is < 0.1).
# 1. Set negative anchors first. They get overwritten below if a GT box is
# matched to them. Skip boxes in crowd areas.
anchor_iou_argmax = np.argmax(overlaps, axis=1)
anchor_iou_max = overlaps[np.arange(overlaps.shape[0]), anchor_iou_argmax]
if anchors.shape[1] == 4:
anchor_class_matches[(anchor_iou_max < 0.1)] = -1
elif anchors.shape[1] == 6:
anchor_class_matches[(anchor_iou_max < 0.01)] = -1
else:
raise ValueError('anchor shape wrong {}'.format(anchors.shape))
# 2. Set an anchor for each GT box (regardless of IoU value).
gt_iou_argmax = np.argmax(overlaps, axis=0)
for ix, ii in enumerate(gt_iou_argmax):
anchor_class_matches[ii] = gt_class_ids[ix]
# 3. Set anchors with high overlap as positive.
above_trhesh_ixs = np.argwhere(anchor_iou_max >= anchor_matching_iou)
anchor_class_matches[above_trhesh_ixs] = gt_class_ids[anchor_iou_argmax[above_trhesh_ixs]]
# Subsample to balance positive anchors.
ids = np.where(anchor_class_matches > 0)[0]
+ # extra == these positive anchors are too many --> reset them to negative ones.
extra = len(ids) - (cf.rpn_train_anchors_per_image // 2)
if extra > 0:
# Reset the extra ones to neutral
- ids = np.random.choice(ids, extra, replace=False)
- anchor_class_matches[ids] = 0
+ extra_ids = np.random.choice(ids, extra, replace=False)
+ anchor_class_matches[extra_ids] = 0
# Leave all negative proposals negative now and sample from them in online hard example mining.
# For positive anchors, compute shift and scale needed to transform them to match the corresponding GT boxes.
- ids = np.where(anchor_class_matches > 0)[0]
+ #ids = np.where(anchor_class_matches > 0)[0]
ix = 0 # index into anchor_delta_targets
for i, a in zip(ids, anchors[ids]):
# closest gt box (it might have IoU < anchor_matching_iou)
gt = gt_boxes[anchor_iou_argmax[i]]
# convert coordinates to center plus width/height.
gt_h = gt[2] - gt[0]
gt_w = gt[3] - gt[1]
gt_center_y = gt[0] + 0.5 * gt_h
gt_center_x = gt[1] + 0.5 * gt_w
# Anchor
a_h = a[2] - a[0]
a_w = a[3] - a[1]
a_center_y = a[0] + 0.5 * a_h
a_center_x = a[1] + 0.5 * a_w
if cf.dim == 2:
anchor_delta_targets[ix] = [
(gt_center_y - a_center_y) / a_h,
(gt_center_x - a_center_x) / a_w,
np.log(gt_h / a_h),
np.log(gt_w / a_w),
]
else:
gt_d = gt[5] - gt[4]
gt_center_z = gt[4] + 0.5 * gt_d
a_d = a[5] - a[4]
a_center_z = a[4] + 0.5 * a_d
anchor_delta_targets[ix] = [
(gt_center_y - a_center_y) / a_h,
(gt_center_x - a_center_x) / a_w,
(gt_center_z - a_center_z) / a_d,
np.log(gt_h / a_h),
np.log(gt_w / a_w),
np.log(gt_d / a_d)
]
# normalize.
anchor_delta_targets[ix] /= cf.rpn_bbox_std_dev
ix += 1
return anchor_class_matches, anchor_delta_targets
def clip_to_window(window, boxes):
"""
window: (y1, x1, y2, x2) / 3D: (z1, z2). The window in the image we want to clip to.
boxes: [N, (y1, x1, y2, x2)] / 3D: (z1, z2)
"""
boxes[:, 0] = boxes[:, 0].clamp(float(window[0]), float(window[2]))
boxes[:, 1] = boxes[:, 1].clamp(float(window[1]), float(window[3]))
boxes[:, 2] = boxes[:, 2].clamp(float(window[0]), float(window[2]))
boxes[:, 3] = boxes[:, 3].clamp(float(window[1]), float(window[3]))
if boxes.shape[1] > 5:
boxes[:, 4] = boxes[:, 4].clamp(float(window[4]), float(window[5]))
boxes[:, 5] = boxes[:, 5].clamp(float(window[4]), float(window[5]))
return boxes
def nms_numpy(box_coords, scores, thresh):
""" non-maximum suppression on 2D or 3D boxes in numpy.
:param box_coords: [y1,x1,y2,x2 (,z1,z2)] with y1<=y2, x1<=x2, z1<=z2.
:param scores: ranking scores (higher score == higher rank) of boxes.
:param thresh: IoU threshold for clustering.
:return:
"""
y1 = box_coords[:, 0]
x1 = box_coords[:, 1]
y2 = box_coords[:, 2]
x2 = box_coords[:, 3]
assert np.all(y1 <= y2) and np.all(x1 <= x2), """"the definition of the coordinates is crucially important here:
coordinates of which maxima are taken need to be the lower coordinates"""
areas = (x2 - x1) * (y2 - y1)
is_3d = box_coords.shape[1] == 6
if is_3d: # 3-dim case
z1 = box_coords[:, 4]
z2 = box_coords[:, 5]
assert np.all(z1<=z2), """"the definition of the coordinates is crucially important here:
coordinates of which maxima are taken need to be the lower coordinates"""
areas *= (z2 - z1)
order = scores.argsort()[::-1]
keep = []
while order.size > 0: # order is the sorted index. maps order to index: order[1] = 24 means (rank1, ix 24)
i = order[0] # highest scoring element
yy1 = np.maximum(y1[i], y1[order]) # highest scoring element still in >order<, is compared to itself, that is okay.
xx1 = np.maximum(x1[i], x1[order])
yy2 = np.minimum(y2[i], y2[order])
xx2 = np.minimum(x2[i], x2[order])
h = np.maximum(0.0, yy2 - yy1)
w = np.maximum(0.0, xx2 - xx1)
inter = h * w
if is_3d:
zz1 = np.maximum(z1[i], z1[order])
zz2 = np.minimum(z2[i], z2[order])
d = np.maximum(0.0, zz2 - zz1)
inter *= d
iou = inter / (areas[i] + areas[order] - inter)
non_matches = np.nonzero(iou <= thresh)[0] # get all elements that were not matched and discard all others.
order = order[non_matches]
keep.append(i)
return keep
def roi_align_3d_numpy(input: np.ndarray, rois, output_size: tuple,
spatial_scale: float = 1., sampling_ratio: int = -1) -> np.ndarray:
""" This fct mainly serves as a verification method for 3D CUDA implementation of RoIAlign, it's highly
inefficient due to the nested loops.
:param input: (ndarray[N, C, H, W, D]): input feature map
:param rois: list (N,K(n), 6), K(n) = nr of rois in batch-element n, single roi of format (y1,x1,y2,x2,z1,z2)
:param output_size:
:param spatial_scale:
:param sampling_ratio:
:return: (List[N, K(n), C, output_size[0], output_size[1], output_size[2]])
"""
out_height, out_width, out_depth = output_size
coord_grid = tuple([np.linspace(0, input.shape[dim] - 1, num=input.shape[dim]) for dim in range(2, 5)])
pooled_rois = [[]] * len(rois)
assert len(rois) == input.shape[0], "batch dim mismatch, rois: {}, input: {}".format(len(rois), input.shape[0])
print("Numpy 3D RoIAlign progress:", end="\n")
for b in range(input.shape[0]):
for roi in tqdm.tqdm(rois[b]):
y1, x1, y2, x2, z1, z2 = np.array(roi) * spatial_scale
roi_height = max(float(y2 - y1), 1.)
roi_width = max(float(x2 - x1), 1.)
roi_depth = max(float(z2 - z1), 1.)
if sampling_ratio <= 0:
sampling_ratio_h = int(np.ceil(roi_height / out_height))
sampling_ratio_w = int(np.ceil(roi_width / out_width))
sampling_ratio_d = int(np.ceil(roi_depth / out_depth))
else:
sampling_ratio_h = sampling_ratio_w = sampling_ratio_d = sampling_ratio # == n points per bin
bin_height = roi_height / out_height
bin_width = roi_width / out_width
bin_depth = roi_depth / out_depth
n_points = sampling_ratio_h * sampling_ratio_w * sampling_ratio_d
pooled_roi = np.empty((input.shape[1], out_height, out_width, out_depth), dtype="float32")
for chan in range(input.shape[1]):
lin_interpolator = scipy.interpolate.RegularGridInterpolator(coord_grid, input[b, chan],
method="linear")
for bin_iy in range(out_height):
for bin_ix in range(out_width):
for bin_iz in range(out_depth):
bin_val = 0.
for i in range(sampling_ratio_h):
for j in range(sampling_ratio_w):
for k in range(sampling_ratio_d):
loc_ijk = [
y1 + bin_iy * bin_height + (i + 0.5) * (bin_height / sampling_ratio_h),
x1 + bin_ix * bin_width + (j + 0.5) * (bin_width / sampling_ratio_w),
z1 + bin_iz * bin_depth + (k + 0.5) * (bin_depth / sampling_ratio_d)]
# print("loc_ijk", loc_ijk)
if not (np.any([c < -1.0 for c in loc_ijk]) or loc_ijk[0] > input.shape[2] or
loc_ijk[1] > input.shape[3] or loc_ijk[2] > input.shape[4]):
for catch_case in range(3):
# catch on-border cases
if int(loc_ijk[catch_case]) == input.shape[catch_case + 2] - 1:
loc_ijk[catch_case] = input.shape[catch_case + 2] - 1
bin_val += lin_interpolator(loc_ijk)
pooled_roi[chan, bin_iy, bin_ix, bin_iz] = bin_val / n_points
pooled_rois[b].append(pooled_roi)
return np.array(pooled_rois)
############################################################
# Pytorch Utility Functions
############################################################
def unique1d(tensor):
- if tensor.size()[0] == 0 or tensor.size()[0] == 1:
+ if tensor.shape[0] == 0 or tensor.shape[0] == 1:
return tensor
tensor = tensor.sort()[0]
unique_bool = tensor[1:] != tensor [:-1]
first_element = torch.tensor([True], dtype=torch.bool, requires_grad=False)
if tensor.is_cuda:
first_element = first_element.cuda()
unique_bool = torch.cat((first_element, unique_bool),dim=0)
- return tensor[unique_bool.data]
+ return tensor[unique_bool]
def log2(x):
"""Implementatin of Log2. Pytorch doesn't have a native implemenation."""
ln2 = Variable(torch.log(torch.FloatTensor([2.0])), requires_grad=False)
if x.is_cuda:
ln2 = ln2.cuda()
return torch.log(x) / ln2
def intersect1d(tensor1, tensor2):
aux = torch.cat((tensor1, tensor2), dim=0)
aux = aux.sort(descending=True)[0]
return aux[:-1][(aux[1:] == aux[:-1]).data]
def shem(roi_probs_neg, negative_count, ohem_poolsize):
"""
stochastic hard example mining: from a list of indices (referring to non-matched predictions),
determine a pool of highest scoring (worst false positives) of size negative_count*ohem_poolsize.
Then, sample n (= negative_count) predictions of this pool as negative examples for loss.
:param roi_probs_neg: tensor of shape (n_predictions, n_classes).
:param negative_count: int.
:param ohem_poolsize: int.
:return: (negative_count). indices refer to the positions in roi_probs_neg. If pool smaller than expected due to
limited negative proposals availabel, this function will return sampled indices of number < negative_count without
throwing an error.
"""
# sort according to higehst foreground score.
probs, order = roi_probs_neg[:, 1:].max(1)[0].sort(descending=True)
select = torch.tensor((ohem_poolsize * int(negative_count), order.size()[0])).min().int()
pool_indices = order[:select]
rand_idx = torch.randperm(pool_indices.size()[0])
return pool_indices[rand_idx[:negative_count].cuda()]
def initialize_weights(net):
"""
Initialize model weights. Current Default in Pytorch (version 0.4.1) is initialization from a uniform distriubtion.
Will expectably be changed to kaiming_uniform in future versions.
"""
init_type = net.cf.weight_init
for m in [module for module in net.modules() if type(module) in [nn.Conv2d, nn.Conv3d,
nn.ConvTranspose2d,
nn.ConvTranspose3d,
nn.Linear]]:
if init_type == 'xavier_uniform':
nn.init.xavier_uniform_(m.weight.data)
if m.bias is not None:
m.bias.data.zero_()
elif init_type == 'xavier_normal':
nn.init.xavier_normal_(m.weight.data)
if m.bias is not None:
m.bias.data.zero_()
elif init_type == "kaiming_uniform":
nn.init.kaiming_uniform_(m.weight.data, mode='fan_out', nonlinearity=net.cf.relu, a=0)
if m.bias is not None:
fan_in, fan_out = nn.init._calculate_fan_in_and_fan_out(m.weight.data)
bound = 1 / np.sqrt(fan_out)
nn.init.uniform_(m.bias, -bound, bound)
elif init_type == "kaiming_normal":
nn.init.kaiming_normal_(m.weight.data, mode='fan_out', nonlinearity=net.cf.relu, a=0)
if m.bias is not None:
fan_in, fan_out = nn.init._calculate_fan_in_and_fan_out(m.weight.data)
bound = 1 / np.sqrt(fan_out)
nn.init.normal_(m.bias, -bound, bound)
class NDConvGenerator(object):
"""
generic wrapper around conv-layers to avoid 2D vs. 3D distinguishing in code.
"""
def __init__(self, dim):
self.dim = dim
def __call__(self, c_in, c_out, ks, pad=0, stride=1, norm=None, relu='relu'):
"""
:param c_in: number of in_channels.
:param c_out: number of out_channels.
:param ks: kernel size.
:param pad: pad size.
:param stride: kernel stride.
:param norm: string specifying type of feature map normalization. If None, no normalization is applied.
:param relu: string specifying type of nonlinearity. If None, no nonlinearity is applied.
:return: convolved feature_map.
"""
if self.dim == 2:
conv = nn.Conv2d(c_in, c_out, kernel_size=ks, padding=pad, stride=stride)
if norm is not None:
if norm == 'instance_norm':
norm_layer = nn.InstanceNorm2d(c_out)
elif norm == 'batch_norm':
norm_layer = nn.BatchNorm2d(c_out)
else:
raise ValueError('norm type as specified in configs is not implemented...')
conv = nn.Sequential(conv, norm_layer)
else:
conv = nn.Conv3d(c_in, c_out, kernel_size=ks, padding=pad, stride=stride)
if norm is not None:
if norm == 'instance_norm':
norm_layer = nn.InstanceNorm3d(c_out)
elif norm == 'batch_norm':
norm_layer = nn.BatchNorm3d(c_out)
else:
raise ValueError('norm type as specified in configs is not implemented... {}'.format(norm))
conv = nn.Sequential(conv, norm_layer)
if relu is not None:
if relu == 'relu':
relu_layer = nn.ReLU(inplace=True)
elif relu == 'leaky_relu':
relu_layer = nn.LeakyReLU(inplace=True)
else:
raise ValueError('relu type as specified in configs is not implemented...')
conv = nn.Sequential(conv, relu_layer)
return conv
def get_one_hot_encoding(y, n_classes):
"""
transform a numpy label array to a one-hot array of the same shape.
:param y: array of shape (b, 1, y, x, (z)).
:param n_classes: int, number of classes to unfold in one-hot encoding.
:return y_ohe: array of shape (b, n_classes, y, x, (z))
"""
dim = len(y.shape) - 2
if dim == 2:
y_ohe = np.zeros((y.shape[0], n_classes, y.shape[2], y.shape[3])).astype('int32')
if dim ==3:
y_ohe = np.zeros((y.shape[0], n_classes, y.shape[2], y.shape[3], y.shape[4])).astype('int32')
for cl in range(n_classes):
y_ohe[:, cl][y[:, 0] == cl] = 1
return y_ohe
def get_dice_per_batch_and_class(pred, y, n_classes):
'''
computes dice scores per batch instance and class.
:param pred: prediction array of shape (b, 1, y, x, (z)) (e.g. softmax prediction with argmax over dim 1)
:param y: ground truth array of shape (b, 1, y, x, (z)) (contains int [0, ..., n_classes]
:param n_classes: int
:return: dice scores of shape (b, c)
'''
pred = get_one_hot_encoding(pred, n_classes)
y = get_one_hot_encoding(y, n_classes)
axes = tuple(range(2, len(pred.shape)))
intersect = np.sum(pred*y, axis=axes)
denominator = np.sum(pred, axis=axes)+np.sum(y, axis=axes) + 1e-8
dice = 2.0*intersect / denominator
return dice
def sum_tensor(input, axes, keepdim=False):
axes = np.unique(axes)
if keepdim:
for ax in axes:
input = input.sum(ax, keepdim=True)
else:
for ax in sorted(axes, reverse=True):
input = input.sum(int(ax))
return input
def batch_dice(pred, y, false_positive_weight=1.0, smooth=1e-6):
'''
compute soft dice over batch. this is a differentiable score and can be used as a loss function.
only dice scores of foreground classes are returned, since training typically
does not benefit from explicit background optimization. Pixels of the entire batch are considered a pseudo-volume to compute dice scores of.
This way, single patches with missing foreground classes can not produce faulty gradients.
:param pred: (b, c, y, x, (z)), softmax probabilities (network output). (c==classes)
:param y: (b, c, y, x, (z)), one-hot-encoded segmentation mask.
:param false_positive_weight: float [0,1]. For weighting of imbalanced classes,
reduces the penalty for false-positive pixels. Can be beneficial sometimes in data with heavy fg/bg imbalances.
:return: soft dice score (float). This function discards the background score and returns the mean of foreground scores.
'''
if len(pred.size()) == 4:
axes = (0, 2, 3)
intersect = sum_tensor(pred * y, axes, keepdim=False)
denom = sum_tensor(false_positive_weight*pred + y, axes, keepdim=False)
return torch.mean(( (2 * intersect + smooth) / (denom + smooth) )[1:]) # only fg dice here.
elif len(pred.size()) == 5:
axes = (0, 2, 3, 4)
intersect = sum_tensor(pred * y, axes, keepdim=False)
denom = sum_tensor(false_positive_weight*pred + y, axes, keepdim=False)
return torch.mean(( (2*intersect + smooth) / (denom + smooth) )[1:]) # only fg dice here.
else:
raise ValueError('wrong input dimension in dice loss')
def batch_dice_mask(pred, y, mask, false_positive_weight=1.0, smooth=1e-6):
'''
compute soft dice over batch. this is a diffrentiable score and can be used as a loss function.
only dice scores of foreground classes are returned, since training typically
does not benefit from explicit background optimization. Pixels of the entire batch are considered a pseudo-volume to compute dice scores of.
This way, single patches with missing foreground classes can not produce faulty gradients.
:param pred: (b, c, y, x, (z)), softmax probabilities (network output).
:param y: (b, c, y, x, (z)), one hote encoded segmentation mask.
:param false_positive_weight: float [0,1]. For weighting of imbalanced classes,
reduces the penalty for false-positive pixels. Can be beneficial sometimes in data with heavy fg/bg imbalances.
:return: soft dice score (float). This function discards the background score and returns the mean of foreground scores.
'''
mask = mask.unsqueeze(1).repeat(1, 2, 1, 1)
if len(pred.size()) == 4:
axes = (0, 2, 3)
intersect = sum_tensor(pred * y * mask, axes, keepdim=False)
denom = sum_tensor(false_positive_weight*pred * mask + y * mask, axes, keepdim=False)
return torch.mean(( (2*intersect + smooth) / (denom + smooth))[1:]) # only fg dice here.
elif len(pred.size()) == 5:
axes = (0, 2, 3, 4)
intersect = sum_tensor(pred * y, axes, keepdim=False)
denom = sum_tensor(false_positive_weight*pred + y, axes, keepdim=False)
return torch.mean(( (2*intersect + smooth) / (denom + smooth) )[1:]) # only fg dice here.
else:
raise ValueError('wrong input dimension in dice loss')
\ No newline at end of file