diff --git a/custom_extensions/nms/setup.py b/custom_extensions/nms/setup.py
index 911e616..94daa19 100644
--- a/custom_extensions/nms/setup.py
+++ b/custom_extensions/nms/setup.py
@@ -1,18 +1,19 @@
"""
Created at 07.11.19 19:12
@author: gregor
"""
import os, sys
from pathlib import Path
from setuptools import setup
from torch.utils import cpp_extension
dir_ = Path(os.path.dirname(sys.argv[0]))
setup(name='nms_extension',
ext_modules=[cpp_extension.CUDAExtension('nms_extension', [str(dir_/'src/nms_interface.cpp'), str(dir_/'src/nms.cu')])],
cmdclass={'build_ext': cpp_extension.BuildExtension}
- )
\ No newline at end of file
+ )
+
diff --git a/datasets/toy/configs.py b/datasets/toy/configs.py
index 8d702f5..6919307 100644
--- a/datasets/toy/configs.py
+++ b/datasets/toy/configs.py
@@ -1,495 +1,495 @@
#!/usr/bin/env python
# Copyright 2019 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
from collections import namedtuple
boxLabel = namedtuple('boxLabel', ["name", "color"])
Label = namedtuple("Label", ['id', 'name', 'shape', 'radius', 'color', 'regression', 'ambiguities', 'gt_distortion'])
binLabel = namedtuple("binLabel", ['id', 'name', 'color', 'bin_vals'])
class Configs(DefaultConfigs):
def __init__(self, server_env=None):
super(Configs, self).__init__(server_env)
#########################
# Prepro #
#########################
self.pp_rootdir = os.path.join('/mnt/HDD2TB/Documents/data/toy', "cyl1ps_dev")
self.pp_npz_dir = self.pp_rootdir+"_npz"
self.pre_crop_size = [320,320,8] #y,x,z; determines pp data shape (2D easily implementable, but only 3D for now)
self.min_2d_radius = 6 #in pixels
self.n_train_samples, self.n_test_samples = 80, 80
# not actually real one-hot encoding (ohe) but contains more info: roi-overlap only within classes.
self.pp_create_ohe_seg = False
self.pp_empty_samples_ratio = 0.1
self.pp_place_radii_mid_bin = True
self.pp_only_distort_2d = True
# outer-most intensity of blurred radii, relative to inner-object intensity. <1 for decreasing, > 1 for increasing.
# e.g.: setting 0.1 means blurred edge has min intensity 10% as large as inner-object intensity.
self.pp_blur_min_intensity = 0.2
self.max_instances_per_sample = 1 #how many max instances over all classes per sample (img if 2d, vol if 3d)
self.max_instances_per_class = self.max_instances_per_sample # how many max instances per image per class
self.noise_scale = 0. # std-dev of gaussian noise
self.ambigs_sampling = "gaussian" #"gaussian" or "uniform"
""" radius_calib: gt distort for calibrating uncertainty. Range of gt distortion is inferable from
image by distinguishing it from the rest of the object.
blurring width around edge will be shifted so that symmetric rel to orig radius.
blurring scale: if self.ambigs_sampling is uniform, distribution's non-zero range (b-a) will be sqrt(12)*scale
since uniform dist has variance (b-a)²/12. b,a will be placed symmetrically around unperturbed radius.
if sampling is gaussian, then scale parameter sets one std dev, i.e., blurring width will be orig_radius * std_dev * 2.
"""
self.ambiguities = {
#set which classes to apply which ambs to below in class labels
#choose out of: 'outer_radius', 'inner_radius', 'radii_relations'.
#kind #probability #scale (gaussian std, relative to unperturbed value)
#"outer_radius": (1., 0.5),
#"outer_radius_xy": (1., 0.5),
#"inner_radius": (0.5, 0.1),
#"radii_relations": (0.5, 0.1),
"radius_calib": (1., 1./6)
}
# shape choices: 'cylinder', 'block'
- self.pp_classes = [Label(1, 'cylinder', 'cylinder', ((6,6,1),(40,40,8)), (*self.blue, 1.), "radius_2d", (), ('radius_calib',)),
- #Label(2, 'block', 'block', ((6,6,1),(40,40,8)), (*self.aubergine,1.), "radii_2d", (), ('radius_calib',))
+ # id, name, shape, radius, color, regression, ambiguities, gt_distortion
+ self.pp_classes = [Label(1, 'cylinder', 'cylinder', ((6,6,1),(40,40,8)), (*self.blue, 1.), "radius_2d", (), ()),
+ #Label(2, 'block', 'block', ((6,6,1),(40,40,8)), (*self.aubergine,1.), "radii_2d", (), ('radius_calib',))
]
#########################
# I/O #
#########################
- #self.data_sourcedir = '/mnt/HDD2TB/Documents/data/toy/cyl1ps_dev'
- self.data_sourcedir = '/mnt/HDD2TB/Documents/data/toy/cyl1ps_exact'
- #self.data_sourcedir = '/mnt/HDD2TB/Documents/data/toy/cyl1ps_ambig_beyond_bin'
+ self.data_sourcedir = '/mnt/HDD2TB/Documents/data/toy/cyl1ps_dev'
+ #self.data_sourcedir = '/mnt/HDD2TB/Documents/data/toy/cyl1ps_exact'
if server_env:
#self.data_sourcedir = '/datasets/data_ramien/toy/cyl1ps_exact_npz'
self.data_sourcedir = '/datasets/data_ramien/toy/cyl1ps_ambig_beyond_bin_npz'
self.test_data_sourcedir = os.path.join(self.data_sourcedir, 'test')
self.data_sourcedir = os.path.join(self.data_sourcedir, "train")
self.info_df_name = 'info_df.pickle'
# one out of ['mrcnn', 'retina_net', 'retina_unet', 'detection_unet', 'ufrcnn', 'detection_fpn'].
- self.model = 'retina_net'
+ self.model = 'detection_unet'
self.model_path = 'models/{}.py'.format(self.model if not 'retina' in self.model else 'retina_net')
self.model_path = os.path.join(self.source_dir, self.model_path)
#########################
# Architecture #
#########################
# one out of [2, 3]. dimension the model operates in.
self.dim = 2
# 'class', 'regression', 'regression_bin', 'regression_ken_gal'
# currently only tested mode is a single-task at a time (i.e., only one task in below list)
# but, in principle, tasks could be combined (e.g., object classes and regression per class)
self.prediction_tasks = ['class',]
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 = 'instance_norm' # one of None, 'instance_norm', 'batch_norm'
self.relu = 'relu'
# one of 'xavier_uniform', 'xavier_normal', or 'kaiming_normal', None (=default = 'kaiming_uniform')
self.weight_init = None
self.regression_n_features = 1 # length of regressor target vector
#########################
# Data Loader #
#########################
self.num_epochs = 32
self.num_train_batches = 120 if self.dim == 2 else 80
self.batch_size = 16 if self.dim == 2 else 8
self.n_cv_splits = 4
# select modalities from preprocessed data
self.channels = [0]
self.n_channels = len(self.channels)
# which channel (mod) to show as bg in plotting, will be extra added to batch if not in self.channels
self.plot_bg_chan = 0
self.crop_margin = [20, 20, 1] # has to be smaller than respective patch_size//2
self.patch_size_2D = self.pre_crop_size[:2]
self.patch_size_3D = self.pre_crop_size[:2]+[8]
# patch_size to be used for training. pre_crop_size is the patch_size before data augmentation.
self.patch_size = self.patch_size_2D if self.dim == 2 else self.patch_size_3D
# ratio of free sampled batch elements before class balancing is triggered
# (>0 to include "empty"/background patches.)
self.batch_random_ratio = 0.2
self.balance_target = "class_targets" if 'class' in self.prediction_tasks else "rg_bin_targets"
self.observables_patient = []
self.observables_rois = []
self.seed = 3 #for generating folds
#############################
# Colors, Classes, Legends #
#############################
self.plot_frequency = 1
binary_bin_labels = [binLabel(1, 'r<=25', (*self.green, 1.), (1,25)),
binLabel(2, 'r>25', (*self.red, 1.), (25,))]
quintuple_bin_labels = [binLabel(1, 'r2-10', (*self.green, 1.), (2,10)),
binLabel(2, 'r10-20', (*self.yellow, 1.), (10,20)),
binLabel(3, 'r20-30', (*self.orange, 1.), (20,30)),
binLabel(4, 'r30-40', (*self.bright_red, 1.), (30,40)),
binLabel(5, 'r>40', (*self.red, 1.), (40,))]
# choose here if to do 2-way or 5-way regression-bin classification
task_spec_bin_labels = quintuple_bin_labels
self.class_labels = [
# regression: regression-task label, either value or "(x,y,z)_radius" or "radii".
# ambiguities: name of above defined ambig to apply to image data (not gt); need to be iterables!
# gt_distortion: name of ambig to apply to gt only; needs to be iterable!
# #id #name #shape #radius #color #regression #ambiguities #gt_distortion
Label( 0, 'bg', None, (0, 0, 0), (*self.white, 0.), (0, 0, 0), (), ())]
if "class" in self.prediction_tasks:
self.class_labels += self.pp_classes
else:
self.class_labels += [Label(1, 'object', 'object', ('various',), (*self.orange, 1.), ('radius_2d',), ("various",), ('various',))]
if any(['regression' in task for task in self.prediction_tasks]):
self.bin_labels = [binLabel(0, 'bg', (*self.white, 1.), (0,))]
self.bin_labels += task_spec_bin_labels
self.bin_id2label = {label.id: label for label in self.bin_labels}
bins = [(min(label.bin_vals), max(label.bin_vals)) for label in self.bin_labels]
self.bin_id2rg_val = {ix: [np.mean(bin)] for ix, bin in enumerate(bins)}
self.bin_edges = [(bins[i][1] + bins[i + 1][0]) / 2 for i in range(len(bins) - 1)]
self.bin_dict = {label.id: label.name for label in self.bin_labels if label.id != 0}
if self.class_specific_seg:
self.seg_labels = self.class_labels
self.box_type2label = {label.name: label for label in self.box_labels}
self.class_id2label = {label.id: label for label in self.class_labels}
self.class_dict = {label.id: label.name for label in self.class_labels if label.id != 0}
self.seg_id2label = {label.id: label for label in self.seg_labels}
self.cmap = {label.id: label.color for label in self.seg_labels}
self.plot_prediction_histograms = True
self.plot_stat_curves = False
self.has_colorchannels = False
self.plot_class_ids = True
self.num_classes = len(self.class_dict)
self.num_seg_classes = len(self.seg_labels)
#########################
# Data Augmentation #
#########################
self.do_aug = True
self.da_kwargs = {
'mirror': True,
'mirror_axes': tuple(np.arange(0, self.dim, 1)),
'do_elastic_deform': False,
'alpha': (500., 1500.),
'sigma': (40., 45.),
'do_rotation': False,
'angle_x': (0., 2 * np.pi),
'angle_y': (0., 0),
'angle_z': (0., 0),
'do_scale': False,
'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)
#########################
# Schedule / Selection #
#########################
# 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 = 220 # if 'all' iterates over entire val_set once.
if self.val_mode == 'val_sampling':
self.num_val_batches = 25 if self.dim==2 else 15
self.save_n_models = 2
self.min_save_thresh = 1 if self.dim == 2 else 1 # =wait time in epochs
if "class" in self.prediction_tasks:
self.model_selection_criteria = {name + "_ap": 1. for name in self.class_dict.values()}
elif any("regression" in task for task in self.prediction_tasks):
self.model_selection_criteria = {name + "_ap": 0.2 for name in self.class_dict.values()}
self.model_selection_criteria.update({name + "_avp": 0.8 for name in self.class_dict.values()})
self.lr_decay_factor = 0.5
self.scheduling_patience = int(self.num_epochs / 5)
self.weight_decay = 1e-5
self.clip_norm = None # number or None
#########################
# Testing / Plotting #
#########################
self.test_aug_axes = (0,1,(0,1)) # None or list: choices are 0,1,(0,1)
self.held_out_test_set = True
self.max_test_patients = "all" # number or "all" for all
self.test_against_exact_gt = not 'exact' in self.data_sourcedir
self.val_against_exact_gt = False # True is an unrealistic --> irrelevant scenario.
self.report_score_level = ['rois'] # 'patient' or 'rois' (incl)
self.patient_class_of_interest = 1
self.patient_bin_of_interest = 2
self.eval_bins_separately = False#"additionally" if not 'class' in self.prediction_tasks else False
self.metrics = ['ap', 'auc', 'dice']
if any(['regression' in task for task in self.prediction_tasks]):
self.metrics += ['avp', 'rg_MAE_weighted', 'rg_MAE_weighted_tp',
'rg_bin_accuracy_weighted', 'rg_bin_accuracy_weighted_tp']
if 'aleatoric' in self.model:
self.metrics += ['rg_uncertainty', 'rg_uncertainty_tp', 'rg_uncertainty_tp_weighted']
self.evaluate_fold_means = True
self.ap_match_ious = [0.5] # threshold(s) for considering a prediction as true positive
self.min_det_thresh = 0.3
self.model_max_iou_resolution = 0.2
# aggregation method for test and val_patient predictions.
# wbc = weighted box clustering as in https://arxiv.org/pdf/1811.08661.pdf,
# nms = standard non-maximum suppression, or None = no clustering
self.clustering = 'wbc'
# iou thresh (exclusive!) for regarding two preds as concerning the same ROI
self.clustering_iou = self.model_max_iou_resolution # has to be larger than desired possible overlap iou of model predictions
self.merge_2D_to_3D_preds = False
self.merge_3D_iou = self.model_max_iou_resolution
self.n_test_plots = 1 # per fold and rank
self.test_n_epochs = self.save_n_models # should be called n_test_ens, since is number of models to ensemble over during testing
# is multiplied by (1 + nr of test augs)
#self.losses_to_monitor += ['class_loss', 'rg_loss']
#########################
# Assertions #
#########################
if not 'class' in self.prediction_tasks:
assert self.num_classes == 1
#########################
# Add model specifics #
#########################
{'mrcnn': self.add_mrcnn_configs, 'mrcnn_aleatoric': self.add_mrcnn_configs,
'retina_net': self.add_mrcnn_configs, 'retina_unet': self.add_mrcnn_configs,
'detection_unet': self.add_det_unet_configs, 'detection_fpn': self.add_det_fpn_configs
}[self.model]()
def rg_val_to_bin_id(self, rg_val):
#only meant for isotropic radii!!
# only 2D radii (x and y dims) or 1D (x or y) are expected
return np.round(np.digitize(rg_val, self.bin_edges).mean())
def add_det_fpn_configs(self):
self.learning_rate = [5 * 1e-4] * self.num_epochs
self.dynamic_lr_scheduling = True
self.scheduling_criterion = 'torch_loss'
self.scheduling_mode = 'min' if "loss" in self.scheduling_criterion else 'max'
self.n_roi_candidates = 4 if self.dim == 2 else 6
# max number of roi candidates to identify per image (slice in 2D, volume in 3D)
# loss mode: either weighted cross entropy ('wce'), batch-wise dice loss ('dice), or the sum of both ('dice_wce')
self.seg_loss_mode = 'wce'
- self.wce_weights = [1] * self.num_seg_classes if 'dice' in self.seg_loss_mode else [0.1, 1, 1]
+ self.wce_weights = [1] * self.num_seg_classes if 'dice' in self.seg_loss_mode else [0.1, 1]
self.fp_dice_weight = 1 if self.dim == 2 else 1
# if <1, false positive predictions in foreground are penalized less.
self.detection_min_confidence = 0.05
# how to determine score of roi: 'max' or 'median'
self.score_det = 'max'
def add_det_unet_configs(self):
self.learning_rate = [5 * 1e-4] * self.num_epochs
self.dynamic_lr_scheduling = True
self.scheduling_criterion = "torch_loss"
self.scheduling_mode = 'min' if "loss" in self.scheduling_criterion else 'max'
# max number of roi candidates to identify per image (slice in 2D, volume in 3D)
self.n_roi_candidates = 4 if self.dim == 2 else 6
# loss mode: either weighted cross entropy ('wce'), batch-wise dice loss ('dice), or the sum of both ('dice_wce')
self.seg_loss_mode = 'wce'
- self.wce_weights = [1] * self.num_seg_classes if 'dice' in self.seg_loss_mode else [0.1, 1, 1]
+ self.wce_weights = [1] * self.num_seg_classes if 'dice' in self.seg_loss_mode else [0.1, 1]
# if <1, false positive predictions in foreground are penalized less.
self.fp_dice_weight = 1 if self.dim == 2 else 1
self.detection_min_confidence = 0.05
# how to determine score of roi: 'max' or 'median'
self.score_det = 'max'
self.init_filts = 32
self.kernel_size = 3 # ks for horizontal, normal convs
self.kernel_size_m = 2 # ks for max pool
self.pad = "same" # "same" or integer, padding of horizontal convs
def add_mrcnn_configs(self):
self.learning_rate = [1e-4] * self.num_epochs
self.dynamic_lr_scheduling = True # with scheduler set in exec
self.scheduling_criterion = max(self.model_selection_criteria, key=self.model_selection_criteria.get)
self.scheduling_mode = 'min' if "loss" in self.scheduling_criterion else 'max'
# number of classes for network heads: n_foreground_classes + 1 (background)
self.head_classes = self.num_classes + 1 if 'class' in self.prediction_tasks else 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)
self.detect_while_training = True
# 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_train = True
self.return_masks_in_val = True
self.return_masks_in_test = True
# feature map strides per pyramid level are inferred from architecture. anchor scales are set accordingly.
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': [[4], [8], [16], [32]], 'z': [[1], [2], [4], [8]]}
# 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 64
# 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 = max(0.8, self.model_max_iou_resolution)
# loss sampling settings.
self.rpn_train_anchors_per_image = 4
self.train_rois_per_image = 6 # per batch_instance
self.roi_positive_ratio = 0.5
self.anchor_matching_iou = 0.8
# k negative example candidates are drawn from a pool of size k*shem_poolsize (stochastic hard-example mining),
# where k<=#positive examples.
self.shem_poolsize = 2
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]]) # y1,x1,y2,x2,z1,z2
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]
self.plot_y_max = 1.5
self.n_plot_rpn_props = 5 if self.dim == 2 else 30 # per batch_instance (slice in 2D / patient in 3D)
# pre-selection in proposal-layer (stage 1) for NMS-speedup. applied per batch element.
self.pre_nms_limit = 2000 if self.dim == 2 else 4000
# 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 as one "batch".
self.roi_chunk_size = 1300 if self.dim == 2 else 500
self.post_nms_rois_training = 200 * (self.head_classes-1) if self.dim == 2 else 400
self.post_nms_rois_inference = 200 * (self.head_classes-1)
# Final selection of detections (refine_detections)
self.model_max_instances_per_batch_element = 9 if self.dim == 2 else 18 # per batch element and class.
self.detection_nms_threshold = self.model_max_iou_resolution # needs to be > 0, otherwise all predictions are one cluster.
self.model_min_confidence = 0.2 # iou for nms in box refining (directly after heads), should be >0 since ths>=x in mrcnn.py
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 == 'retina_net' or self.model == 'retina_unet':
# whether to use focal loss or SHEM for loss-sample selection
self.focal_loss = False
# implement extra anchor-scales according to https://arxiv.org/abs/1708.02002
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 = (500 if self.dim == 2 else 6250) * self.batch_size
# anchor matching iou is lower than in Mask R-CNN according to https://arxiv.org/abs/1708.02002
self.anchor_matching_iou = 0.7
if self.model == 'retina_unet':
self.operate_stride1 = True
diff --git a/datasets/toy/generate_toys.py b/datasets/toy/generate_toys.py
index 7d430d5..081eed1 100644
--- a/datasets/toy/generate_toys.py
+++ b/datasets/toy/generate_toys.py
@@ -1,388 +1,388 @@
#!/usr/bin/env python
# Copyright 2019 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.
# ==============================================================================
""" Generate a data set of toy examples. Examples can be cylinders, spheres, blocks, diamonds.
Distortions may be applied, e.g., noise to the radius ground truths.
Settings are configured in configs file.
"""
import plotting as plg
import os
import shutil
import warnings
import time
from multiprocessing import Pool
import numpy as np
import pandas as pd
import data_manager as dmanager
for msg in ["RuntimeWarning: divide by zero encountered in true_divide.*",]:
warnings.filterwarnings("ignore", msg)
class ToyGenerator(object):
""" Generator of toy data set.
A train and a test split with certain nr of samples are created and saved to disk. Samples can contain varying
number of objects. Objects have shapes cylinder or block (diamond, ellipsoid, torus not fully implemented).
self.mp_args holds image split and id, objects are then randomly drawn into each image. Multi-processing is
enabled for parallel creation of images, final .npy-files can then be converted to .npz.
"""
def __init__(self, cf):
"""
:param cf: configs file holding object specifications and output directories.
"""
self.cf = cf
self.n_train, self.n_test = cf.n_train_samples, cf.n_test_samples
self.sample_size = cf.pre_crop_size
self.dim = len(self.sample_size)
self.class_radii = np.array([label.radius for label in self.cf.pp_classes if label.id!=0])
self.class_id2label = {label.id: label for label in self.cf.pp_classes}
self.mp_args = []
# count sample ids consecutively over train, test splits within on dataset (one shape kind)
self.last_s_id = 0
for split in ["train", "test"]:
self.set_splits_info(split)
def set_splits_info(self, split):
""" Set info for data set splits, i.e., directory and nr of samples.
:param split: name of split, in {"train", "test"}.
"""
out_dir = os.path.join(self.cf.pp_rootdir, split)
os.makedirs(out_dir, exist_ok=True)
n_samples = self.n_train if "train" in split else self.n_test
self.mp_args+= [[out_dir, self.last_s_id+running_id] for running_id in range(n_samples)]
self.last_s_id+= n_samples
def generate_sample_radii(self, class_ids, shapes):
# the radii set in labels are ranges to sample from in the form [(min_x,min_y,min_z), (max_x,max_y,max_z)]
all_radii = []
for ix, cl_radii in enumerate([self.class_radii[cl_id - 1].transpose() for cl_id in class_ids]):
if "cylinder" in shapes[ix] or "block" in shapes[ix]:
# maintain 2D aspect ratio
sample_radii = [np.random.uniform(*cl_radii[0])] * 2
assert len(sample_radii) == 2, "upper sr {}, cl_radii {}".format(sample_radii, cl_radii)
if self.cf.pp_place_radii_mid_bin:
bef_conv_r = np.copy(sample_radii)
bin_id = self.cf.rg_val_to_bin_id(bef_conv_r)
assert np.isscalar(bin_id)
sample_radii = self.cf.bin_id2rg_val[bin_id]*2
assert len(sample_radii) == 2, "mid before sr {}, sr {}, rgv2bid {}, cl_radii {}, bid2rgval {}".format(bef_conv_r, sample_radii, bin_id, cl_radii,
self.cf.bin_id2rg_val[bin_id])
else:
raise NotImplementedError("requested object shape {}".format(shapes[ix]))
if self.dim == 3:
assert len(sample_radii) == 2, "lower sr {}, cl_radii {}".format(sample_radii, cl_radii)
#sample_radii += [np.random.uniform(*cl_radii[2])]
sample_radii = np.concatenate((sample_radii, np.random.uniform(*cl_radii[2], size=1)))
all_radii.append(sample_radii)
return all_radii
def apply_gt_distort(self, class_id, radii, radii_divs, outer_min_radii=None, outer_max_radii=None):
""" Apply a distortion to the ground truth (gt). This is motivated by investigating the effects of noisy labels.
GTs that can be distorted are the object radii and ensuing GT quantities like segmentation and regression
targets.
:param class_id: class id of object.
:param radii: radii of object. This is in the abstract sense, s.t. for a block-shaped object radii give the side
lengths.
:param radii_divs: radii divisors, i.e., fractions to take from radii to get inner radii of hole-shaped objects,
like a torus.
:param outer_min_radii: min radii assignable when distorting gt.
:param outer_max_radii: max radii assignable when distorting gt.
:return:
"""
applied_gt_distort = False
for ambig in self.class_id2label[class_id].gt_distortion:
if self.cf.ambiguities[ambig][0] > np.random.rand():
if ambig == "outer_radius":
radii = radii * abs(np.random.normal(1., self.cf.ambiguities["outer_radius"][1]))
applied_gt_distort = True
if ambig == "radii_relations":
radii = radii * abs(np.random.normal(1.,self.cf.ambiguities["radii_relations"][1],size=len(radii)))
applied_gt_distort = True
if ambig == "inner_radius":
radii_divs = radii_divs * abs(np.random.normal(1., self.cf.ambiguities["inner_radius"][1]))
applied_gt_distort = True
if ambig == "radius_calib":
if self.cf.ambigs_sampling=="uniform":
radii = abs(np.random.uniform(outer_min_radii, outer_max_radii))
elif self.cf.ambigs_sampling=="gaussian":
distort = abs(np.random.normal(1, scale=self.cf.ambiguities["radius_calib"][1], size=None))
assert len(radii) == self.dim, "radii {}".format(radii)
radii *= [distort, distort, 1.] if self.cf.pp_only_distort_2d else distort
applied_gt_distort = True
return radii, radii_divs, applied_gt_distort
def draw_object(self, img, seg, undistorted_seg, ics, regress_targets, undistorted_rg_targets, applied_gt_distort,
roi_ix, class_id, shape, radii, center):
""" Draw a single object into the given image and add it to the corresponding ground truths.
:param img: image (volume) to hold the object.
:param seg: pixel-wise labelling of the image, possibly distorted if gt distortions are applied.
:param undistorted_seg: certainly undistorted, i.e., exact segmentation of object.
:param ics: indices which mark the positions within the image.
:param regress_targets: regression targets (e.g., 2D radii of object), evtly distorted.
:param undistorted_rg_targets: undistorted regression targets.
:param applied_gt_distort: boolean, whether or not gt distortion was applied.
:param roi_ix: running index of object in whole image.
:param class_id: class id of object.
:param shape: shape of object (e.g., whether to draw a cylinder, or block, or ...).
:param radii: radii of object (in an abstract sense, i.e., radii are side lengths in case of block shape).
:param center: center of object in image coordinates.
:return: img, seg, undistorted_seg, regress_targets, undistorted_rg_targets, applied_gt_distort, which are now
extended are amended to reflect the new object.
"""
radii_blur = hasattr(self.cf, "ambiguities") and hasattr(self.class_id2label[class_id],
"gt_distortion") and 'radius_calib' in \
self.class_id2label[class_id].gt_distortion
if radii_blur:
blur_width = self.cf.ambiguities['radius_calib'][1]
if self.cf.ambigs_sampling == "uniform":
blur_width *= np.sqrt(12)
if self.cf.pp_only_distort_2d:
outer_max_radii = np.concatenate((radii[:2] + blur_width * radii[:2], [radii[2]]))
outer_min_radii = np.concatenate((radii[:2] - blur_width * radii[:2], [radii[2]]))
#print("belt width ", outer_max_radii - outer_min_radii)
else:
outer_max_radii = radii + blur_width * radii
outer_min_radii = radii - blur_width * radii
else:
outer_max_radii, outer_min_radii = radii, radii
if "ellipsoid" in shape or "torus" in shape:
# sphere equation: (x-h)**2 + (y-k)**2 - (z-l)**2 = r**2
# ellipsoid equation: ((x-h)/a)**2+((y-k)/b)**2+((z-l)/c)**2 <= 1; a, b, c the "radii"/ half-length of principal axes
obj = ((ics - center) / radii) ** 2
elif "diamond" in shape:
# diamond equation: (|x-h|)/a+(|y-k|)/b+(|z-l|)/c <= 1
obj = abs(ics - center) / radii
elif "cylinder" in shape:
# cylinder equation:((x-h)/a)**2 + ((y-k)/b)**2 <= 1 while |z-l| <= c
obj = ((ics - center).astype("float64") / radii) ** 2
# set z values s.t. z slices outside range are sorted out
obj[:, -1] = np.where(abs((ics - center)[:, -1]) <= radii[2], 0., 1.1)
if radii_blur:
inner_obj = ((ics - center).astype("float64") / outer_min_radii) ** 2
inner_obj[:, -1] = np.where(abs((ics - center)[:, -1]) <= outer_min_radii[2], 0., 1.1)
outer_obj = ((ics - center).astype("float64") / outer_max_radii) ** 2
outer_obj[:, -1] = np.where(abs((ics - center)[:, -1]) <= outer_max_radii[2], 0., 1.1)
# radial dists: sqrt( (x-h)**2 + (y-k)**2 + (z-l)**2 )
obj_radial_dists = np.sqrt(np.sum((ics - center).astype("float64")**2, axis=1))
elif "block" in shape:
# block equation: (|x-h|)/a+(|y-k|)/b <= 1 while |z-l| <= c
obj = abs(ics - center) / radii
obj[:, -1] = np.where(abs((ics - center)[:, -1]) <= radii[2], 0., 1.1)
if radii_blur:
inner_obj = abs(ics - center) / outer_min_radii
inner_obj[:, -1] = np.where(abs((ics - center)[:, -1]) <= outer_min_radii[2], 0., 1.1)
outer_obj = abs(ics - center) / outer_max_radii
outer_obj[:, -1] = np.where(abs((ics - center)[:, -1]) <= outer_max_radii[2], 0., 1.1)
obj_radial_dists = np.sum(abs(ics - center), axis=1).astype("float64")
else:
raise Exception("Invalid object shape '{}'".format(shape))
# create the "original" GT, i.e., the actually true object and draw it into undistorted seg.
obj = (np.sum(obj, axis=1) <= 1)
obj = obj.reshape(seg[0].shape)
slices_to_discard = np.where(np.count_nonzero(np.count_nonzero(obj, axis=0), axis=0) <= self.cf.min_2d_radius)[0]
obj[..., slices_to_discard] = 0
undistorted_radii = np.copy(radii)
undistorted_seg[class_id][obj] = roi_ix + 1
obj = obj.astype('float64')
if radii_blur:
inner_obj = np.sum(inner_obj, axis=1) <= 1
outer_obj = (np.sum(outer_obj, axis=1) <= 1) & ~inner_obj
obj_radial_dists[outer_obj] = obj_radial_dists[outer_obj] / max(obj_radial_dists[outer_obj])
intensity_slope = self.cf.pp_blur_min_intensity - 1.
# intensity(r) = (i(r_max)-i(0))/r_max * r + i(0), where i(0)==1.
obj_radial_dists[outer_obj] = obj_radial_dists[outer_obj] * intensity_slope + 1.
inner_obj = inner_obj.astype('float64')
#outer_obj, obj_radial_dists = outer_obj.reshape(seg[0].shape), obj_radial_dists.reshape(seg[0].shape)
inner_obj += np.where(outer_obj, obj_radial_dists, 0.)
obj = inner_obj.reshape(seg[0].shape)
if not np.any(obj):
print("An object was completely discarded due to min 2d radius requirement, discarded slices: {}.".format(
slices_to_discard))
# draw the evtly blurred obj into image.
img += obj * (class_id + 1.)
if hasattr(self.cf, "ambiguities") and hasattr(self.class_id2label[class_id], "gt_distortion"):
radii_divs = [None] # dummy since not implemented yet
radii, radii_divs, applied_gt_distort = self.apply_gt_distort(class_id, radii, radii_divs,
outer_min_radii, outer_max_radii)
if applied_gt_distort:
if "ellipsoid" in shape or "torus" in shape:
obj = ((ics - center) / radii) ** 2
elif 'diamond' in shape:
obj = abs(ics - center) / radii
elif "cylinder" in shape:
obj = ((ics - center) / radii) ** 2
obj[:, -1] = np.where(abs((ics - center)[:, -1]) <= radii[2], 0., 1.1)
elif "block" in shape:
obj = abs(ics - center) / radii
obj[:, -1] = np.where(abs((ics - center)[:, -1]) <= radii[2], 0., 1.1)
obj = (np.sum(obj, axis=1) <= 1).reshape(seg[0].shape)
obj[..., slices_to_discard] = False
if self.class_id2label[class_id].regression == "radii":
regress_targets.append(radii)
undistorted_rg_targets.append(undistorted_radii)
elif self.class_id2label[class_id].regression == "radii_2d":
regress_targets.append(radii[:2])
undistorted_rg_targets.append(undistorted_radii[:2])
elif self.class_id2label[class_id].regression == "radius_2d":
regress_targets.append(radii[:1])
undistorted_rg_targets.append(undistorted_radii[:1])
else:
regress_targets.append(self.class_id2label[class_id].regression)
undistorted_rg_targets.append(self.class_id2label[class_id].regression)
seg[class_id][obj.astype('bool')] = roi_ix + 1
return img, seg, undistorted_seg, regress_targets, undistorted_rg_targets, applied_gt_distort
def create_sample(self, args):
""" Create a single sample and save to file. One sample is one image (volume) containing none, one, or multiple
objects.
:param args: out_dir: directory where to save sample, s_id: id of the sample.
:return: specs that identify this single created image
"""
out_dir, s_id = args
print('processing {} {}'.format(out_dir, s_id))
img = np.random.normal(loc=0.0, scale=self.cf.noise_scale, size=self.sample_size)
img[img<0.] = 0.
# one-hot-encoded seg
seg = np.zeros((self.cf.num_classes+1, *self.sample_size)).astype('uint8')
undistorted_seg = np.copy(seg)
applied_gt_distort = False
if hasattr(self.cf, "pp_empty_samples_ratio") and self.cf.pp_empty_samples_ratio >= np.random.rand():
# generate fully empty sample
class_ids, regress_targets, undistorted_rg_targets = [], [], []
else:
class_choices = np.repeat(np.arange(1, self.cf.num_classes+1), self.cf.max_instances_per_class)
n_insts = np.random.randint(1, self.cf.max_instances_per_sample + 1)
class_ids = np.random.choice(class_choices, size=n_insts, replace=False)
shapes = np.array([self.class_id2label[cl_id].shape for cl_id in class_ids])
all_radii = self.generate_sample_radii(class_ids, shapes)
# reorder s.t. larger objects are drawn first (in order to not fully cover smaller objects)
order = np.argsort(-1*np.prod(all_radii,axis=1))
class_ids = class_ids[order]; all_radii = np.array(all_radii)[order]; shapes = shapes[order]
regress_targets, undistorted_rg_targets = [], []
# indices ics equal positions within img/volume
ics = np.argwhere(np.ones(seg[0].shape))
for roi_ix, class_id in enumerate(class_ids):
radii = all_radii[roi_ix]
# enforce distance between object center and image edge relative to radii.
margin_r_divisor = (2, 2, 4)
center = [np.random.randint(radii[dim] / margin_r_divisor[dim], img.shape[dim] -
radii[dim] / margin_r_divisor[dim]) for dim in range(len(img.shape))]
img, seg, undistorted_seg, regress_targets, undistorted_rg_targets, applied_gt_distort = \
self.draw_object(img, seg, undistorted_seg, ics, regress_targets, undistorted_rg_targets, applied_gt_distort,
roi_ix, class_id, shapes[roi_ix], radii, center)
fg_slices = np.where(np.sum(np.sum(np.sum(seg,axis=0), axis=0), axis=0))[0]
if self.cf.pp_create_ohe_seg:
img = img[np.newaxis]
else:
# choosing rois to keep by smaller radius==higher prio needs to be ensured during roi generation,
# smaller objects need to be drawn later (==higher roi id)
seg = seg.max(axis=0)
seg_ids = np.unique(seg)
if len(seg_ids) != len(class_ids) + 1:
# in this case an object was completely covered by a succeeding object
print("skipping corrupt sample")
print("seg ids {}, class_ids {}".format(seg_ids, class_ids))
return None
if not applied_gt_distort:
assert np.all(np.flatnonzero(img>0) == np.flatnonzero(seg>0))
assert np.all(np.array(regress_targets).flatten()==np.array(undistorted_rg_targets).flatten())
out_path = os.path.join(out_dir, '{}.npy'.format(s_id))
np.save(out_path, img.astype('float16')); np.save(os.path.join(out_dir, '{}_seg.npy'.format(s_id)), seg)
if hasattr(self.cf, 'ambiguities') and \
np.any([hasattr(label, "gt_distortion") and len(label.gt_distortion)>0 for label in self.class_id2label.values()]):
undist_out_path = os.path.join(out_dir, '{}_exact_seg.npy'.format(s_id))
if not self.cf.pp_create_ohe_seg:
undistorted_seg = undistorted_seg.max(axis=0)
np.save(undist_out_path, undistorted_seg)
return [out_dir, out_path, class_ids, regress_targets, fg_slices, undistorted_rg_targets, str(s_id)]
def create_sets(self, processes=os.cpu_count()):
""" Create whole training and test set, save to files under given directory cf.out_dir.
:param processes: nr of parallel processes.
"""
print('starting creation of {} images'.format(len(self.mp_args)))
shutil.copyfile("configs.py", os.path.join(self.cf.pp_rootdir, 'applied_configs.py'))
pool = Pool(processes=processes)
imgs_info = pool.map(self.create_sample, self.mp_args, chunksize=1)
pool.close()
pool.join()
imgs_info = [img for img in imgs_info if img is not None]
print("created a total of {} samples.".format(len(imgs_info)))
self.df = pd.DataFrame.from_records(imgs_info, columns=['out_dir', 'path', 'class_ids', 'regression_vectors',
'fg_slices', 'undistorted_rg_vectors', 'pid'])
for out_dir, group_df in self.df.groupby("out_dir"):
group_df.to_pickle(os.path.join(out_dir, 'info_df.pickle'))
def convert_copy_npz(self):
""" Convert a copy of generated .npy-files to npz and save in .npz-directory given in configs.
"""
if hasattr(self.cf, "pp_npz_dir") and self.cf.pp_npz_dir:
for out_dir, group_df in self.df.groupby("out_dir"):
rel_dir = os.path.relpath(out_dir, self.cf.pp_rootdir).split(os.sep)
npz_out_dir = os.path.join(self.cf.pp_npz_dir, str(os.sep).join(rel_dir))
print("npz out dir: ", npz_out_dir)
os.makedirs(npz_out_dir, exist_ok=True)
group_df.to_pickle(os.path.join(npz_out_dir, 'info_df.pickle'))
dmanager.pack_dataset(out_dir, npz_out_dir, recursive=True, verbose=False)
else:
print("Did not convert .npy-files to .npz because npz directory not set in configs.")
if __name__ == '__main__':
import configs as cf
- cf = cf.configs()
+ cf = cf.Configs()
total_stime = time.time()
toy_gen = ToyGenerator(cf)
toy_gen.create_sets()
toy_gen.convert_copy_npz()
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))
diff --git a/default_configs.py b/default_configs.py
index c2d16e2..e68a3ed 100644
--- a/default_configs.py
+++ b/default_configs.py
@@ -1,202 +1,203 @@
#!/usr/bin/env python
# Copyright 2019 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.
# ==============================================================================
"""Default Configurations script. Avoids changing configs of all experiments if general settings are to be changed."""
import os
from collections import namedtuple
boxLabel = namedtuple('boxLabel', ["name", "color"])
class DefaultConfigs:
def __init__(self, server_env=None, dim=2):
self.server_env = server_env
self.cuda_benchmark = True
+ self.sysmetrics_interval = -1 # set > 0 to record system metrics to tboard with this time span in seconds.
#########################
# I/O #
#########################
self.dim = dim
# int [0 < dataset_size]. select n patients from dataset for prototyping.
self.select_prototype_subset = None
# some default paths.
self.source_dir = os.path.dirname(os.path.realpath(__file__)) # current dir.
self.backbone_path = os.path.join(self.source_dir, 'models/backbone.py')
self.input_df_name = 'info_df.pickle'
if server_env:
self.select_prototype_subset = None
#########################
# Colors/legends #
#########################
# in part from solarized theme.
self.black = (0.1, 0.05, 0.)
self.gray = (0.514, 0.580, 0.588)
self.beige = (1., 1., 0.85)
self.white = (0.992, 0.965, 0.890)
self.green = (0.659, 0.792, 0.251) # [168, 202, 64]
self.dark_green = (0.522, 0.600, 0.000) # [133.11, 153. , 0. ]
self.cyan = (0.165, 0.631, 0.596) # [ 42.075, 160.905, 151.98 ]
self.bright_blue = (0.85, 0.95, 1.)
self.blue = (0.149, 0.545, 0.824) # [ 37.995, 138.975, 210.12 ]
self.dkfz_blue = (0, 75. / 255, 142. / 255)
self.dark_blue = (0.027, 0.212, 0.259) # [ 6.885, 54.06 , 66.045]
self.purple = (0.424, 0.443, 0.769) # [108.12 , 112.965, 196.095]
self.aubergine = (0.62, 0.21, 0.44) # [ 157, 53 , 111]
self.magenta = (0.827, 0.212, 0.510) # [210.885, 54.06 , 130.05 ]
self.coral = (1., 0.251, 0.4) # [255,64,102]
self.bright_red = (1., 0.15, 0.1) # [255, 38.25, 25.5]
self.brighter_red = (0.863, 0.196, 0.184) # [220.065, 49.98 , 46.92 ]
self.red = (0.87, 0.05, 0.01) # [ 223, 13, 2]
self.dark_red = (0.6, 0.04, 0.005)
self.orange = (0.91, 0.33, 0.125) # [ 232.05 , 84.15 , 31.875]
self.dark_orange = (0.796, 0.294, 0.086) #[202.98, 74.97, 21.93]
self.yellow = (0.95, 0.9, 0.02) # [ 242.25, 229.5 , 5.1 ]
self.dark_yellow = (0.710, 0.537, 0.000) # [181.05 , 136.935, 0. ]
self.color_palette = [self.blue, self.dark_blue, self.aubergine, self.green, self.yellow, self.orange, self.red,
self.cyan, self.black]
self.box_labels = [
# name color
boxLabel("det", self.blue),
boxLabel("prop", self.gray),
boxLabel("pos_anchor", self.cyan),
boxLabel("neg_anchor", self.cyan),
boxLabel("neg_class", self.green),
boxLabel("pos_class", self.aubergine),
boxLabel("gt", self.red)
] # neg and pos in a medical sense, i.e., pos=positive diagnostic finding
self.box_type2label = {label.name: label for label in self.box_labels}
self.box_color_palette = {label.name: label.color for label in self.box_labels}
# whether the input data is mono-channel or RGB/rgb
self.has_colorchannels = False
#########################
# Data Loader #
#########################
#random seed for fold_generator and batch_generator.
self.seed = 0
#number of threads for multithreaded tasks like batch generation, wcs, merge2dto3d
self.n_workers = 16 if server_env else os.cpu_count()
self.create_bounding_box_targets = True
self.class_specific_seg = True # False if self.model=="mrcnn" else True
#########################
# Architecture #
#########################
self.prediction_tasks = ["class"] # 'class', 'regression_class', 'regression_kendall', 'regression_feindt'
self.weight_decay = 0.0
# nonlinearity to be applied after convs with nonlinearity. one of 'relu' or 'leaky_relu'
self.relu = 'relu'
# if True initializes weights as specified in model script. else use default Pytorch init.
self.weight_init = None
# if True adds high-res decoder levels to feature pyramid: P1 + P0. (e.g. set to true in retina_unet configs)
self.operate_stride1 = False
#########################
# Optimization #
#########################
self.optimizer = "ADAM" # "ADAM" or "SGD" or implemented additionals
#########################
# Schedule #
#########################
# number of folds in cross validation.
self.n_cv_splits = 5
#########################
# Testing / Plotting #
#########################
# perform mirroring at test time. (only XY. Z not done to not blow up predictions times).
self.test_aug = True
# if True, test data lies in a separate folder and is not part of the cross validation.
self.held_out_test_set = False
# if hold-out test set: eval each fold's parameters separately on the test set
self.eval_test_fold_wise = True
# if held_out_test_set provided, ensemble predictions over models of all trained cv-folds.
self.ensemble_folds = False
# what metrics to evaluate
self.metrics = ['ap']
# whether to evaluate fold means when evaluating over more than one fold
self.evaluate_fold_means = False
# how often (in nr of epochs) to plot example batches during train/val
self.plot_frequency = 1
# color specifications for all box_types in prediction_plot.
self.box_color_palette = {'det': 'b', 'gt': 'r', 'neg_class': 'purple',
'prop': 'w', 'pos_class': 'g', 'pos_anchor': 'c', 'neg_anchor': 'c'}
# scan over confidence score in evaluation to optimize it on the validation set.
self.scan_det_thresh = False
# plots roc-curves / prc-curves in evaluation.
self.plot_stat_curves = False
# if True: evaluate average precision per patient id and average over per-pid results,
# instead of computing one ap over whole data set.
self.per_patient_ap = False
# threshold for clustering 2D box predictions to 3D Cubes. Overlap is computed in XY.
self.merge_3D_iou = 0.1
#########################
# MRCNN #
#########################
# if True, mask loss is not applied. used for data sets, where no pixel-wise annotations are provided.
self.frcnn_mode = False
self.return_masks_in_train = False
# if True, unmolds masks in Mask R-CNN to full-res for plotting/monitoring.
self.return_masks_in_val = False
self.return_masks_in_test = False # needed if doing instance segmentation. evaluation not yet implemented.
# add P6 to Feature Pyramid Network.
self.sixth_pooling = False
#########################
# RetinaNet #
#########################
self.focal_loss = False
self.focal_loss_gamma = 2.
diff --git a/exec.py b/exec.py
index f02d8ab..0c03071 100644
--- a/exec.py
+++ b/exec.py
@@ -1,350 +1,344 @@
#!/usr/bin/env python
# Copyright 2019 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. this where all routines come together and the only script you need to call.
refer to parse args below to see options for execution.
"""
import plotting as plg
import os
import warnings
import argparse
import time
import torch
import utils.exp_utils as utils
from evaluator import Evaluator
from predictor import Predictor
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 true_divide.*"]:
warnings.filterwarnings("ignore", msg)
def train(cf, logger):
"""
performs the training routine for a given fold. saves plots and selected parameters to the experiment dir
- specified in the configs.
-
+ specified in the configs. logs to file and tensorboard.
"""
logger.info('performing training in {}D over fold {} on experiment {} with model {}'.format(
cf.dim, cf.fold, cf.exp_dir, cf.model))
logger.time("train_val")
# -------------- inits and settings -----------------
net = model.net(cf, logger).cuda()
if cf.optimizer == "ADAM":
optimizer = torch.optim.Adam(net.parameters(), lr=cf.learning_rate[0], weight_decay=cf.weight_decay)
elif cf.optimizer == "SGD":
optimizer = torch.optim.SGD(net.parameters(), lr=cf.learning_rate[0], weight_decay=cf.weight_decay, momentum=0.3)
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)
starting_epoch = 1
if cf.resume_from_checkpoint:
starting_epoch = utils.load_checkpoint(cf.resume_from_checkpoint, net, optimizer)
logger.info('resumed from checkpoint {} at epoch {}'.format(cf.resume_from_checkpoint, starting_epoch))
# prepare monitoring
monitor_metrics = utils.prepare_monitoring(cf)
logger.info('loading dataset and initializing batch generators...')
batch_gen = data_loader.get_train_generators(cf, logger)
# -------------- training -----------------
for epoch in range(starting_epoch, cf.num_epochs + 1):
logger.info('starting training epoch {}/{}'.format(epoch, cf.num_epochs))
logger.time("train_epoch")
net.train()
train_results_list = []
train_evaluator = Evaluator(cf, logger, mode='train')
for i in range(cf.num_train_batches):
logger.time("train_batch_loadfw")
batch = next(batch_gen['train'])
batch_gen['train'].generator.stats['roi_counts'] += batch['roi_counts']
batch_gen['train'].generator.stats['empty_samples_count'] += batch['empty_samples_count']
logger.time("train_batch_loadfw")
logger.time("train_batch_netfw")
results_dict = net.train_forward(batch)
logger.time("train_batch_netfw")
logger.time("train_batch_bw")
optimizer.zero_grad()
results_dict['torch_loss'].backward()
if cf.clip_norm:
- torch.nn.utils.clip_grad_norm_(net.parameters(), cf.clip_norm, norm_type=2) #gradient clipping
+ torch.nn.utils.clip_grad_norm_(net.parameters(), cf.clip_norm, norm_type=2) # gradient clipping
optimizer.step()
- train_results_list.append(({k:v for k,v in results_dict.items() if k != "seg_preds"}, batch["pid"])) #slim res dict
+ train_results_list.append(({k:v for k,v in results_dict.items() if k != "seg_preds"}, batch["pid"])) # slim res dict
if not cf.server_env:
print("\rFinished training batch " +
"{}/{} in {:.1f}s ({:.2f}/{:.2f} forw load/net, {:.2f} backw).".format(i+1, cf.num_train_batches,
logger.get_time("train_batch_loadfw")+
logger.get_time("train_batch_netfw")
+logger.time("train_batch_bw"),
logger.get_time("train_batch_loadfw",reset=True),
logger.get_time("train_batch_netfw", reset=True),
logger.get_time("train_batch_bw", reset=True)), end="", flush=True)
print()
#--------------- train eval ----------------
if (epoch-1)%cf.plot_frequency==0:
# view an example batch
logger.time("train_plot")
plg.view_batch(cf, batch, results_dict, has_colorchannels=cf.has_colorchannels, show_gt_labels=True,
out_file=os.path.join(cf.plot_dir, 'batch_example_train_{}.png'.format(cf.fold)))
- logger.info("generated train-example plot in {:.2f}s".format(logger.get_time("train_plot", reset=True)))
+ logger.info("generated train-example plot in {:.2f}s".format(logger.time("train_plot")))
logger.time("evals")
_, monitor_metrics['train'] = train_evaluator.evaluate_predictions(train_results_list, monitor_metrics['train'])
- #np_loss, torch_loss = train_loss_running_mean / cf.num_train_batches, monitor_metrics['train']["loss"][-1]
- #assert np_loss/torch_loss-1<0.005, "{} vs {}".format(np_loss, torch_loss)
logger.time("evals")
logger.time("train_epoch", toggle=False)
del train_results_list
+
#----------- validation ------------
logger.info('starting validation in mode {}.'.format(cf.val_mode))
logger.time("val_epoch")
with torch.no_grad():
net.eval()
val_results_list = []
val_evaluator = Evaluator(cf, logger, mode=cf.val_mode)
val_predictor = Predictor(cf, net, logger, mode='val')
for i in range(batch_gen['n_val']):
logger.time("val_batch")
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, batch["pid"]])
if not cf.server_env:
print("\rFinished validation {} {}/{} in {:.1f}s.".format('patient' if cf.val_mode=='val_patient' else 'batch',
i + 1, batch_gen['n_val'],
logger.time("val_batch")), end="", flush=True)
print()
#------------ val eval -------------
- logger.time("val_plot")
if (epoch - 1) % cf.plot_frequency == 0:
+ logger.time("val_plot")
plg.view_batch(cf, batch, results_dict, has_colorchannels=cf.has_colorchannels, show_gt_labels=True,
out_file=os.path.join(cf.plot_dir, 'batch_example_val_{}.png'.format(cf.fold)))
- logger.time("val_plot")
+ logger.info("generated val plot in {:.2f}s".format(logger.time("val_plot")))
logger.time("evals")
_, monitor_metrics['val'] = val_evaluator.evaluate_predictions(val_results_list, monitor_metrics['val'])
model_selector.run_model_selection(net, optimizer, monitor_metrics, epoch)
del val_results_list
#----------- monitoring -------------
monitor_metrics.update({"lr":
{str(g) : group['lr'] for (g, group) in enumerate(optimizer.param_groups)}})
logger.metrics2tboard(monitor_metrics, global_step=epoch)
logger.time("evals")
logger.info('finished epoch {}/{}, took {:.2f}s. train total: {:.2f}s, average: {:.2f}s. val total: {:.2f}s, average: {:.2f}s.'.format(
epoch, cf.num_epochs, logger.get_time("train_epoch")+logger.time("val_epoch"), logger.get_time("train_epoch"),
logger.get_time("train_epoch", reset=True)/cf.num_train_batches, logger.get_time("val_epoch"),
logger.get_time("val_epoch", reset=True)/batch_gen["n_val"]))
- logger.info("time for evals: {:.2f}s, val plot {:.2f}s".format(logger.get_time("evals", reset=True), logger.get_time("val_plot", reset=True)))
+ logger.info("time for evals: {:.2f}s".format(logger.get_time("evals", reset=True)))
#-------------- scheduling -----------------
if not cf.dynamic_lr_scheduling:
for param_group in optimizer.param_groups:
param_group['lr'] = cf.learning_rate[epoch-1]
else:
scheduler.step(monitor_metrics["val"][cf.scheduling_criterion][-1])
logger.time("train_val")
logger.info("Training and validating over {} epochs took {}".format(cf.num_epochs, logger.get_time("train_val", format="hms", reset=True)))
batch_gen['train'].generator.print_stats(logger, plot=True)
def test(cf, logger, max_fold=None):
"""performs testing for a given fold (or held out set). saves stats in evaluator.
"""
logger.time("test_fold")
logger.info('starting testing model of fold {} in exp {}'.format(cf.fold, cf.exp_dir))
net = model.net(cf, logger).cuda()
batch_gen = data_loader.get_test_generator(cf, logger)
test_predictor = Predictor(cf, net, logger, mode='test')
test_results_list = test_predictor.predict_test_set(batch_gen, return_results = not hasattr(
cf, "eval_test_separately") or not cf.eval_test_separately)
if test_results_list is not None:
test_evaluator = Evaluator(cf, logger, mode='test')
test_evaluator.evaluate_predictions(test_results_list)
test_evaluator.score_test_df(max_fold=max_fold)
- mins, secs = divmod(logger.get_time("test_fold"), 60)
- h, mins = divmod(mins, 60)
- t = "{:d}h:{:02d}m:{:02d}s".format(int(h), int(mins), int(secs))
-
- logger.info('Testing of fold {} took {}.'.format(cf.fold, t))
+ logger.info('Testing of fold {} took {}.'.format(cf.fold, logger.get_time("test_fold", reset=True, format="hms")))
if __name__ == '__main__':
stime = time.time()
parser = argparse.ArgumentParser()
parser.add_argument('-m', '--mode', type=str, default='train_test', help='one out of: create_exp, analysis, train, train_test, or test')
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='/home/gregor/Documents/regrcnn/datasets/toy/experiments/dev',
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_from_checkpoint', type=str, default=None,
help='path to checkpoint. if resuming from checkpoint, the desired fold still needs to be parsed via --folds.')
parser.add_argument('--dataset_name', type=str, default='toy', help="path to the dataset-specific code in source_dir/datasets")
parser.add_argument('-d', '--dev', default=False, action='store_true', help="development mode: shorten everything")
args = parser.parse_args()
args.dataset_name = os.path.join("datasets", args.dataset_name) if not "datasets" in args.dataset_name else args.dataset_name
folds = args.folds
resume_from_checkpoint = None if args.resume_from_checkpoint in ['None', 'none'] else args.resume_from_checkpoint
if args.mode == 'create_exp':
cf = utils.prep_exp(args.dataset_name, args.exp_dir, args.server_env, use_stored_settings=False)
- logger = utils.get_logger(cf.exp_dir, cf.server_env)
+ logger = utils.get_logger(cf.exp_dir, cf.server_env, -1)
logger.info('created experiment directory at {}'.format(args.exp_dir))
elif args.mode == 'train' or args.mode == 'train_test':
cf = utils.prep_exp(args.dataset_name, 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
torch.backends.cudnn.benchmark = cf.dim==3
else:
torch.backends.cudnn.benchmark = cf.cuda_benchmark
if args.data_dest is not None:
cf.data_dest = args.data_dest
- logger = utils.get_logger(cf.exp_dir, cf.server_env)
+ logger = utils.get_logger(cf.exp_dir, cf.server_env, cf.sysmetrics_interval)
data_loader = utils.import_module('data_loader', os.path.join(args.dataset_name, '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:
"""k-fold cross-validation: the dataset is split into k equally-sized folds, one used for validation,
one for testing, the rest for training. This loop iterates k-times over the dataset, cyclically moving the
splits. k==folds, fold in [0,folds) says which split is used for testing.
"""
cf.fold_dir = os.path.join(cf.exp_dir, 'fold_{}'.format(fold))
cf.fold, logger.fold = fold, fold
cf.resume_from_checkpoint = resume_from_checkpoint
if not os.path.exists(cf.fold_dir):
os.mkdir(cf.fold_dir)
train(cf, logger)
cf.resume_from_checkpoint = None
if args.mode == 'train_test':
test(cf, logger)
elif args.mode == 'test':
cf = utils.prep_exp(args.dataset_name, args.exp_dir, args.server_env, use_stored_settings=True, is_training=False)
if args.data_dest is not None:
cf.data_dest = args.data_dest
logger = utils.get_logger(cf.exp_dir, cf.server_env)
data_loader = utils.import_module('data_loader', os.path.join(args.dataset_name, 'data_loader.py'))
model = utils.import_module('model', cf.model_path)
logger.info("loaded model from {}".format(cf.model_path))
fold_dirs = sorted([os.path.join(cf.exp_dir, f) for f in os.listdir(cf.exp_dir) if
os.path.isdir(os.path.join(cf.exp_dir, f)) and f.startswith("fold")])
if folds is None:
folds = range(cf.n_cv_splits)
if args.dev:
folds = folds[:2]
cf.batch_size, cf.max_test_patients, cf.test_n_epochs = 1 if cf.dim==2 else 1, 2, 2
else:
torch.backends.cudnn.benchmark = cf.cuda_benchmark
for fold in folds:
cf.fold = fold
cf.fold_dir = os.path.join(cf.exp_dir, 'fold_{}'.format(cf.fold))
if cf.fold_dir in fold_dirs:
test(cf, logger, max_fold=max([int(f[-1]) for f in fold_dirs]))
else:
logger.info("Skipping fold {} since no model parameters found.".format(fold))
# load raw predictions saved by predictor during testing, run aggregation algorithms and evaluation.
elif args.mode == 'analysis':
""" analyse already saved predictions.
"""
cf = utils.prep_exp(args.dataset_name, args.exp_dir, args.server_env, use_stored_settings=True, is_training=False)
- logger = utils.get_logger(cf.exp_dir, cf.server_env)
+ logger = utils.get_logger(cf.exp_dir, cf.server_env, -1)
if cf.held_out_test_set and not cf.eval_test_fold_wise:
predictor = Predictor(cf, net=None, logger=logger, mode='analysis')
results_list = predictor.load_saved_predictions()
logger.info('starting evaluation...')
cf.fold = 0
evaluator = Evaluator(cf, logger, mode='test')
evaluator.evaluate_predictions(results_list)
evaluator.score_test_df(max_fold=0)
else:
fold_dirs = sorted([os.path.join(cf.exp_dir, f) for f in os.listdir(cf.exp_dir) if
os.path.isdir(os.path.join(cf.exp_dir, f)) and f.startswith("fold")])
if args.dev:
fold_dirs = fold_dirs[:1]
if folds is None:
folds = range(cf.n_cv_splits)
for fold in folds:
cf.fold = fold
cf.fold_dir = os.path.join(cf.exp_dir, 'fold_{}'.format(cf.fold))
if cf.fold_dir in fold_dirs:
predictor = Predictor(cf, net=None, logger=logger, mode='analysis')
results_list = predictor.load_saved_predictions()
# results_list[x][1] is pid, results_list[x][0] is list of len samples-per-patient, each entry hlds
# list of boxes per that sample, i.e., len(results_list[x][y][0]) would be nr of boxes in sample y of patient x
logger.info('starting evaluation...')
evaluator = Evaluator(cf, logger, mode='test')
evaluator.evaluate_predictions(results_list)
max_fold = max([int(f[-1]) for f in fold_dirs])
evaluator.score_test_df(max_fold=max_fold)
else:
logger.info("Skipping fold {} since no model parameters found.".format(fold))
else:
raise ValueError('mode "{}" specified in args is not implemented.'.format(args.mode))
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
torch.cuda.empty_cache()
diff --git a/models/detection_unet.py b/models/detection_unet.py
index 142f560..20394ba 100644
--- a/models/detection_unet.py
+++ b/models/detection_unet.py
@@ -1,545 +1,545 @@
import warnings
import os
import shutil
import time
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import utils.exp_utils as utils
import utils.model_utils as mutils
'''
Use nn.DataParallel to use more than one GPU
'''
def center_crop_2D_image_batched(img, crop_size):
# from batch generator tools from https://github.com/MIC-DKFZ/batchgenerators
# dim 0 is batch, dim 1 is channel, dim 2 and 3 are x y
center = np.array(img.shape[2:]) / 2.
if not hasattr(crop_size, "__iter__"):
center_crop = [int(crop_size)] * (len(img.shape) - 2)
else:
center_crop = np.array(crop_size)
assert len(center_crop) == (len(
img.shape) - 2), "If you provide a list/tuple as center crop make sure it has the same len as your data has dims (2d)"
return img[:, :, int(center[0] - center_crop[0] / 2.):int(center[0] + center_crop[0] / 2.),
int(center[1] - center_crop[1] / 2.):int(center[1] + center_crop[1] / 2.)]
def center_crop_3D_image_batched(img, crop_size):
# dim 0 is batch, dim 1 is channel, dim 2, 3 and 4 are x y z
center = np.array(img.shape[2:]) / 2.
if not hasattr(crop_size, "__iter__"):
center_crop = np.array([int(crop_size)] * (len(img.shape) - 2))
else:
center_crop = np.array(crop_size)
assert len(center_crop) == (len(
img.shape) - 2), "If you provide a list/tuple as center crop make sure it has the same len as your data has dims (3d)"
return img[:, :, int(center[0] - center_crop[0] / 2.):int(center[0] + center_crop[0] / 2.),
int(center[1] - center_crop[1] / 2.):int(center[1] + center_crop[1] / 2.),
int(center[2] - center_crop[2] / 2.):int(center[2] + center_crop[2] / 2.)]
def centercrop_vol(tensor, size):
""":param tensor: tensor whose last two dimensions should be centercropped to size
:param size: 2- or 3-int tuple of target (height, width(,depth))
"""
dim = len(size)
if dim==2:
center_crop_2D_image_batched(tensor, size)
elif dim==3:
center_crop_2D_image_batched(tensor, size)
else:
raise Exception("invalid size argument {} encountered in centercrop".format(size))
"""this below worked so fine, when optional z-dim was first spatial dim instead of last
h_, w_ = size[0], size[1] #target size
(h,w) = tensor.size()[-2:] #orig size
dh, dw = h-h_, w-w_ #deltas
if dim == 3:
d_ = size[2]
d = tensor.size()[-3]
dd = d-d_
if h_<h:
tensor = tensor[...,dh//2:-int(math.ceil(dh/2.)),:] #crop height
elif h_>=h:
print("no h crop")
warn.warn("no height crop applied since target dims larger equal orig dims")
if w_<w:
tensor = tensor[...,dw//2:-int(math.ceil(dw/2.))]
elif w_>=w:
warn.warn("no width crop applied since target dims larger equal orig dims")
if dim == 3:
if d_ < d:
tensor = tensor[..., dd // 2:-int(math.ceil(dd / 2.)),:,:]
elif d_ >= d:
warn.warn("no depth crop applied since target dims larger equal orig dims")
"""
return tensor
def dimcalc_conv2D(dims,F=3,s=1,pad="same"):
r"""
:param dims: orig width, height as (2,)-np.array
:param F: quadratic kernel size
:param s: stride
:param pad: pad
"""
if pad=="same":
pad = (F-1)//2
h, w = dims[0], dims[1]
return np.floor([(h + 2*pad-F)/s+1, (w+ 2*pad-F)/s+1])
def dimcalc_transconv2D(dims,F=2,s=2):
r"""
:param dims: orig width, height as (2,)-np.array
:param F: quadratic kernel size
:param s: stride
"""
h, w = dims[0], dims[1]
return np.array([(h-1)*s+F, (w-1)*s+F])
def dimcalc_Unet_std(init_dims, F=3, F_pool=2, F_up=2, s=1, s_pool=2, s_up=2, pad=0):
r"""Calculate theoretic dimensions of feature maps throughout layers of this U-net.
"""
dims = np.array(init_dims)
print("init dims: ", dims)
def down(dims):
for i in range(2):
dims = dimcalc_conv2D(dims, F=F, s=s, pad=pad)
dims = dimcalc_conv2D(dims, F=F_pool, s=s_pool)
return dims.astype(int)
def up(dims):
for i in range(2):
dims = dimcalc_conv2D(dims, F=F, s=s, pad=pad)
dims = dimcalc_transconv2D(dims, F=F_up,s=s_up)
return dims.astype(int)
stage = 1
for i in range(4):
dims = down(dims)
print("stage ", stage, ": ", dims)
stage+=1
for i in range(4):
dims = up(dims)
print("stage ", stage, ": ", dims)
stage+=1
for i in range(2):
dims = dimcalc_conv2D(dims,F=F,s=s, pad=pad).astype(int)
print("final output size: ", dims)
return dims
def dimcalc_Unet(init_dims, F=3, F_pool=2, F_up=2, s=1, s_pool=2, s_up=2, pad=0):
r"""Calculate theoretic dimensions of feature maps throughout layers of this U-net.
"""
dims = np.array(init_dims)
print("init dims: ", dims)
def down(dims):
for i in range(3):
dims = dimcalc_conv2D(dims, F=F, s=s, pad=pad)
dims = dimcalc_conv2D(dims, F=F_pool, s=s_pool)
return dims.astype(int)
def up(dims):
dims = dimcalc_transconv2D(dims, F=F_up,s=s_up)
for i in range(3):
dims = dimcalc_conv2D(dims, F=F, s=s, pad=pad)
return dims.astype(int)
stage = 1
for i in range(6):
dims = down(dims)
print("stage ", stage, ": ", dims)
stage+=1
for i in range(3):
dims = dimcalc_conv2D(dims, F=F, s=s, pad=pad)
for i in range(6):
dims = up(dims)
print("stage ", stage, ": ", dims)
stage+=1
dims = dims.astype(int)
print("final output size: ", dims)
return dims
class horiz_conv(nn.Module):
def __init__(self, in_chans, out_chans, kernel_size, c_gen, norm, pad=0, relu="relu", bottleneck=True):
super(horiz_conv, self).__init__()
#TODO maybe make res-block?
if bottleneck:
bottleneck = int(np.round((in_chans+out_chans)*3/8))
#print("bottleneck:", bottleneck)
else:
bottleneck = out_chans
self.conv = nn.Sequential(
c_gen(in_chans, bottleneck, kernel_size, pad=pad, norm=norm, relu=relu), #TODO maybe use norm only on last conv?
c_gen(bottleneck, out_chans, kernel_size, pad=pad, norm=norm, relu=relu), #TODO maybe make bottleneck?
#c_gen(out_chans, out_chans, kernel_size, pad=pad, norm=norm, relu=relu),
)
def forward(self, x):
x = self.conv(x)
return x
class up(nn.Module):
def __init__(self, in_chans, out_chans, kernel_size, interpol, c_gen, norm, pad=0, relu="relu", stride_ip=2):
super(up, self).__init__()
self.dim = c_gen.dim
self.upsample = interpol(stride_ip, "bilinear") if self.dim==2 else interpol(stride_ip, "trilinear") #TODO check if fits with spatial dims order in data
self.reduce_chans = c_gen(in_chans, out_chans, ks=1, norm=norm, relu=None)
self.horiz = horiz_conv(out_chans*2, out_chans, kernel_size, c_gen, norm=norm, pad=pad, relu=relu)
def forward(self, x, skip_inp):
#TODO maybe add highway weights in skips?
x = self.upsample(x)
x = self.reduce_chans(x)
#print("shape x, skip", x.shape, skip_inp.shape)
targ_size = x.size()[-self.dim:] #ft map x,y,z (spatial)
skip_inp = centercrop_vol(skip_inp, targ_size)
assert targ_size == skip_inp.size()[-self.dim:], "corresp. skip and forward dimensions don't match"
x = torch.cat((x,skip_inp),dim=1)
x = self.horiz(x)
return x
class net(nn.Module):
r"""U-Net with few more steps than standard.
Dimensions:
feature maps have dims ...xhxwxd, d=feature map depth, h, w = orig
img height, width. h,w each are downsized by unpadded forward-convs and pooling,
upsized by upsampling or upconvolution.
If :math:`F\times F` is the single kernel_size and stride is :math:`s\geq 1`,
:math:`k` is the number of kernels in the conv, i.e. the resulting feature map depth,
(all may differ between operations), then
:Forward Conv: input :math:`h \times w \times d` is converted to
.. math:: \left[ (h-F)//s+1 \right] \times \left[ (w-F)//s+1 \right] \times k
:Pooling: input :math:`h \times w \times d` is converted to
.. math:: \left[ (h-F)//s+1 \right] \times \left[ (w-F)//s+1 \right] \times d,
pooling filters have no depths => orig depths preserved.
:Up-Conv.: input :math:`h \times w \times d` is converted to
.. math:: \left[ (h-1)s + F \right] \times \left[ (w-1)s + F \right] \times k
"""
def down(self, in_chans, out_chans, kernel_size, kernel_size_m, pad=0, relu="relu",maintain_z=False):
"""generate encoder block
:param in_chans:
:param out_chans:
:param kernel_size:
:param pad:
:return:
"""
if maintain_z and self.dim==3:
stride_pool = (2,2,1)
if not hasattr(kernel_size_m, "__iter__"):
kernel_size_m = [kernel_size_m]*self.dim
kernel_size_m = (*kernel_size_m[:-1], 1)
else:
stride_pool = 2
module = nn.Sequential(
nn.MaxPool2d(kernel_size_m, stride=stride_pool) if self.dim == 2 else nn.MaxPool3d(
kernel_size_m, stride=stride_pool),
#--> needs stride 2 in z in upsampling as well!
horiz_conv(in_chans, out_chans, kernel_size, self.c_gen, self.norm, pad, relu=relu)
)
return module
def up(self, in_chans, out_chans, kernel_size, pad=0, relu="relu", maintain_z=False):
"""generate decoder block
:param in_chans:
:param out_chans:
:param kernel_size:
:param pad:
:param relu:
:return:
"""
if maintain_z and self.dim==3:
stride_ip = (2,2,1)
else:
stride_ip = 2
module = up(in_chans, out_chans, kernel_size, self.Interpolator, self.c_gen, norm=self.norm, pad=pad,
relu=relu, stride_ip=stride_ip)
return module
def __init__(self, cf, logger):
super(net, self).__init__()
self.cf = cf
self.dim = cf.dim
self.norm = cf.norm
self.logger = logger
backbone = utils.import_module('bbone', cf.backbone_path)
self.c_gen = backbone.ConvGenerator(cf.dim)
self.Interpolator = backbone.Interpolate
#down = DownBlockGen(cf.dim)
#up = UpBlockGen(cf.dim, backbone.Interpolate)
down = self.down
up = self.up
pad = cf.pad
if pad=="same":
pad = (cf.kernel_size-1)//2
self.dims = "not yet recorded"
self.is_cuda = False
self.init = horiz_conv(len(cf.channels), cf.init_filts, cf.kernel_size, self.c_gen, self.norm, pad=pad,
relu=cf.relu)
self.down1 = down(cf.init_filts, cf.init_filts*2, cf.kernel_size, cf.kernel_size_m, pad=pad, relu=cf.relu)
self.down2 = down(cf.init_filts*2, cf.init_filts*4, cf.kernel_size, cf.kernel_size_m, pad=pad, relu=cf.relu)
self.down3 = down(cf.init_filts*4, cf.init_filts*6, cf.kernel_size, cf.kernel_size_m, pad=pad, relu=cf.relu)
self.down4 = down(cf.init_filts*6, cf.init_filts*8, cf.kernel_size, cf.kernel_size_m, pad=pad, relu=cf.relu,
maintain_z=True)
self.down5 = down(cf.init_filts*8, cf.init_filts*12, cf.kernel_size, cf.kernel_size_m, pad=pad, relu=cf.relu,
maintain_z=True)
#self.down6 = down(cf.init_filts*10, cf.init_filts*14, cf.kernel_size, cf.kernel_size_m, pad=pad, relu=cf.relu)
#self.up1 = up(cf.init_filts*14, cf.init_filts*10, cf.kernel_size, pad=pad, relu=cf.relu)
self.up2 = up(cf.init_filts*12, cf.init_filts*8, cf.kernel_size, pad=pad, relu=cf.relu, maintain_z=True)
self.up3 = up(cf.init_filts*8, cf.init_filts*6, cf.kernel_size, pad=pad, relu=cf.relu, maintain_z=True)
self.up4 = up(cf.init_filts*6, cf.init_filts*4, cf.kernel_size, pad=pad, relu=cf.relu)
self.up5 = up(cf.init_filts*4, cf.init_filts*2, cf.kernel_size, pad=pad, relu=cf.relu)
self.up6 = up(cf.init_filts*2, cf.init_filts, cf.kernel_size, pad=pad, relu=cf.relu)
self.seg = self.c_gen(cf.init_filts, cf.num_seg_classes, 1, norm=None, relu=None) #TODO maybe apply norm too?
# initialize parameters
if self.cf.weight_init == "custom":
logger.info("Tried to use custom weight init which is not defined. Using pytorch default.")
elif self.cf.weight_init:
mutils.initialize_weights(self)
else:
logger.info("using default pytorch weight init")
def forward(self, x):
r'''Forward application of network-function.
:param x: input to the network, expected as torch.tensor of dims
.. math:: batch\_size \times channels \times height \times width
requires_grad should be True for training
'''
#self.dims = np.array([x.size()[-self.dim-1:]])
x1 = self.init(x)
#self.dims = np.vstack((self.dims, x1.size()[-self.dim-1:]))
#---downwards---
x2 = self.down1(x1)
#self.dims = np.vstack((self.dims, x2.size()[-self.dim-1:]))
x3 = self.down2(x2)
#self.dims = np.vstack((self.dims, x3.size()[-self.dim-1:]))
x4 = self.down3(x3)
#self.dims = np.vstack((self.dims, x4.size()[-self.dim-1:]))
x5 = self.down4(x4)
#self.dims = np.vstack((self.dims, x5.size()[-self.dim-1:]))
#x6 = self.down5(x5)
#self.dims = np.vstack((self.dims, x6.size()[-self.dim-1:]))
#---bottom---
x = self.down5(x5)
#self.dims = np.vstack((self.dims, x.size()[-self.dim-1:]))
#---upwards---
#x = self.up1(x, x6)
#self.dims = np.vstack((self.dims, x.size()[-self.dim-1:]))
x = self.up2(x, x5)
#self.dims = np.vstack((self.dims, x.size()[-self.dim-1:]))
x = self.up3(x, x4)
#self.dims = np.vstack((self.dims, x.size()[-self.dim-1:]))
x = self.up4(x, x3)
#self.dims = np.vstack((self.dims, x.size()[-self.dim-1:]))
x = self.up5(x, x2)
#self.dims = np.vstack((self.dims, x.size()[-self.dim-1:]))
x = self.up6(x, x1)
#self.dims = np.vstack((self.dims, x.size()[-self.dim-1:]))
# ---final---
x = self.seg(x)
#self.dims = np.vstack((self.dims, x.size()[-self.dim-1:]))
seg_logits = x
out_box_coords, out_scores = [], []
seg_probs = F.softmax(seg_logits.detach(), dim=1).cpu().data.numpy()
#seg_probs = F.softmax(seg_logits, dim=1)
assert seg_logits.shape[1]==self.cf.num_seg_classes
for cl in range(1, seg_logits.shape[1]):
hard_mask = np.copy(seg_probs).argmax(1)
#hard_mask = seg_probs.clone().argmax(1)
hard_mask[hard_mask != cl] = 0
hard_mask[hard_mask == cl] = 1
# perform connected component analysis on argmaxed predictions,
# draw boxes around components and return coordinates.
box_coords, rois = mutils.get_coords(hard_mask, self.cf.n_roi_candidates, self.cf.dim)
# for each object, choose the highest softmax score (in the respective class)
# of all pixels in the component as object score.
scores = [[] for b_inst in range(x.shape[0])] # np.zeros((out_features.shape[0], self.cf.n_roi_candidates))
for b_inst, brois in enumerate(rois):
for nix, nroi in enumerate(brois):
score_det = np.max if self.cf.score_det == "max" else np.median # score determination
scores[b_inst].append(score_det(seg_probs[b_inst, cl][nroi > 0]))
out_box_coords.append(box_coords)
out_scores.append(scores)
return seg_logits, out_box_coords, out_scores
# noinspection PyCallingNonCallable
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.
:param kwargs:
: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]
'torch_loss': 1D torch tensor for backprop.
'class_loss': classification loss for monitoring. here: dummy array, since no classification conducted.
"""
- img = torch.from_numpy(batch["data"]).cuda()
+ img = torch.from_numpy(batch["data"]).float().cuda()
seg = torch.from_numpy(batch["seg"]).long().cuda()
- seg_ohe = torch.from_numpy(mutils.get_one_hot_encoding(batch['seg'], self.cf.num_seg_classes)).cuda()
+ seg_ohe = torch.from_numpy(mutils.get_one_hot_encoding(batch['seg'], self.cf.num_seg_classes)).float().cuda()
results_dict = {}
seg_logits, box_coords, scores = self.forward(img)
# no extra class loss applied in this model. pass dummy tensor for monitoring.
results_dict['class_loss'] = np.nan
results_dict['boxes'] = [[] for _ in range(img.shape[0])]
for cix in range(len(self.cf.class_dict.keys())):
for bix in range(img.shape[0]):
for rix in range(len(scores[cix][bix])):
if scores[cix][bix][rix] > self.cf.detection_min_confidence:
results_dict['boxes'][bix].append({'box_coords': np.copy(box_coords[cix][bix][rix]),
'box_score': scores[cix][bix][rix],
'box_pred_class_id': cix + 1, # add 0 for background.
'box_type': 'det',
})
for bix in range(img.shape[0]): #bix = batch-element index
for tix in range(len(batch['bb_target'][bix])): #target index
gt_box = {'box_coords': batch['bb_target'][bix][tix], 'box_type': 'gt'}
for name in self.cf.roi_items:
gt_box.update({name: batch[name][bix][tix]})
results_dict['boxes'][bix].append(gt_box)
# compute segmentation loss as either weighted cross entropy, dice loss, or the sum of both.
seg_pred = F.softmax(seg_logits, 1)
loss = torch.tensor([0.], dtype=torch.float, requires_grad=False).cuda()
if self.cf.seg_loss_mode == 'dice' or self.cf.seg_loss_mode == 'dice_wce':
loss += 1 - mutils.batch_dice(seg_pred, seg_ohe.float(),
false_positive_weight=float(self.cf.fp_dice_weight))
if self.cf.seg_loss_mode == 'wce' or self.cf.seg_loss_mode == 'dice_wce':
loss += F.cross_entropy(seg_logits, seg[:, 0], weight=torch.FloatTensor(self.cf.wce_weights).cuda(),
- reduction='elementwise_mean')
+ reduction='mean')
results_dict['torch_loss'] = loss
seg_pred = seg_pred.argmax(dim=1).unsqueeze(dim=1).cpu().data.numpy()
results_dict['seg_preds'] = seg_pred
if 'dice' in self.cf.metrics:
results_dict['batch_dices'] = mutils.dice_per_batch_and_class(seg_pred, batch["seg"],
self.cf.num_seg_classes, convert_to_ohe=True)
#print("batch dice scores ", results_dict['batch_dices'] )
# self.logger.info("loss: {0:.2f}".format(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'
:param kwargs:
: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]
"""
img = torch.FloatTensor(batch['data']).cuda()
seg_logits, box_coords, scores = self.forward(img)
results_dict = {}
results_dict['boxes'] = [[] for b_inst in range(img.shape[0])]
for cix in range(len(box_coords)): #class index
for bix in range(img.shape[0]): #batch instance
for rix in range(len(scores[cix][bix])): #range(self.cf.n_roi_candidates):
if scores[cix][bix][rix] > self.cf.detection_min_confidence:
results_dict['boxes'][bix].append({'box_coords': np.copy(box_coords[cix][bix][rix]),
'box_score': scores[cix][bix][rix],
'box_pred_class_id': cix + 1,
'box_type': 'det'})
# carry probs instead of preds to use for multi-model voting in predictor
results_dict['seg_preds'] = F.softmax(seg_logits, dim=1).cpu().data.numpy()
return results_dict
def actual_dims(self, print_=True):
r"""Return dimensions of actually calculated layers at beginning of each block.
"""
if print_:
print("dimensions as recorded in forward pass: ")
for stage in range(len(self.dims)):
print("Stage ", stage, ": ", self.dims[stage])
return self.dims
def cuda(self, device=None):
r"""Moves all model parameters and buffers to the GPU.
This also makes associated parameters and buffers different objects. So
it should be called before constructing optimizer if the module will
live on GPU while being optimized.
Arguments:
device (int, optional): if specified, all parameters will be
copied to that device
Returns:
Module: self
"""
try:
self.loss_f = self.loss_f.cuda()
except:
pass
self.is_cuda = True
return self._apply(lambda t: t.cuda(device))
def cpu(self):
r"""Moves all model parameters and buffers to the CPU.
Returns:
Module: self
"""
self.is_cuda = False
return self._apply(lambda t: t.cpu())
\ No newline at end of file
diff --git a/models/retina_net.py b/models/retina_net.py
index 7c6aa9b..d618e5a 100644
--- a/models/retina_net.py
+++ b/models/retina_net.py
@@ -1,785 +1,779 @@
#!/usr/bin/env python
# Copyright 2019 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"""
import utils.model_utils as mutils
import utils.exp_utils as utils
import sys
sys.path.append('../')
-# from cuda_functions.nms_2D.pth_nms import nms_gpu as nms_2D
-# from cuda_functions.nms_3D.pth_nms import nms_gpu as nms_3D
from custom_extensions.nms import nms
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils
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, norm=cf.norm)
self.conv_2 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu, norm=cf.norm)
self.conv_3 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu, norm=cf.norm)
self.conv_4 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu, norm=cf.norm)
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.shape[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, norm=cf.norm)
self.conv_2 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu, norm=cf.norm)
self.conv_3 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu, norm=cf.norm)
self.conv_4 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu, norm=cf.norm)
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.shape[0], -1, self.dim * 2)
return [bb_logits]
class RoIRegressor(nn.Module):
def __init__(self, cf, conv, rg_feats):
"""
Builds the RoI-item-regression sub-network. Regression items can be, e.g., malignancy scores of tumors.
"""
super(RoIRegressor, self).__init__()
self.dim = conv.dim
n_input_channels = cf.end_filts
n_features = cf.n_rpn_features
self.rg_feats = rg_feats
n_output_channels = cf.n_anchors_per_pos * self.rg_feats
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, norm=cf.norm)
self.conv_2 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu, norm=cf.norm)
self.conv_3 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu, norm=cf.norm)
self.conv_4 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu, norm=cf.norm)
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)
x = self.conv_final(x)
axes = (0, 2, 3, 1) if self.dim == 2 else (0, 2, 3, 4, 1)
x = x.permute(*axes)
x = x.contiguous()
x = x.view(x.shape[0], -1, self.rg_feats)
return [x]
############################################################
# Loss Functions
############################################################
#
def compute_class_loss(anchor_matches, class_pred_logits, shem_poolsize=20):
"""
:param anchor_matches: (n_anchors). [-1, 0, 1] 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 not 0 in pos_indices.size():
pos_indices = pos_indices.squeeze(1)
roi_logits_pos = class_pred_logits[pos_indices]
targets_pos = anchor_matches[pos_indices].detach()
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 not 0 in neg_indices.size():
neg_indices = neg_indices.squeeze(1)
roi_logits_neg = class_pred_logits[neg_indices]
negative_count = np.max((1, pos_indices.cpu().data.numpy().size))
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, who 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 unused 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: tensor (n_anchors). value in [-1, 0, class_ids] for negative, neutral, and positive matched anchors.
i.e., positively matched anchors are marked by class_id >0
:return: loss: torch 1D tensor.
"""
if not 0 in torch.nonzero(anchor_matches>0).shape:
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.shape[0], :].detach()
# Smooth L1 loss
loss = F.smooth_l1_loss(pred_deltas, target_deltas)
else:
loss = torch.FloatTensor([0]).cuda()
return loss
def compute_rg_loss(tasks, target, pred, 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, 1] for negative, neutral, and positive matched anchors.
:return: loss: torch 1D tensor.
"""
if not 0 in target.shape and not 0 in torch.nonzero(anchor_matches>0).shape:
indices = torch.nonzero(anchor_matches>0).squeeze(1)
# Pick rgs that contribute to the loss
pred = pred[indices]
# Trim target
target = target[:pred.shape[0]].detach()
if 'regression_bin' in tasks:
loss = F.cross_entropy(pred, target.long())
else:
loss = F.smooth_l1_loss(pred, target)
else:
loss = torch.FloatTensor([0]).cuda()
return loss
def compute_focal_class_loss(anchor_matches, class_pred_logits, gamma=2.):
""" Focal Loss FL = -(1-q)^g log(q) with q = pred class probability.
:param anchor_matches: (n_anchors). [-1, 0, class] for negative, neutral, and positive matched anchors.
:param class_pred_logits: (n_anchors, n_classes). logits from classifier sub-network.
:param gamma: g in above formula, good results with g=2 in original paper.
:return: loss: torch tensor
:return: focal loss
"""
# Positive and Negative anchors contribute to the loss,
# but neutral anchors (match value = 0) don't.
pos_indices = torch.nonzero(anchor_matches > 0).squeeze(-1) # dim=-1 instead of 1 or 0 to cover empty matches.
neg_indices = torch.nonzero(anchor_matches == -1).squeeze(-1)
target_classes = torch.cat( (anchor_matches[pos_indices].long(), torch.LongTensor([0] * neg_indices.shape[0]).cuda()) )
non_neutral_indices = torch.cat( (pos_indices, neg_indices) )
q = F.softmax(class_pred_logits[non_neutral_indices], dim=1) # q shape: (n_non_neutral_anchors, n_classes)
# one-hot encoded target classes: keep only the pred probs of the correct class. it will receive incentive to be maximized.
# log(q_i) where i = target class --> FL shape (n_anchors,)
# need to transform to indices into flattened tensor to use torch.take
target_locs_flat = q.shape[1] * torch.arange(q.shape[0]).cuda() + target_classes
q = torch.take(q, target_locs_flat)
FL = torch.log(q) # element-wise log
FL *= -(1-q)**gamma
# take mean over all considered anchors
FL = FL.sum() / FL.shape[0]
return FL
def refine_detections(anchors, probs, deltas, regressions, 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 regressions: (n_proposals, n_classes, n_rg_feats)
: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, pred_regr))
"""
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]]
pre_nms_regressions = regressions[keep_arr[:, 0]]
keep = torch.arange(pre_nms_scores.size()[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
- # if cf.dim == 2:
- # class_keep = nms_2D(torch.cat((ix_rois, ix_scores.unsqueeze(1)), dim=1), cf.detection_nms_threshold)
- # else:
- # class_keep = nms_3D(torch.cat((ix_rois, ix_scores.unsqueeze(1)), dim=1), cf.detection_nms_threshold)
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),
pre_nms_regressions[keep]), dim=1)
return result
def gt_anchor_matching(cf, anchors, gt_boxes, gt_class_ids=None, gt_regressions=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)
gt_regressions: [num_gt_rgs, n_rg_feats], if None empty rg_targets are returned
Returns:
anchor_class_matches: [N] (int32) matches between anchors and GT boxes. class_id = positive anchor,
-1 = negative anchor, 0 = neutral. i.e., positively matched anchors are marked by class_id (which is >0).
anchor_delta_targets: [N, (dy, dx, (dz), log(dh), log(dw), (log(dd)))] Anchor bbox deltas.
anchor_rg_targets: [n_anchors, n_rg_feats]
"""
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))
if gt_regressions is not None:
if 'regression_bin' in cf.prediction_tasks:
anchor_rg_targets = np.zeros((cf.rpn_train_anchors_per_image,))
else:
anchor_rg_targets = np.zeros((cf.rpn_train_anchors_per_image, cf.regression_n_features))
else:
anchor_rg_targets = np.array([])
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, anchor_rg_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 = mutils.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_thresh_ixs = np.argwhere(anchor_iou_max >= anchor_matching_iou)
anchor_class_matches[above_thresh_ixs] = gt_class_ids[anchor_iou_argmax[above_thresh_ixs]]
# Subsample to balance positive anchors.
ids = np.where(anchor_class_matches > 0)[0]
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
# Leave all negative proposals negative for now and sample from them later 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]
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
if gt_regressions is not None:
anchor_rg_targets[ix] = gt_regressions[anchor_iou_argmax[i]]
ix += 1
return anchor_class_matches, anchor_delta_targets, anchor_rg_targets
############################################################
# RetinaNet Class
############################################################
class net(nn.Module):
"""Encapsulates the RetinaNet model functionality.
"""
def __init__(self, cf, logger):
"""
cf: A Sub-class of the cf class
model_dir: Directory to save training logs and trained weights
"""
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")
self.debug_acm = []
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 divisible by 2 at least 5 times "
"to avoid fractions when downscaling and upscaling."
"For example, use 256, 320, 384, 448, 512, ... etc. ")
backbone = utils.import_module('bbone', self.cf.backbone_path)
self.logger.info("loaded backbone from {}".format(self.cf.backbone_path))
conv = backbone.ConvGenerator(self.cf.dim)
# 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).cuda()
self.classifier = Classifier(self.cf, conv).cuda()
self.bb_regressor = BBRegressor(self.cf, conv).cuda()
if 'regression' in self.cf.prediction_tasks:
self.roi_regressor = RoIRegressor(self.cf, conv, self.cf.regression_n_features).cuda()
elif 'regression_bin' in self.cf.prediction_tasks:
# classify into bins of regression values
self.roi_regressor = RoIRegressor(self.cf, conv, len(self.cf.bin_labels)).cuda()
else:
self.roi_regressor = lambda x: [torch.tensor([]).cuda()]
if self.cf.model == 'retina_unet':
self.final_conv = conv(self.cf.end_filts, self.cf.num_seg_classes, ks=1, pad=0, norm=self.cf.norm, relu=None)
def forward(self, img):
"""
:param img: input img (b, c, y, x, (z)).
"""
# Feature extraction
fpn_outs = self.fpn(img)
if self.cf.model == 'retina_unet':
seg_logits = self.final_conv(fpn_outs[0])
selected_fmaps = [fpn_outs[i + 1] for i in self.cf.pyramid_levels]
else:
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, roi_reg_layer_outputs = [], [], [] # list of lists
for p in selected_fmaps:
class_layer_outputs.append(self.classifier(p))
bb_reg_layer_outputs.append(self.bb_regressor(p))
roi_reg_layer_outputs.append(self.roi_regressor(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]
if not 0 == roi_reg_layer_outputs[0][0].shape[0]:
rg_outputs = list(zip(*roi_reg_layer_outputs))
rg_outputs = [torch.cat(list(o), dim=1) for o in rg_outputs][0]
else:
if self.cf.dim == 2:
n_feats = np.array([p.shape[-2] * p.shape[-1] * self.cf.n_anchors_per_pos for p in selected_fmaps]).sum()
else:
n_feats = np.array([p.shape[-3]*p.shape[-2]*p.shape[-1]*self.cf.n_anchors_per_pos for p in selected_fmaps]).sum()
rg_outputs = torch.zeros((selected_fmaps[0].shape[0], n_feats, self.cf.regression_n_features),
dtype=torch.float32).fill_(float('NaN')).cuda()
# 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])
flat_rg_outputs = rg_outputs.view(-1, rg_outputs.shape[-1])
detections = refine_detections(self.anchors, flat_class_softmax, flat_bb_outputs, flat_rg_outputs, batch_ixs,
self.cf)
return detections, class_logits, bb_outputs, rg_outputs, seg_logits
def get_results(self, 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,
pred_regression)
: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, 1] only fg. vs. bg for now.
class-specific return of masks will come with implementation of instance segmentation evaluation.
"""
detections = detections.cpu().data.numpy()
batch_ixs = detections[:, self.cf.dim*2]
detections = [detections[batch_ixs == ix] for ix in range(img_shape[0])]
if box_results_list == None: # for test_forward, where no previous list exists.
box_results_list = [[] for _ in range(img_shape[0])]
for ix in range(img_shape[0]):
if not 0 in detections[ix].shape:
boxes = detections[ix][:, :2 * self.cf.dim].astype(np.int32)
class_ids = detections[ix][:, 2 * self.cf.dim + 1].astype(np.int32)
scores = detections[ix][:, 2 * self.cf.dim + 2]
regressions = detections[ix][:, 2 * self.cf.dim + 3:]
# Filter out detections with zero area. Often only happens in early
# stages of training when the network weights are still a bit random.
if self.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)
regressions = np.delete(regressions, exclude_ix, axis=0)
if not 0 in boxes.shape:
for ix2, score in enumerate(scores):
if score >= self.cf.model_min_confidence:
box = {'box_type': 'det', 'box_coords': boxes[ix2], 'box_score': score,
'box_pred_class_id': class_ids[ix2]}
if "regression_bin" in self.cf.prediction_tasks:
# in this case, regression preds are actually the rg_bin_ids --> map to rg value the bin stands for
box['rg_bin'] = regressions[ix2].argmax()
box['regression'] = self.cf.bin_id2rg_val[box['rg_bin']]
else:
box['regression'] = regressions[ix2]
if hasattr(self.cf, "rg_val_to_bin_id") and \
any(['regression' in task for task in self.cf.prediction_tasks]):
box['rg_bin'] = self.cf.rg_val_to_bin_id(regressions[ix2])
box_results_list[ix].append(box)
results_dict = {}
results_dict['boxes'] = box_results_list
if seg_logits is None:
# output dummy segmentation for retina_net.
out_logits_shape = list(img_shape)
out_logits_shape[1] = self.cf.num_seg_classes
results_dict['seg_preds'] = np.zeros(out_logits_shape, dtype=np.float16)
#todo: try with seg_preds=None? as to not carry heavy dummy preds.
else:
# output label maps for retina_unet.
results_dict['seg_preds'] = F.softmax(seg_logits, 1).cpu().data.numpy()
return results_dict
def train_forward(self, batch, is_validation=False):
"""
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].
'torch_loss': 1D torch tensor for backprop.
'class_loss': classification loss for monitoring.
"""
img = batch['data']
gt_class_ids = batch['class_targets']
gt_boxes = batch['bb_target']
if 'regression' in self.cf.prediction_tasks:
gt_regressions = batch["regression_targets"]
elif 'regression_bin' in self.cf.prediction_tasks:
gt_regressions = batch["rg_bin_targets"]
else:
gt_regressions = None
var_seg_ohe = torch.FloatTensor(mutils.get_one_hot_encoding(batch['seg'], self.cf.num_seg_classes)).cuda()
var_seg = torch.LongTensor(batch['seg']).cuda()
img = torch.from_numpy(img).float().cuda()
torch_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, pred_rgs, 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 tix in range(len(gt_boxes[b])):
gt_box = {'box_type': 'gt', 'box_coords': batch['bb_target'][b][tix]}
for name in self.cf.roi_items:
gt_box.update({name: batch[name][b][tix]})
box_results_list[b].append(gt_box)
# match gt boxes with anchors to generate targets.
anchor_class_match, anchor_target_deltas, anchor_target_rgs = gt_anchor_matching(
self.cf, self.np_anchors, gt_boxes[b], gt_class_ids[b], gt_regressions[b] if gt_regressions is not None else None)
# 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([])
anchor_target_rgs = np.array([])
anchor_class_match = torch.from_numpy(anchor_class_match).cuda()
anchor_target_deltas = torch.from_numpy(anchor_target_deltas).float().cuda()
anchor_target_rgs = torch.from_numpy(anchor_target_rgs).float().cuda()
if self.cf.focal_loss:
# compute class loss as focal loss as suggested in original publication, but multi-class.
class_loss = compute_focal_class_loss(anchor_class_match, class_logits[b], gamma=self.cf.focal_loss_gamma)
# sparing appendix of negative anchors for monitoring as not really relevant
else:
# compute class loss with SHEM.
class_loss, neg_anchor_ix = compute_class_loss(anchor_class_match, class_logits[b])
# 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'})
rg_loss = compute_rg_loss(self.cf.prediction_tasks, anchor_target_rgs, pred_rgs[b], anchor_class_match)
bbox_loss = compute_bbox_loss(anchor_target_deltas, pred_deltas[b], anchor_class_match)
torch_loss += (class_loss + bbox_loss + rg_loss) / img.shape[0]
results_dict = self.get_results(img.shape, detections, seg_logits, box_results_list)
results_dict['seg_preds'] = results_dict['seg_preds'].argmax(axis=1).astype('uint8')[:, np.newaxis]
if self.cf.model == 'retina_unet':
seg_loss_dice = 1 - mutils.batch_dice(F.softmax(seg_logits, dim=1),var_seg_ohe)
seg_loss_ce = F.cross_entropy(seg_logits, var_seg[:, 0])
torch_loss += (seg_loss_dice + seg_loss_ce) / 2
#self.logger.info("loss: {0:.2f}, class: {1:.2f}, bbox: {2:.2f}, seg dice: {3:.3f}, seg ce: {4:.3f}, "
# "mean pixel preds: {5:.5f}".format(torch_loss.item(), batch_class_loss.item(), batch_bbox_loss.item(),
# seg_loss_dice.item(), seg_loss_ce.item(), np.mean(results_dict['seg_preds'])))
if 'dice' in self.cf.metrics:
results_dict['batch_dices'] = mutils.dice_per_batch_and_class(
results_dict['seg_preds'], batch["seg"], self.cf.num_seg_classes, convert_to_ohe=True)
#else:
#self.logger.info("loss: {0:.2f}, class: {1:.2f}, bbox: {2:.2f}".format(
# torch_loss.item(), class_loss.item(), bbox_loss.item()))
results_dict['torch_loss'] = torch_loss
results_dict['class_loss'] = class_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': actually contain seg probabilities since evaluated to seg_preds (via argmax) in predictor.
or dummy seg logits for real retina net (detection only)
"""
img = torch.from_numpy(batch['data']).float().cuda()
detections, _, _, _, seg_logits = self.forward(img)
results_dict = self.get_results(img.shape, detections, seg_logits)
return results_dict
\ No newline at end of file
diff --git a/predictor.py b/predictor.py
index 6b92782..c1f70e9 100644
--- a/predictor.py
+++ b/predictor.py
@@ -1,1005 +1,1007 @@
#!/usr/bin/env python
# Copyright 2019 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
from multiprocessing import Pool
import pickle
import time
import copy
import numpy as np
import torch
from scipy.stats import norm
from collections import OrderedDict
import pandas as pd
import plotting as plg
import utils.model_utils as mutils
import utils.exp_utils as utils
def get_mirrored_patch_crops(patch_crops, org_img_shape):
mirrored_patch_crops = []
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])
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])
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
def get_mirrored_patch_crops_ax_dep(patch_crops, org_img_shape, mirror_axes):
mirrored_patch_crops = []
for ax_ix, axes in enumerate(mirror_axes):
if isinstance(axes, (int, float)) and int(axes) == 0:
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])
elif isinstance(axes, (int, float)) and int(axes) == 1:
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])
elif hasattr(axes, "__iter__") and (tuple(axes) == (0, 1) or tuple(axes) == (1, 0)):
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])
else:
raise Exception("invalid mirror axes {} in get mirrored patch crops".format(axes))
return mirrored_patch_crops
def apply_wbc_to_patient(inputs):
"""
wrapper around prediction box consolidation: weighted box clustering (wbc). 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.
"""
regress_flag, in_patient_results_list, pid, class_dict, clustering_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])
try:
box_patch_id = np.array([b[1]['patch_id'] for b in boxes])
except KeyError: #backward compatibility for already saved pred results ... omg
box_patch_id = np.array([b[1]['ens_ix'] for b in boxes])
box_regressions = np.array([b[1]['regression'] for b in boxes]) if regress_flag else None
box_rg_bins = np.array([b[1]['rg_bin'] if 'rg_bin' in b[1].keys() else float('NaN') for b in boxes])
box_rg_uncs = np.array([b[1]['rg_uncertainty'] if 'rg_uncertainty' in b[1].keys() else float('NaN') for b in boxes])
if 0 not in box_scores.shape:
keep_scores, keep_coords, keep_n_missing, keep_regressions, keep_rg_bins, keep_rg_uncs = \
weighted_box_clustering(box_coords, box_scores, box_center_factor, box_n_overlaps, box_rg_bins, box_rg_uncs,
box_regressions, box_patch_id, clustering_iou, n_ens)
for boxix in range(len(keep_scores)):
clustered_box = {'box_type': 'det', 'box_coords': keep_coords[boxix],
'box_score': keep_scores[boxix], 'cluster_n_missing': keep_n_missing[boxix],
'box_pred_class_id': cl}
if regress_flag:
clustered_box.update({'regression': keep_regressions[boxix],
'rg_uncertainty': keep_rg_uncs[boxix],
'rg_bin': keep_rg_bins[boxix]})
out_patient_results_list[bix].append(clustered_box)
# 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 weighted_box_clustering(box_coords, scores, box_pc_facts, box_n_ovs, box_rg_bins, box_rg_uncs,
box_regress, 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 within 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 box_coords: y1, x1, y2, x2, (z1), (z2).
:param scores: confidence scores.
:param box_pc_facts: patch-center factors from position on patch tiles.
:param box_n_ovs: number of patch overlaps at box position.
:param box_rg_bins: regression bin predictions.
:param box_rg_uncs: (n_dets,) regression uncertainties (from model mrcnn_aleatoric).
:param box_regress: (n_dets, n_regression_features).
:param box_patch_id: ensemble index.
:param thresh: threshold for iou_matching.
:param n_ens: number of models, that are ensembled. (-> number of expected predictions 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 box_coords.shape[1] == 4 else 3
y1 = box_coords[:,0]
x1 = box_coords[:,1]
y2 = box_coords[:,2]
x2 = box_coords[:,3]
areas = (y2 - y1 + 1) * (x2 - x1 + 1)
if dim == 3:
z1 = box_coords[:, 4]
z2 = box_coords[:, 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_scores = []
keep_coords = []
keep_n_missing = []
keep_regress = []
keep_rg_bins = []
keep_rg_uncs = []
while order.size > 0:
i = order[0] # highest scoring element
yy1 = np.maximum(y1[i], y1[order])
xx1 = np.maximum(x1[i], x1[order])
yy2 = np.minimum(y2[i], y2[order])
xx2 = np.minimum(x2[i], x2[order])
w = np.maximum(0, xx2 - xx1 + 1)
h = np.maximum(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, zz2 - zz1 + 1)
inter *= d
# overlap between currently highest scoring box and all boxes.
ovr = inter / (areas[i] + areas[order] - inter)
ovr_fl = inter.astype('float64') / (areas[i] + areas[order] - inter.astype('float64'))
assert np.all(ovr==ovr_fl), "ovr {}\n ovr_float {}".format(ovr, ovr_fl)
# get all the predictions that match the current box to build one cluster.
matches = np.nonzero(ovr > thresh)[0]
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 scores 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 weighted 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))
if box_regress is not None:
# compute wt. avg. of regression vectors (component-wise average)
avg_regress = np.sum(box_regress[order[matches]] * match_scores[:, np.newaxis], axis=0) / np.sum(
match_scores)
avg_rg_bins = np.round(np.sum(box_rg_bins[order[matches]] * match_scores) / np.sum(match_scores))
avg_rg_uncs = np.sum(box_rg_uncs[order[matches]] * match_scores) / np.sum(match_scores)
else:
avg_regress = np.array(float('NaN'))
avg_rg_bins = np.array(float('NaN'))
avg_rg_uncs = np.array(float('NaN'))
# 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)
keep_n_missing.append((n_missing_preds / n_expected_preds * 100)) # relative
keep_regress.append(avg_regress)
keep_rg_uncs.append(avg_rg_uncs)
keep_rg_bins.append(avg_rg_bins)
# get index of all elements that were not matched and discard all others.
inds = np.nonzero(ovr <= thresh)[0]
inds_where = np.where(ovr<=thresh)[0]
assert np.all(inds == inds_where), "inds_nonzero {} \ninds_where {}".format(inds, inds_where)
order = order[inds]
return keep_scores, keep_coords, keep_n_missing, keep_regress, keep_rg_bins, keep_rg_uncs
def apply_nms_to_patient(inputs):
in_patient_results_list, pid, class_dict, iou_thresh = inputs
out_patient_results_list = []
# collect box predictions over batch dimension (slices) and store slice info as slice_ids.
for batch in in_patient_results_list:
batch_el_boxes = []
for cl in list(class_dict.keys()):
det_boxes = [box for box in batch if (box['box_type'] == 'det' and box['box_pred_class_id'] == cl)]
box_coords = np.array([box['box_coords'] for box in det_boxes])
box_scores = np.array([box['box_score'] for box in det_boxes])
if 0 not in box_scores.shape:
keep_ix = mutils.nms_numpy(box_coords, box_scores, iou_thresh)
else:
keep_ix = []
batch_el_boxes += [det_boxes[ix] for ix in keep_ix]
batch_el_boxes += [box for box in batch if box['box_type'] == 'gt']
out_patient_results_list.append(batch_el_boxes)
assert len(in_patient_results_list) == len(out_patient_results_list), "batch dim needs to be maintained, in: {}, out {}".format(len(in_patient_results_list), len(out_patient_results_list))
return [out_patient_results_list, pid]
def nms_2to3D(dets, thresh):
"""
Merges 2D boxes to 3D cubes. For this purpose, boxes of all slices are regarded as lying in one slice.
An adaptation of Non-maximum suppression is applied where clusters are found (like in NMS) with the extra constraint
that suppressed boxes have to have 'connected' z coordinates w.r.t the core slice (cluster center, highest
scoring box, the prevailing box). 'connected' z-coordinates are determined
as the z-coordinates with predictions until the first coordinate for which no prediction is 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 suppressed. 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]
assert np.all(y1 <= y2) and np.all(x1 <= x2), """"the definition of the coordinates is crucially important here:
where maximum is taken needs to be the lower coordinate"""
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: 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: 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 + 1)
w = np.maximum(0.0, xx2 - xx1 + 1)
inter = h * w
iou = inter / (areas[i] + areas[order] - inter)
matches = np.argwhere(
iou > 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_holes = [ii for ii in np.arange(core_slice, np.max(slice_ids)) if ii not in slice_ids]
lower_holes = [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_holes) if len(upper_holes) > 0 else np.max(slice_ids)
min_valid_slice_id = np.max(lower_holes) if len(lower_holes) > 0 else np.min(slice_ids)
z_matches = matches[(slice_ids <= max_valid_slice_id) & (slice_ids >= min_valid_slice_id)]
# expand by one z voxel since box content is surrounded w/o overlap, i.e., z-content computed as z2-z1
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 apply_2d_3d_merging_to_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()):
det_boxes, slice_ids = [], []
# collect box predictions over batch dimension (slices) and store slice info as slice_ids.
for batch_ix, batch in enumerate(in_patient_results_list):
batch_element_det_boxes = [(ix, box) for ix, box in enumerate(batch) if
(box['box_type'] == 'det' and box['box_pred_class_id'] == cl)]
det_boxes += batch_element_det_boxes
slice_ids += [batch_ix] * len(batch_element_det_boxes)
box_coords = np.array([batch[1]['box_coords'] for batch in det_boxes])
box_scores = np.array([batch[1]['box_score'] for batch in det_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})
for kix, kz in zip(keep_ix, keep_z):
keep_box = det_boxes[kix][1]
keep_box['box_coords'] = list(keep_box['box_coords']) + kz
out_patient_results_list.append(keep_box)
gt_boxes = [box for b in in_patient_results_list for box in b if box['box_type'] == 'gt']
if len(gt_boxes) > 0:
assert np.all([len(box["box_coords"]) == 6 for box in gt_boxes]), "expanded preds to 3D but GT is 2D."
out_patient_results_list += gt_boxes
return [[out_patient_results_list], pid] # additional list wrapping is extra batch dim.
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: apply_2d_3d_merging_to_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 patient-wise ground truth to the final predictions to be passed to the evaluator.
"""
def __init__(self, cf, net, logger, mode):
self.cf = cf
self.batch_size = cf.batch_size
self.logger = logger
self.mode = mode
self.net = net
self.n_ens = 1
self.rank_ix = '0'
self.regress_flag = any(['regression' in task for task in self.cf.prediction_tasks])
if self.cf.merge_2D_to_3D_preds:
assert self.cf.dim == 2, "Merge 2Dto3D only valid for 2D preds, but current dim is {}.".format(self.cf.dim)
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_axes is not None:
self.n_ens *= (len(self.cf.test_aug_axes)+1)
self.example_plot_dir = os.path.join(cf.test_dir, "example_plots")
os.makedirs(self.example_plot_dir, exist_ok=True)
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))
- loss / class_loss (only in validation mode)
"""
#self.logger.info('forwarding (patched) patient with shape: {}'.format(batch['data'].shape))
img = batch['data']
if img.shape[0] <= self.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']]
elif self.mode == 'test':
results_dict = self.net.test_forward(batch, return_masks=self.cf.return_masks_in_test)
else: #needs batch tiling
split_ixs = np.split(np.arange(img.shape[0]), np.arange(img.shape[0])[::self.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':
# if hasattr(self.cf, "losses_to_monitor"):
# loss_names = self.cf.losses_to_monitor
# else:
# loss_names = {name for dic in chunk_dicts for name in dic if 'loss' in name}
# estimate patient loss by mean over batch_chunks. Most similar to training loss.
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])
# 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 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))
- monitor_values (only in validation mode)
returned dict is a flattened version with 1 batch instance (3D) or slices (2D)
"""
if patch_crops is not None:
#print("patch_crops not None, applying patch center factor")
patches_dict = self.batch_tiling_forward(batch)
results_dict = {'boxes': [[] for _ in range(batch['original_img_shape'][0])]}
#bc of ohe--> channel dim of seg has size num_classes
out_seg_shape = list(batch['original_img_shape'])
out_seg_shape[1] = patches_dict["seg_preds"].shape[1]
out_seg_preds = np.zeros(out_seg_shape, dtype=np.float16)
patch_overlap_map = np.zeros_like(out_seg_preds, dtype='uint8')
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]
patch_overlap_map[:, :, pc[0]:pc[1], pc[2]:pc[3], pc[4]:pc[5]] += 1
elif self.cf.dim == 2:
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
out_seg_preds[patch_overlap_map > 0] /= patch_overlap_map[patch_overlap_map > 0]
results_dict['seg_preds'] = out_seg_preds
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 down-weighted for consolidation, using the 'box_patch_center_factor', which is
# obtained by a gaussian distribution over positions in the patch and average over spatial dimensions.
# Also the info 'box_n_overlaps' is stored for consolidation, which represents the amount of
# overlapping patches at the box's position.
c = box['box_coords']
#box_centers = np.array([(c[ii] + c[ii+2])/2 for ii in range(len(c)//2)])
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':
results_dict['torch_loss'] = patches_dict['torch_loss']
results_dict['class_loss'] = patches_dict['class_loss']
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 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))
- loss / class_loss (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_axes is not None:
if isinstance(self.cf.test_aug_axes, (int, float)):
self.cf.test_aug_axes = (self.cf.test_aug_axes,)
#assert np.all(np.array(self.cf.test_aug_axes)<self.cf.dim), "test axes {} need to be spatial axes".format(self.cf.test_aug_axes)
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_ax_dep(patch_crops, batch['original_img_shape'],
self.cf.test_aug_axes)
self.logger.info("mirrored patch crop coords for patched patient in test augs!")
else:
mirrored_patch_crops = [None] * 3
img = np.copy(batch['data'])
for n_aug, sp_axis in enumerate(self.cf.test_aug_axes):
#sp_axis = np.array(axis) #-2 #spatial axis index
axis = np.array(sp_axis)+2
if isinstance(sp_axis, (int, float)):
# mirroring along one axis at a time
batch['data'] = np.flip(img, axis=axis).copy()
chunk_dict = self.spatial_tiling_forward(batch, mirrored_patch_crops[n_aug], n_aug=str(n_aug))
# 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[sp_axis] = org_img_shape[axis] - chunk_dict['boxes'][ix][boxix]['box_coords'][sp_axis+2]
coords[sp_axis+2] = org_img_shape[axis] - chunk_dict['boxes'][ix][boxix]['box_coords'][sp_axis]
assert coords[2] >= coords[0], [coords, chunk_dict['boxes'][ix][boxix]['box_coords']]
assert coords[3] >= coords[1], [coords, chunk_dict['boxes'][ix][boxix]['box_coords']]
chunk_dict['boxes'][ix][boxix]['box_coords'] = coords
# re-transform segmentation predictions.
chunk_dict['seg_preds'] = np.flip(chunk_dict['seg_preds'], axis=axis)
elif hasattr(sp_axis, "__iter__") and tuple(sp_axis)==(0,1) or tuple(sp_axis)==(1,0):
#NEED: mirrored patch crops are given as [(y-axis), (x-axis), (y-,x-axis)], obey this order!
# mirroring along two axes at same time
batch['data'] = np.flip(np.flip(img, axis=axis[0]), axis=axis[1]).copy()
chunk_dict = self.spatial_tiling_forward(batch, mirrored_patch_crops[n_aug], n_aug=str(n_aug))
# 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[sp_axis[0]] = org_img_shape[axis[0]] - chunk_dict['boxes'][ix][boxix]['box_coords'][sp_axis[0]+2]
coords[sp_axis[0]+2] = org_img_shape[axis[0]] - chunk_dict['boxes'][ix][boxix]['box_coords'][sp_axis[0]]
coords[sp_axis[1]] = org_img_shape[axis[1]] - chunk_dict['boxes'][ix][boxix]['box_coords'][sp_axis[1]+2]
coords[sp_axis[1]+2] = org_img_shape[axis[1]] - chunk_dict['boxes'][ix][boxix]['box_coords'][sp_axis[1]]
assert coords[2] >= coords[0], [coords, chunk_dict['boxes'][ix][boxix]['box_coords']]
assert coords[3] >= coords[1], [coords, chunk_dict['boxes'][ix][boxix]['box_coords']]
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=axis[0]), axis=axis[1]).copy()
else:
raise Exception("Invalid axis type {} in test augs".format(type(axis)))
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])])
results_dict['seg_preds'] = np.stack([dic['seg_preds'] for dic in results_list], axis=1)
# needs segs probs in seg_preds entry:
results_dict['seg_preds'] = np.sum(results_dict['seg_preds'], axis=1) #add up seg probs from different augs per class
if self.mode == 'val':
results_dict['torch_loss'] = results_list[0]['torch_loss']
results_dict['class_loss'] = results_list[0]['class_loss']
return results_dict
def load_saved_predictions(self):
"""loads raw predictions saved by self.predict_test_set. aggregates and/or merges 2D boxes to 3D cubes for
evaluation (if model predicts 2D but evaluation is run in 3D), according to settings config.
:return: 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.
- 'batch_dices': dice scores as recorded in raw prediction results.
- 'seg_preds': not implemented yet. could replace dices by seg preds to have raw seg info available, however
would consume critically large memory amount. todo evaluation of instance/semantic segmentation.
"""
results_file = 'pred_results.pkl' if not self.cf.held_out_test_set else 'pred_results_held_out.pkl'
if not self.cf.held_out_test_set or self.cf.eval_test_fold_wise:
self.logger.info("loading saved predictions of fold {}".format(self.cf.fold))
with open(os.path.join(self.cf.fold_dir, results_file), 'rb') as handle:
results_list = pickle.load(handle)
box_results_list = [(res_dict["boxes"], pid) for res_dict, pid in results_list]
da_factor = len(self.cf.test_aug_axes)+1 if self.cf.test_aug_axes is not None else 1
self.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), self.n_ens))
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, results_file), '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))
#it's usually not right to filter out gts here, is it?
da_factor = len(self.cf.test_aug_axes)+1 if self.cf.test_aug_axes is not None else 1
self.n_ens = self.cf.test_n_epochs * da_factor * folds_loaded
# -------------- aggregation of boxes via clustering -----------------
if self.cf.clustering == "wbc":
self.logger.info('applying WBC to test-set predictions with iou {} and n_ens {} over {} patients'.format(
self.cf.clustering_iou, self.n_ens, len(box_results_list)))
mp_inputs = [[self.regress_flag, ii[0], ii[1], self.cf.class_dict, self.cf.clustering_iou, self.n_ens] for ii
in box_results_list]
del box_results_list
pool = Pool(processes=self.cf.n_workers)
box_results_list = pool.map(apply_wbc_to_patient, mp_inputs, chunksize=1)
pool.close()
pool.join()
del mp_inputs
elif self.cf.clustering == "nms":
self.logger.info('applying standard NMS to test-set predictions with iou {} over {} patients.'.format(
self.cf.clustering_iou, len(box_results_list)))
pool = Pool(processes=self.cf.n_workers)
mp_inputs = [[ii[0], ii[1], self.cf.class_dict, self.cf.clustering_iou] for ii in box_results_list]
del box_results_list
box_results_list = pool.map(apply_nms_to_patient, mp_inputs, chunksize=1)
pool.close()
pool.join()
del mp_inputs
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=self.cf.n_workers)
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(apply_2d_3d_merging_to_patient, mp_inputs, chunksize=1)
pool.close()
pool.join()
del mp_inputs
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 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))
- loss / class_loss (only in validation mode)
"""
if self.mode=="test":
self.logger.info('predicting patient {} for fold {} '.format(np.unique(batch['pid']), self.cf.fold))
# True if patient is provided in patches and predictions need to be tiled.
self.patched_patient = 'patch_crop_coords' in list(batch.keys())
# forward batch through prediction pipeline.
results_dict = self.data_aug_forward(batch)
#has seg probs in entry 'seg_preds'
if self.mode == 'val':
for b in range(batch['patient_bb_target'].shape[0]):
for t in range(len(batch['patient_bb_target'][b])):
gt_box = {'box_type': 'gt', 'box_coords': batch['patient_bb_target'][b][t],
'class_targets': batch['patient_class_targets'][b][t]}
for name in self.cf.roi_items:
gt_box.update({name : batch['patient_'+name][b][t]})
results_dict['boxes'][b].append(gt_box)
if 'dice' in self.cf.metrics:
if self.patched_patient:
assert 'patient_seg' in batch.keys(), "Results_dict preds are in original patient shape."
results_dict['batch_dices'] = mutils.dice_per_batch_and_class(
results_dict['seg_preds'], batch["patient_seg"] if self.patched_patient else batch['seg'],
self.cf.num_seg_classes, convert_to_ohe=True)
if self.patched_patient and self.cf.clustering == "wbc":
wbc_input = [self.regress_flag, results_dict['boxes'], 'dummy_pid', self.cf.class_dict, self.cf.clustering_iou, self.n_ens]
results_dict['boxes'] = apply_wbc_to_patient(wbc_input)[0]
elif self.patched_patient:
nms_inputs = [results_dict['boxes'], 'dummy_pid', self.cf.class_dict, self.cf.clustering_iou]
results_dict['boxes'] = apply_nms_to_patient(nms_inputs)[0]
if self.cf.merge_2D_to_3D_preds:
results_dict['2D_boxes'] = results_dict['boxes']
merge_dims_inputs = [results_dict['boxes'], 'dummy_pid', self.cf.class_dict, self.cf.merge_3D_iou]
results_dict['boxes'] = apply_2d_3d_merging_to_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 evaluation of instance/semantic segmentation.
"""
# -------------- raw predicting -----------------
dict_of_patients_results = OrderedDict()
set_of_result_types = set()
# 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_params.pth'.format(epoch)) 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)
with torch.no_grad():
plot_batches = np.random.choice(np.arange(batch_gen['n_test']), size=self.cf.n_test_plots, replace=False)
for i in range(batch_gen['n_test']):
batch = next(batch_gen['test'])
pid = np.unique(batch['pid'])
assert len(pid)==1
pid = pid[0]
if not pid in dict_of_patients_results.keys(): # store batch info in patient entry of results dict.
dict_of_patients_results[pid] = {}
dict_of_patients_results[pid]['results_dicts'] = []
dict_of_patients_results[pid]['patient_bb_target'] = batch['patient_bb_target']
for name in self.cf.roi_items:
dict_of_patients_results[pid]["patient_"+name] = batch["patient_"+name]
stime = time.time()
results_dict = self.predict_patient(batch) #only holds "boxes", "seg_preds"
# needs ohe seg probs in seg_preds entry:
results_dict['seg_preds'] = np.argmax(results_dict['seg_preds'], axis=1)[:,np.newaxis]
self.logger.info("predicting patient {} with weight rank {} (progress: {}/{}) took {:.2f}s".format(
str(pid), rank_ix, (rank_ix)*batch_gen['n_test']+(i+1), len(weight_paths)*batch_gen['n_test'], time.time()-stime))
if i in plot_batches and (not self.patched_patient or 'patient_data' in batch.keys()):
try:
# view qualitative results of random test case
+ self.logger.time("test_plot")
out_file = os.path.join(self.example_plot_dir,
'batch_example_test_{}_rank_{}.png'.format(self.cf.fold, rank_ix))
plg.view_batch(self.cf, batch, res_dict=results_dict, out_file=out_file,
show_seg_ids='dice' in self.cf.metrics,
has_colorchannels=self.cf.has_colorchannels, show_gt_labels=True)
+ self.logger.info("generated example test plot {} in {:.2f}s".format(os.path.basename(out_file), self.logger.time("test_plot")))
except Exception as e:
self.logger.info("WARNING: error in view_batch: {}".format(e))
if 'dice' in self.cf.metrics:
if self.patched_patient:
assert 'patient_seg' in batch.keys(), "Results_dict preds are in original patient shape."
results_dict['batch_dices'] = mutils.dice_per_batch_and_class( results_dict['seg_preds'],
batch["patient_seg"] if self.patched_patient else batch['seg'],
self.cf.num_seg_classes, convert_to_ohe=True)
dict_of_patients_results[pid]['results_dicts'].append({k:v for k,v in results_dict.items()
if k in ["boxes", "batch_dices"]})
# collect result types to know which ones to look for when saving
set_of_result_types.update(dict_of_patients_results[pid]['results_dicts'][-1].keys())
# -------------- re-order, save raw results -----------------
self.logger.info('finished predicting test set. starting aggregation of predictions.')
results_per_patient = []
for pid, p_dict in dict_of_patients_results.items():
# dict_of_patients_results[pid]['results_list'] has length batch['n_test']
results_dict = {}
# collect all boxes/seg_preds of same batch_instance over temporal instances.
b_size = len(p_dict['results_dicts'][0]["boxes"])
for res_type in [rtype for rtype in set_of_result_types if rtype in ["boxes", "batch_dices"]]:#, "seg_preds"]]:
if not 'batch' in res_type: #assume it's results on batch-element basis
results_dict[res_type] = [[item for rank_dict in p_dict['results_dicts'] for item in rank_dict[res_type][batch_instance]]
for batch_instance in range(b_size)]
else:
results_dict[res_type] = []
for dict in p_dict['results_dicts']:
if 'dice' in res_type:
item = dict[res_type] #dict['batch_dices'] has shape (num_seg_classes,)
assert len(item) == self.cf.num_seg_classes, \
"{}, {}".format(len(item), self.cf.num_seg_classes)
else:
raise NotImplementedError
results_dict[res_type].append(item)
# rdict[dice] shape (n_rank_epochs (n_saved_ranks), nsegclasses)
# calc mean over test epochs so inline with shape from sampling
results_dict[res_type] = np.mean(results_dict[res_type], axis=0) #maybe error type with other than dice
if not hasattr(self.cf, "eval_test_separately") or not self.cf.eval_test_separately:
# add unpatched 2D or 3D (if dim==3 or merge_2D_to_3D) ground truth boxes for evaluation.
for b in range(p_dict['patient_bb_target'].shape[0]):
for targ in range(len(p_dict['patient_bb_target'][b])):
gt_box = {'box_type': 'gt', 'box_coords':p_dict['patient_bb_target'][b][targ],
'class_targets': p_dict['patient_class_targets'][b][targ]}
for name in self.cf.roi_items:
gt_box.update({name: p_dict["patient_"+name][b][targ]})
results_dict['boxes'][b].append(gt_box)
results_per_patient.append([results_dict, pid])
out_string = 'pred_results_held_out' if self.cf.held_out_test_set else 'pred_results'
with open(os.path.join(self.cf.fold_dir, '{}.pkl'.format(out_string)), 'wb') as handle:
pickle.dump(results_per_patient, handle)
if return_results:
# -------------- results processing, clustering, etc. -----------------
final_patient_box_results = [ (res_dict["boxes"], pid) for res_dict,pid in results_per_patient ]
if self.cf.clustering == "wbc":
self.logger.info('applying WBC to test-set predictions with iou = {} and n_ens = {}.'.format(
self.cf.clustering_iou, self.n_ens))
mp_inputs = [[self.regress_flag, ii[0], ii[1], self.cf.class_dict, self.cf.clustering_iou, self.n_ens] for ii in final_patient_box_results]
del final_patient_box_results
pool = Pool(processes=self.cf.n_workers)
final_patient_box_results = pool.map(apply_wbc_to_patient, mp_inputs, chunksize=1)
pool.close()
pool.join()
del mp_inputs
elif self.cf.clustering == "nms":
self.logger.info('applying standard NMS to test-set predictions with iou = {}.'.format(self.cf.clustering_iou))
pool = Pool(processes=self.cf.n_workers)
mp_inputs = [[ii[0], ii[1], self.cf.class_dict, self.cf.clustering_iou] for ii in final_patient_box_results]
del final_patient_box_results
final_patient_box_results = pool.map(apply_nms_to_patient, mp_inputs, chunksize=1)
pool.close()
pool.join()
del mp_inputs
if self.cf.merge_2D_to_3D_preds:
self.logger.info('applying 2D-to-3D merging to test-set predictions with iou = {}.'.format(self.cf.merge_3D_iou))
mp_inputs = [[ii[0], ii[1], self.cf.class_dict, self.cf.merge_3D_iou] for ii in final_patient_box_results]
del final_patient_box_results
pool = Pool(processes=self.cf.n_workers)
final_patient_box_results = pool.map(apply_2d_3d_merging_to_patient, mp_inputs, chunksize=1)
pool.close()
pool.join()
del mp_inputs
# 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]
# results_per_patient = [(res_dict["boxes"] = boxes, pid) for (boxes,pid) in final_patient_box_results]
return results_per_patient # holds list of (results_dict, pid)
diff --git a/requirements.txt b/requirements.txt
index 1161169..88d3644 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,113 +1,66 @@
-absl-py==0.7.1
-alabaster==0.7.11
-asn1crypto==0.24.0
-astor==0.7.1
-astroid==2.0.1
-Babel==2.6.0
+absl-py==0.8.1
backcall==0.1.0
-batchgenerators==0.18.2
-bleach==2.1.3
-certifi==2018.4.16
-cffi==1.11.5
+batchgenerators==0.19.3
+cachetools==3.1.1
+certifi==2019.11.28
chardet==3.0.4
-cloudpickle==0.5.3
-cryptography==2.3
cycler==0.10.0
-Cython==0.29.6
-dask==0.18.1
-decorator==4.3.0
-docutils==0.14
-entrypoints==0.2.3
-future==0.16.0
-gast==0.2.0
-grpcio==1.13.0
-h5py==2.9.0
-html5lib==1.0.1
-idna==2.7
-imagesize==1.0.0
-isort==4.3.4
-jedi==0.12.1
-jeepney==0.3.1
-Jinja2==2.10.1
-joblib==0.13.2
-jsonschema==2.6.0
-Keras-Applications==1.0.7
-Keras-Preprocessing==1.0.9
-keyring==13.2.1
-kiwisolver==1.0.1
-lazy-object-proxy==1.3.1
+Cython==0.29.14
+decorator==4.4.1
+future==0.18.2
+google-auth==1.7.1
+google-auth-oauthlib==0.4.1
+grpcio==1.25.0
+idna==2.8
+imageio==2.6.1
+ipython==7.9.0
+ipython-genutils==0.2.0
+jedi==0.15.1
+joblib==0.14.0
+kiwisolver==1.1.0
linecache2==1.0.0
-Markdown==2.6.11
-MarkupSafe==1.0
-matplotlib==3.0.3
-mccabe==0.6.1
-mistune==0.8.3
-mock==2.0.0
-nbconvert==5.3.1
-nbformat==4.4.0
-networkx==2.1
-nibabel==2.3.0
-nilearn==0.4.2
-numpy==1.14.5
-numpydoc==0.8.0
-nvidia-ml-py3==7.352.0
-packaging==17.1
-pandas==0.24.2
-pandocfilters==1.4.2
-parso==0.3.1
-pathlib==1.0.1
-pbr==5.1.3
-pexpect==4.6.0
-pickleshare==0.7.4
+Markdown==3.1.1
+matplotlib==3.1.2
+networkx==2.4
+ninja==1.9.0.post1
+nms-extension==0.0.0
+numpy==1.17.3
+oauthlib==3.1.0
+pandas==0.25.3
+parso==0.5.1
+pexpect==4.7.0
+pickleshare==0.7.5
Pillow==6.2.1
-prompt-toolkit==2.0.9
-protobuf==3.7.1
-psutil==5.4.6
+prompt-toolkit==2.0.10
+protobuf==3.10.0
+psutil==5.6.5
ptyprocess==0.6.0
-pycodestyle==2.4.0
-pycparser==2.18
-pyflakes==2.0.0
-Pygments==2.2.0
-pylint==2.0.1
-PyOpenGL==3.1.0
-pyparsing==2.2.0
-PyQt5==5.9.2
-python-dateutil==2.7.3
-pytz==2018.5
-PyWavelets==0.5.2
-pyzmq==17.1.0
-QtAwesome==0.4.4
-qtconsole==4.3.1
-QtPy==1.4.2
+pyasn1==0.4.8
+pyasn1-modules==0.2.7
+Pygments==2.4.2
+pyparsing==2.4.4
+python-dateutil==2.8.1
+pytz==2019.3
+PyWavelets==1.1.1
+RegRCNN==0.0.2
requests==2.22.0
-rope==0.14.0
-scikit-image==0.14.0
+requests-oauthlib==1.3.0
+RoIAlign-extension-2D==0.0.0
+RoIAlign-extension-3D==0.0.0
+rsa==4.0
+scikit-image==0.16.2
scikit-learn==0.21.3
-scipy==1.1.0
-SecretStorage==3.0.1
-simplegeneric==0.8.1
-SimpleITK==1.2.2
-sip==4.19.8
-six==1.11.0
-snowballstemmer==1.2.1
-Sphinx==1.7.6
-sphinxcontrib-websupport==1.1.0
-tensorboard==1.13.1
-tensorboardX==1.6
-tensorflow==1.13.1
-tensorflow-estimator==1.13.0
-termcolor==1.1.0
-testpath==0.3.1
-toolz==0.9.0
-torch==0.4.1
-torchvision==0.2.1
-tornado==5.1
+scipy==1.3.1
+SimpleITK==1.2.3
+six==1.13.0
+tensorboard==2.0.2
+threadpoolctl==1.1.0
+torch==1.3.1
+torchvision==0.4.2
+tqdm==4.39.0
traceback2==1.4.0
-traitlets==4.3.2
-typed-ast==1.1.0
+traitlets==4.3.3
unittest2==1.1.0
-urllib3==1.25.3
+urllib3==1.25.7
wcwidth==0.1.7
-webencodings==0.5.1
-Werkzeug==0.15.5
-wrapt==1.10.11
+Werkzeug==0.16.0
diff --git a/utils/exp_utils.py b/utils/exp_utils.py
index 68de0d3..138cdb2 100644
--- a/utils/exp_utils.py
+++ b/utils/exp_utils.py
@@ -1,629 +1,632 @@
#!/usr/bin/env python
# Copyright 2019 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 plotting as plg
import sys
import os
import subprocess
import threading
import pickle
import importlib.util
import psutil
import time
import logging
from torch.utils.tensorboard import SummaryWriter
from collections import OrderedDict
import numpy as np
import pandas as pd
import torch
def import_module(name, path):
"""
correct way of importing a module dynamically in python 3.
:param name: name given to module instance.
:param path: path to module.
:return: module: returned module instance.
"""
spec = importlib.util.spec_from_file_location(name, path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def save_obj(obj, name):
"""Pickle a python object."""
with open(name + '.pkl', 'wb') as f:
pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)
def load_obj(file_path):
with open(file_path, 'rb') as handle:
return pickle.load(handle)
def IO_safe(func, *args, _tries=5, _raise=True, **kwargs):
""" Wrapper calling function func with arguments args and keyword arguments kwargs to catch input/output errors
on cluster.
:param func: function to execute (intended to be read/write operation to a problematic cluster drive, but can be
any function).
:param args: positional args of func.
:param kwargs: kw args of func.
:param _tries: how many attempts to make executing func.
"""
for _try in range(_tries):
try:
return func(*args, **kwargs)
except OSError as e: # to catch cluster issues with network drives
if _raise:
raise e
else:
print("After attempting execution {} time{}, following error occurred:\n{}".format(_try+1,"" if _try==0 else "s", e))
continue
-
def query_nvidia_gpu(device_id, d_keyword=None, no_units=False):
"""
:param device_id:
:param d_keyword: -d, --display argument (keyword(s) for selective display), all are selected if None
:return: dict of gpu-info items
"""
cmd = ['nvidia-smi', '-i', str(device_id), '-q']
if d_keyword is not None:
cmd += ['-d', d_keyword]
outp = subprocess.check_output(cmd).strip().decode('utf-8').split("\n")
outp = [x for x in outp if len(x)>0]
headers = [ix for ix, item in enumerate(outp) if len(item.split(":"))==1] + [len(outp)]
out_dict = {}
for lix, hix in enumerate(headers[:-1]):
head = outp[hix].strip().replace(" ", "_").lower()
out_dict[head] = {}
for lix2 in range(hix, headers[lix+1]):
try:
key, val = [x.strip().lower() for x in outp[lix2].split(":")]
if no_units:
val = val.split()[0]
out_dict[head][key] = val
except:
pass
return out_dict
class CombinedPrinter(object):
"""combined print function.
prints to logger and/or file if given, to normal print if non given.
"""
def __init__(self, logger=None, file=None):
if logger is None and file is None:
self.out = [print]
elif logger is None:
self.out = [file.write]
elif file is None:
self.out = [logger.info]
else:
self.out = [logger.info, file.write]
def __call__(self, string):
for fct in self.out:
fct(string)
class Nvidia_GPU_Logger(object):
def __init__(self):
self.count = None
def get_vals(self):
cmd = ['nvidia-settings', '-t', '-q', 'GPUUtilization']
gpu_util = subprocess.check_output(cmd).strip().decode('utf-8').split(",")
gpu_util = dict([f.strip().split("=") for f in gpu_util])
cmd[-1] = 'UsedDedicatedGPUMemory'
gpu_used_mem = subprocess.check_output(cmd).strip().decode('utf-8')
current_vals = {"gpu_mem_alloc": gpu_used_mem, "gpu_graphics_util": int(gpu_util['graphics']),
"gpu_mem_util": gpu_util['memory'], "time": time.time()}
return current_vals
def loop(self):
i = 0
while True:
self.get_vals()
self.log["time"].append(time.time())
self.log["gpu_util"].append(self.current_vals["gpu_graphics_util"])
- if self.count != None:
+ if self.count is not None:
i += 1
- if i == count:
+ if i == self.count:
exit(0)
time.sleep(self.interval)
def start(self, interval=1.):
self.interval = interval
self.start_time = time.time()
self.log = {"time": [], "gpu_util": []}
if self.interval is not None:
thread = threading.Thread(target=self.loop)
thread.daemon = True
thread.start()
class CombinedLogger(object):
"""Combine console and tensorboard logger and record system metrics.
"""
def __init__(self, name, log_dir, server_env=True, fold="", sysmetrics_interval=2):
self.pylogger = logging.getLogger(name)
self.tboard = SummaryWriter(log_dir=log_dir)
self.times = {}
self.fold = fold
- # monitor system metrics (cpu, mem, ...)
+
+ self.pylogger.setLevel(logging.DEBUG)
+ self.log_file = os.path.join(log_dir, 'exec.log')
+ self.pylogger.addHandler(logging.FileHandler(self.log_file))
if not server_env:
+ self.pylogger.addHandler(ColorHandler())
+ else:
+ self.pylogger.addHandler(logging.StreamHandler())
+ self.pylogger.propagate = False
+
+ # monitor system metrics (cpu, mem, ...)
+ if not server_env and sysmetrics_interval>0:
self.sysmetrics = pd.DataFrame(columns=["global_step", "rel_time", r"CPU (%)", "mem_used (GB)", r"mem_used (%)",
r"swap_used (GB)", r"gpu_utilization (%)"], dtype="float16")
for device in range(torch.cuda.device_count()):
self.sysmetrics["mem_allocd (GB) by torch on {:10s}".format(torch.cuda.get_device_name(device))] = np.nan
self.sysmetrics["mem_cached (GB) by torch on {:10s}".format(torch.cuda.get_device_name(device))] = np.nan
self.sysmetrics_start(sysmetrics_interval)
+ else:
+ print("NOT logging sysmetrics")
def __getattr__(self, attr):
"""delegate all undefined method requests to objects of
this class in order pylogger, tboard (first find first serve).
E.g., combinedlogger.add_scalars(...) should trigger self.tboard.add_scalars(...)
"""
for obj in [self.pylogger, self.tboard]:
if attr in dir(obj):
return getattr(obj, attr)
raise AttributeError("CombinedLogger has no attribute {}".format(attr))
def time(self, name, toggle=None):
"""record time-spans as with a stopwatch.
:param name:
:param toggle: True^=On: start time recording, False^=Off: halt rec. if None determine from current status.
:return: either start-time or last recorded interval
"""
if toggle is None:
if name in self.times.keys():
toggle = not self.times[name]["toggle"]
else:
toggle = True
if toggle:
if not name in self.times.keys():
self.times[name] = {"total": 0, "last":0}
elif self.times[name]["toggle"] == toggle:
- print("restarting running stopwatch")
+ self.info("restarting running stopwatch")
self.times[name]["last"] = time.time()
self.times[name]["toggle"] = toggle
return time.time()
else:
if toggle == self.times[name]["toggle"]:
self.info("WARNING: tried to stop stopped stop watch: {}.".format(name))
self.times[name]["last"] = time.time()-self.times[name]["last"]
self.times[name]["total"] += self.times[name]["last"]
self.times[name]["toggle"] = toggle
return self.times[name]["last"]
def get_time(self, name=None, kind="total", format=None, reset=False):
"""
:param name:
:param kind: 'total' or 'last'
:param format: None for float, "hms"/"ms" for (hours), mins, secs as string
:param reset: reset time after retrieving
:return:
"""
if name is None:
times = self.times
if reset:
self.reset_time()
return times
else:
+ if self.times[name]["toggle"]:
+ self.time(name, toggle=False)
time = self.times[name][kind]
if format == "hms":
m, s = divmod(time, 60)
h, m = divmod(m, 60)
time = "{:d}h:{:02d}m:{:02d}s".format(int(h), int(m), int(s))
elif format == "ms":
m, s = divmod(time, 60)
time = "{:02d}m:{:02d}s".format(int(m), int(s))
if reset:
self.reset_time(name)
return time
def reset_time(self, name=None):
if name is None:
self.times = {}
else:
del self.times[name]
def sysmetrics_update(self, global_step=None):
if global_step is None:
global_step = time.strftime("%x_%X")
mem = psutil.virtual_memory()
mem_used = (mem.total-mem.available)
gpu_vals = self.gpu_logger.get_vals()
rel_time = time.time()-self.sysmetrics_start_time
self.sysmetrics.loc[len(self.sysmetrics)] = [global_step, rel_time,
psutil.cpu_percent(), mem_used/1024**3, mem_used/mem.total*100,
psutil.swap_memory().used/1024**3, int(gpu_vals['gpu_graphics_util']),
*[torch.cuda.memory_allocated(d)/1024**3 for d in range(torch.cuda.device_count())],
*[torch.cuda.memory_cached(d)/1024**3 for d in range(torch.cuda.device_count())]
]
return self.sysmetrics.loc[len(self.sysmetrics)-1].to_dict()
def sysmetrics2tboard(self, metrics=None, global_step=None, suptitle=None):
tag = "per_time"
if metrics is None:
metrics = self.sysmetrics_update(global_step=global_step)
tag = "per_epoch"
if suptitle is not None:
suptitle = str(suptitle)
elif self.fold!="":
suptitle = "Fold_"+str(self.fold)
if suptitle is not None:
self.tboard.add_scalars(suptitle+"/System_Metrics/"+tag, {k:v for (k,v) in metrics.items() if (k!="global_step"
and k!="rel_time")}, global_step)
def sysmetrics_loop(self):
try:
os.nice(-19)
+ self.info("Logging system metrics with superior process priority.")
except:
- print("System-metrics logging has no superior process priority.")
+ self.info("Logging system metrics WITHOUT superior process priority.")
while True:
metrics = self.sysmetrics_update()
self.sysmetrics2tboard(metrics, global_step=metrics["rel_time"])
#print("thread alive", self.thread.is_alive())
time.sleep(self.sysmetrics_interval)
def sysmetrics_start(self, interval):
- if interval is not None:
+ if interval is not None and interval>0:
self.sysmetrics_interval = interval
self.gpu_logger = Nvidia_GPU_Logger()
self.sysmetrics_start_time = time.time()
self.thread = threading.Thread(target=self.sysmetrics_loop)
self.thread.daemon = True
self.thread.start()
def sysmetrics_save(self, out_file):
-
self.sysmetrics.to_pickle(out_file)
def metrics2tboard(self, metrics, global_step=None, suptitle=None):
"""
:param metrics: {'train': dataframe, 'val':df}, df as produced in
evaluator.py.evaluate_predictions
"""
#print("metrics", metrics)
if global_step is None:
global_step = len(metrics['train'][list(metrics['train'].keys())[0]])-1
if suptitle is not None:
suptitle = str(suptitle)
else:
suptitle = "Fold_"+str(self.fold)
for key in ['train', 'val']:
#series = {k:np.array(v[-1]) for (k,v) in metrics[key].items() if not np.isnan(v[-1]) and not 'Bin_Stats' in k}
loss_series = {}
unc_series = {}
bin_stat_series = {}
mon_met_series = {}
for tag,val in metrics[key].items():
val = val[-1] #maybe remove list wrapping, recording in evaluator?
if 'bin_stats' in tag.lower() and not np.isnan(val):
bin_stat_series["{}".format(tag.split("/")[-1])] = val
elif 'uncertainty' in tag.lower() and not np.isnan(val):
unc_series["{}".format(tag)] = val
elif 'loss' in tag.lower() and not np.isnan(val):
loss_series["{}".format(tag)] = val
elif not np.isnan(val):
mon_met_series["{}".format(tag)] = val
self.tboard.add_scalars(suptitle+"/Binary_Statistics/{}".format(key), bin_stat_series, global_step)
self.tboard.add_scalars(suptitle + "/Uncertainties/{}".format(key), unc_series, global_step)
self.tboard.add_scalars(suptitle + "/Losses/{}".format(key), loss_series, global_step)
self.tboard.add_scalars(suptitle+"/Monitor_Metrics/{}".format(key), mon_met_series, global_step)
self.tboard.add_scalars(suptitle + "/Learning_Rate", metrics["lr"], global_step)
return
def batchImgs2tboard(self, batch, results_dict, cmap, boxtype2color, img_bg=False, global_step=None):
raise NotImplementedError("not up-to-date, problem with importing plotting-file, torchvision dependency.")
if len(batch["seg"].shape)==5: #3D imgs
slice_ix = np.random.randint(batch["seg"].shape[-1])
seg_gt = plg.to_rgb(batch['seg'][:,0,:,:,slice_ix], cmap)
seg_pred = plg.to_rgb(results_dict['seg_preds'][:,0,:,:,slice_ix], cmap)
mod_img = plg.mod_to_rgb(batch["data"][:,0,:,:,slice_ix]) if img_bg else None
elif len(batch["seg"].shape)==4:
seg_gt = plg.to_rgb(batch['seg'][:,0,:,:], cmap)
seg_pred = plg.to_rgb(results_dict['seg_preds'][:,0,:,:], cmap)
mod_img = plg.mod_to_rgb(batch["data"][:,0]) if img_bg else None
else:
raise Exception("batch content has wrong format: {}".format(batch["seg"].shape))
#from here on only works in 2D
seg_gt = np.transpose(seg_gt, axes=(0,3,1,2)) #previous shp: b,x,y,c
seg_pred = np.transpose(seg_pred, axes=(0,3,1,2))
seg = np.concatenate((seg_gt, seg_pred), axis=0)
# todo replace torchvision (tv) dependency
seg = tv.utils.make_grid(torch.from_numpy(seg), nrow=2)
self.tboard.add_image("Batch seg, 1st col: gt, 2nd: pred.", seg, global_step=global_step)
if img_bg:
bg_img = np.transpose(mod_img, axes=(0,3,1,2))
else:
bg_img = seg_gt
box_imgs = plg.draw_boxes_into_batch(bg_img, results_dict["boxes"], boxtype2color)
box_imgs = tv.utils.make_grid(torch.from_numpy(box_imgs), nrow=4)
self.tboard.add_image("Batch bboxes", box_imgs, global_step=global_step)
return
def __del__(self): # otherwise might produce multiple prints e.g. in ipython console
for hdlr in self.pylogger.handlers:
hdlr.close()
self.tboard.close()
self.pylogger.handlers = []
del self.pylogger
-def get_logger(exp_dir, server_env=False, sysmetrics_interval=2):
+def get_logger(exp_dir, server_env=False, sysmetrics_interval=-1):
log_dir = os.path.join(exp_dir, "logs")
- logger = CombinedLogger('medical_detection', os.path.join(log_dir, "tboard"), server_env=server_env,
+ logger = CombinedLogger('Reg R-CNN', os.path.join(log_dir, "tboard"), server_env=server_env,
sysmetrics_interval=sysmetrics_interval)
- logger.setLevel(logging.DEBUG)
- log_file = os.path.join(log_dir, 'exec.log')
-
- logger.addHandler(logging.FileHandler(log_file))
- if not server_env:
- logger.addHandler(ColorHandler())
- else:
- logger.addHandler(logging.StreamHandler())
- logger.pylogger.propagate = False
- print('Logging to {}'.format(log_file))
-
+ print("logging to {}".format(logger.log_file))
return logger
def prep_exp(dataset_path, exp_path, server_env, use_stored_settings=True, is_training=True):
"""
I/O handling, creating of experiment folder structure. Also creates a snapshot of configs/model scripts and copies them to the exp_dir.
This way the exp_dir contains all info needed to conduct an experiment, independent to changes in actual source code. Thus, training/inference of this experiment can be started at anytime.
Therefore, the model script is copied back to the source code dir as tmp_model (tmp_backbone).
Provides robust structure for cloud deployment.
:param dataset_path: path to source code for specific data set. (e.g. medicaldetectiontoolkit/lidc_exp)
:param exp_path: path to experiment directory.
:param server_env: boolean flag. pass to configs script for cloud deployment.
:param use_stored_settings: boolean flag. When starting training: If True, starts training from snapshot in existing
experiment directory, else creates experiment directory on the fly using configs/model scripts from source code.
:param is_training: boolean flag. distinguishes train vs. inference mode.
:return: configs object.
"""
if is_training:
if use_stored_settings:
cf_file = import_module('cf', os.path.join(exp_path, 'configs.py'))
cf = cf_file.Configs(server_env)
# in this mode, previously saved model and backbone need to be found in exp dir.
if not os.path.isfile(os.path.join(exp_path, 'model.py')) or \
not os.path.isfile(os.path.join(exp_path, 'backbone.py')):
raise Exception("Selected use_stored_settings option but no model and/or backbone source files exist in exp dir.")
cf.model_path = os.path.join(exp_path, 'model.py')
cf.backbone_path = os.path.join(exp_path, 'backbone.py')
else: # this case overwrites settings files in exp dir, i.e., default_configs, configs, backbone, model
if not os.path.exists(exp_path):
os.mkdir(exp_path)
# run training with source code info and copy snapshot of model to exp_dir for later testing (overwrite scripts if exp_dir already exists.)
subprocess.call('cp {} {}'.format('default_configs.py', os.path.join(exp_path, 'default_configs.py')), shell=True)
subprocess.call('cp {} {}'.format(os.path.join(dataset_path, 'configs.py'), os.path.join(exp_path, 'configs.py')), shell=True)
cf_file = import_module('cf_file', os.path.join(dataset_path, 'configs.py'))
cf = cf_file.Configs(server_env)
subprocess.call('cp {} {}'.format(cf.model_path, os.path.join(exp_path, 'model.py')), shell=True)
subprocess.call('cp {} {}'.format(cf.backbone_path, os.path.join(exp_path, 'backbone.py')), shell=True)
if os.path.isfile(os.path.join(exp_path, "fold_ids.pickle")):
subprocess.call('rm {}'.format(os.path.join(exp_path, "fold_ids.pickle")), shell=True)
else: # testing, use model and backbone stored in exp dir.
cf_file = import_module('cf', os.path.join(exp_path, 'configs.py'))
cf = cf_file.Configs(server_env)
cf.model_path = os.path.join(exp_path, 'model.py')
cf.backbone_path = os.path.join(exp_path, 'backbone.py')
cf.exp_dir = exp_path
cf.test_dir = os.path.join(cf.exp_dir, 'test')
cf.plot_dir = os.path.join(cf.exp_dir, 'plots')
if not os.path.exists(cf.test_dir):
os.mkdir(cf.test_dir)
if not os.path.exists(cf.plot_dir):
os.mkdir(cf.plot_dir)
cf.experiment_name = exp_path.split("/")[-1]
cf.dataset_name = dataset_path
cf.server_env = server_env
cf.created_fold_id_pickle = False
return cf
class ModelSelector:
'''
saves a checkpoint after each epoch as 'last_state' (can be loaded to continue interrupted training).
saves the top-k (k=cf.save_n_models) ranked epochs. In inference, predictions of multiple epochs can be ensembled
to improve performance.
'''
def __init__(self, cf, logger):
self.cf = cf
self.saved_epochs = [-1] * cf.save_n_models
self.logger = logger
def run_model_selection(self, net, optimizer, monitor_metrics, epoch):
"""rank epoch via weighted mean from self.cf.model_selection_criteria: {criterion : weight}
:param net:
:param optimizer:
:param monitor_metrics:
:param epoch:
:return:
"""
crita = self.cf.model_selection_criteria #shorter alias
non_nan_scores = {}
for criterion in crita.keys():
#exclude first entry bc its dummy None entry
non_nan_scores[criterion] = [0 if (ii is None or np.isnan(ii)) else ii for ii in monitor_metrics['val'][criterion]][1:]
n_epochs = len(non_nan_scores[criterion])
epochs_scores = []
for e_ix in range(n_epochs):
epochs_scores.append(np.sum([weight * non_nan_scores[criterion][e_ix] for
criterion,weight in crita.items()])/len(crita.keys()))
# ranking of epochs according to model_selection_criterion
epoch_ranking = np.argsort(epochs_scores)[::-1] + 1 #epochs start at 1
# if set in configs, epochs < min_save_thresh are discarded from saving process.
epoch_ranking = epoch_ranking[epoch_ranking >= self.cf.min_save_thresh]
# check if current epoch is among the top-k epchs.
if epoch in epoch_ranking[:self.cf.save_n_models]:
if self.cf.server_env:
IO_safe(torch.save, net.state_dict(), os.path.join(self.cf.fold_dir, '{}_best_params.pth'.format(epoch)))
# save epoch_ranking to keep info for inference.
IO_safe(np.save, os.path.join(self.cf.fold_dir, 'epoch_ranking'), epoch_ranking[:self.cf.save_n_models])
else:
torch.save(net.state_dict(), os.path.join(self.cf.fold_dir, '{}_best_params.pth'.format(epoch)))
np.save(os.path.join(self.cf.fold_dir, 'epoch_ranking'), epoch_ranking[:self.cf.save_n_models])
self.logger.info(
"saving current epoch {} at rank {}".format(epoch, np.argwhere(epoch_ranking == epoch)))
# delete params of the epoch that just fell out of the top-k epochs.
for se in [int(ii.split('_')[0]) for ii in os.listdir(self.cf.fold_dir) if 'best_params' in ii]:
if se in epoch_ranking[self.cf.save_n_models:]:
subprocess.call('rm {}'.format(os.path.join(self.cf.fold_dir, '{}_best_params.pth'.format(se))),
shell=True)
self.logger.info('deleting epoch {} at rank {}'.format(se, np.argwhere(epoch_ranking == se)))
state = {
'epoch': epoch,
'state_dict': net.state_dict(),
'optimizer': optimizer.state_dict(),
}
if self.cf.server_env:
IO_safe(torch.save, state, os.path.join(self.cf.fold_dir, 'last_state.pth'))
else:
torch.save(state, os.path.join(self.cf.fold_dir, 'last_state.pth'))
def load_checkpoint(checkpoint_path, net, optimizer):
checkpoint = torch.load(checkpoint_path)
net.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
return checkpoint['epoch']
def prepare_monitoring(cf):
"""
creates dictionaries, where train/val metrics are stored.
"""
metrics = {}
# first entry for loss dict accounts for epoch starting at 1.
metrics['train'] = OrderedDict()# [(l_name, [np.nan]) for l_name in cf.losses_to_monitor] )
metrics['val'] = OrderedDict()# [(l_name, [np.nan]) for l_name in cf.losses_to_monitor] )
metric_classes = []
if 'rois' in cf.report_score_level:
metric_classes.extend([v for k, v in cf.class_dict.items()])
if hasattr(cf, "eval_bins_separately") and cf.eval_bins_separately:
metric_classes.extend([v for k, v in cf.bin_dict.items()])
if 'patient' in cf.report_score_level:
metric_classes.extend(['patient_'+cf.class_dict[cf.patient_class_of_interest]])
if hasattr(cf, "eval_bins_separately") and cf.eval_bins_separately:
metric_classes.extend(['patient_' + cf.bin_dict[cf.patient_bin_of_interest]])
for cl in metric_classes:
for m in cf.metrics:
metrics['train'][cl + '_' + m] = [np.nan]
metrics['val'][cl + '_' + m] = [np.nan]
return metrics
class _AnsiColorizer(object):
"""
A colorizer is an object that loosely wraps around a stream, allowing
callers to write text to the stream in a particular color.
Colorizer classes must implement C{supported()} and C{write(text, color)}.
"""
_colors = dict(black=30, red=31, green=32, yellow=33,
blue=34, magenta=35, cyan=36, white=37, default=39)
def __init__(self, stream):
self.stream = stream
@classmethod
def supported(cls, stream=sys.stdout):
"""
A class method that returns True if the current platform supports
coloring terminal output using this method. Returns False otherwise.
"""
if not stream.isatty():
return False # auto color only on TTYs
try:
import curses
except ImportError:
return False
else:
try:
try:
return curses.tigetnum("colors") > 2
except curses.error:
curses.setupterm()
return curses.tigetnum("colors") > 2
except:
raise
# guess false in case of error
return False
def write(self, text, color):
"""
Write the given text to the stream in the given color.
@param text: Text to be written to the stream.
@param color: A string label for a color. e.g. 'red', 'white'.
"""
color = self._colors[color]
self.stream.write('\x1b[%sm%s\x1b[0m' % (color, text))
class ColorHandler(logging.StreamHandler):
def __init__(self, stream=sys.stdout):
super(ColorHandler, self).__init__(_AnsiColorizer(stream))
def emit(self, record):
msg_colors = {
logging.DEBUG: "green",
logging.INFO: "default",
logging.WARNING: "red",
logging.ERROR: "red"
}
color = msg_colors.get(record.levelno, "blue")
self.stream.write(record.msg + "\n", color)
diff --git a/utils/model_utils.py b/utils/model_utils.py
index 326cdf6..d415f84 100644
--- a/utils/model_utils.py
+++ b/utils/model_utils.py
@@ -1,1472 +1,1454 @@
#!/usr/bin/env python
# Copyright 2019 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 time
import warnings
warnings.filterwarnings('ignore', '.*From scipy 0.13.0, the output shape of zoom()*')
import numpy as np
-import math
import scipy.misc
import scipy.ndimage
from scipy.ndimage.measurements import label as lb
import torch
-from torch.autograd import Variable
-
-# from cuda_functions.nms_2D.pth_nms import nms_gpu as nms_2D
-# from cuda_functions.nms_3D.pth_nms import nms_gpu as nms_3D
-# from cuda_functions.roi_align_2D.roi_align.crop_and_resize import CropAndResizeFunction as ra2D
-# from cuda_functions.roi_align_3D.roi_align.crop_and_resize import CropAndResizeFunction as ra3D
from custom_extensions.nms import nms
from custom_extensions.roi_align import roi_align
############################################################
# Segmentation Processing
############################################################
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 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')
elif dim == 3:
y_ohe = np.zeros((y.shape[0], n_classes, y.shape[2], y.shape[3], y.shape[4])).astype('int32')
else:
raise Exception("invalid dimensions {} encountered".format(y.shape))
for cl in np.arange(n_classes):
y_ohe[:, cl][y[:, 0] == cl] = 1
return y_ohe
def dice_per_batch_inst_and_class(pred, y, n_classes, convert_to_ohe=True, smooth=1e-8):
- #actually per batch_instance not batch
'''
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)
'''
if convert_to_ohe:
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)
dice = (2.0*intersect + smooth) / (denominator + smooth)
return dice
def dice_per_batch_and_class(pred, targ, n_classes, convert_to_ohe=True, smooth=1e-8):
'''
computes dice scores per batch and class.
:param pred: prediction array of shape (b, 1, y, x, (z)) (e.g. softmax prediction with argmax over dim 1)
:param targ: ground truth array of shape (b, 1, y, x, (z)) (contains int [0, ..., n_classes])
:param n_classes: int
:param smooth: Laplacian smooth, https://en.wikipedia.org/wiki/Additive_smoothing
:return: dice scores of shape (b, c)
'''
if convert_to_ohe:
pred = get_one_hot_encoding(pred, n_classes)
targ = get_one_hot_encoding(targ, n_classes)
axes = (0, *list(range(2, len(pred.shape)))) #(0,2,3(,4))
intersect = np.sum(pred * targ, axis=axes)
denominator = np.sum(pred, axis=axes) + np.sum(targ, axis=axes)
dice = (2.0 * intersect + smooth) / (denominator + smooth)
assert dice.shape==(n_classes,), "dice shp {}".format(dice.shape)
return dice
def batch_dice(pred, y, false_positive_weight=1.0, eps=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).
: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 mena of foreground scores.
'''
# todo also use additive smooth here instead of eps?
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 / (denom + eps))[1:]) #only fg dice here.
if 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 / (denom + eps))[1:]) #only fg dice here.
else:
raise ValueError('wrong input dimension in dice loss')
############################################################
# 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.
:return: (#boxes1, #boxes2), ious of each box of 1 machted with each of 2
"""
# 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(boxes2.shape[0]):
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]) #only y,x
full_mask[y1:y2, x1:x2] = mask
return full_mask
def unmold_mask_2D_torch(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).float(), (x2 - x1).float()]
zoom_factor = [i / j for i, j in zip(out_zoom, mask.shape)]
mask = mask.unsqueeze(0).unsqueeze(0)
mask = torch.nn.functional.interpolate(mask, scale_factor=zoom_factor)
mask = mask[0][0]
#mask = scipy.ndimage.zoom(mask.cpu().numpy(), zoom_factor, order=1).astype(np.float32)
#mask = torch.from_numpy(mask).cuda()
# Put the mask in the right location.
full_mask = torch.zeros(image_shape[:2]) # only y,x
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
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 + 1) * (y2 - y1 + 1)
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 + 1)
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 + 1)
w = np.maximum(0.0, xx2 - xx1 + 1)
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 + 1)
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.
#print("iou keep {}: {}, non_matches {}".format(i, iou, order[non_matches]))
order = order[non_matches]
keep.append(i)
#print("total keep", keep)
return keep
############################################################
# M-RCNN
############################################################
def refine_proposals(rpn_pred_probs, rpn_pred_deltas, proposal_count, batch_anchors, cf):
"""
Receives anchor scores and selects a subset to pass as proposals
to the second stage. Filtering is done based on anchor scores and
non-max suppression to remove overlaps. It also applies bounding
box refinment details to anchors.
:param rpn_pred_probs: (b, n_anchors, 2)
:param rpn_pred_deltas: (b, n_anchors, (y, x, (z), log(h), log(w), (log(d))))
:return: batch_normalized_props: Proposals in normalized coordinates (b, proposal_count, (y1, x1, y2, x2, (z1), (z2), score))
:return: batch_out_proposals: Box coords + RPN foreground scores
for monitoring/plotting (b, proposal_count, (y1, x1, y2, x2, (z1), (z2), score))
"""
std_dev = torch.from_numpy(cf.rpn_bbox_std_dev[None]).float().cuda()
norm = torch.from_numpy(cf.scale).float().cuda()
anchors = batch_anchors.clone()
batch_scores = rpn_pred_probs[:, :, 1]
# norm deltas
batch_deltas = rpn_pred_deltas * std_dev
batch_normalized_props = []
batch_out_proposals = []
# loop over batch dimension.
for ix in range(batch_scores.shape[0]):
scores = batch_scores[ix]
deltas = batch_deltas[ix]
# improve performance by trimming to top anchors by score
# and doing the rest on the smaller subset.
pre_nms_limit = min(cf.pre_nms_limit, anchors.size()[0])
scores, order = scores.sort(descending=True)
order = order[:pre_nms_limit]
scores = scores[:pre_nms_limit]
deltas = deltas[order, :]
# apply deltas to anchors to get refined anchors and filter with non-maximum suppression.
if batch_deltas.shape[-1] == 4:
boxes = apply_box_deltas_2D(anchors[order, :], deltas)
boxes = clip_boxes_2D(boxes, cf.window)
else:
boxes = apply_box_deltas_3D(anchors[order, :], deltas)
boxes = clip_boxes_3D(boxes, cf.window)
# boxes are y1,x1,y2,x2, torchvision-nms requires x1,y1,x2,y2, but consistent swap x<->y is irrelevant.
keep = nms.nms(boxes, scores, cf.rpn_nms_threshold)
keep = keep[:proposal_count]
boxes = boxes[keep, :]
rpn_scores = scores[keep][:, None]
# pad missing boxes with 0.
if boxes.shape[0] < proposal_count:
n_pad_boxes = proposal_count - boxes.shape[0]
zeros = torch.zeros([n_pad_boxes, boxes.shape[1]]).cuda()
boxes = torch.cat([boxes, zeros], dim=0)
zeros = torch.zeros([n_pad_boxes, rpn_scores.shape[1]]).cuda()
rpn_scores = torch.cat([rpn_scores, zeros], dim=0)
# concat box and score info for monitoring/plotting.
batch_out_proposals.append(torch.cat((boxes, rpn_scores), 1).cpu().data.numpy())
# normalize dimensions to range of 0 to 1.
normalized_boxes = boxes / norm
assert torch.all(normalized_boxes <= 1), "normalized box coords >1 found"
# add again batch dimension
batch_normalized_props.append(torch.cat((normalized_boxes, rpn_scores), 1).unsqueeze(0))
batch_normalized_props = torch.cat(batch_normalized_props)
batch_out_proposals = np.array(batch_out_proposals)
return batch_normalized_props, batch_out_proposals
def pyramid_roi_align(feature_maps, rois, pool_size, pyramid_levels, dim):
"""
Implements ROI Pooling on multiple levels of the feature pyramid.
:param feature_maps: list of feature maps, each of shape (b, c, y, x , (z))
:param rois: proposals (normalized coords.) as returned by RPN. contain info about original batch element allocation.
(n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ixs)
:param pool_size: list of poolsizes in dims: [x, y, (z)]
:param pyramid_levels: list. [0, 1, 2, ...]
:return: pooled: pooled feature map rois (n_proposals, c, poolsize_y, poolsize_x, (poolsize_z))
Output:
Pooled regions in the shape: [num_boxes, height, width, channels].
The width and height are those specific in the pool_shape in the layer
constructor.
"""
boxes = rois[:, :dim*2]
batch_ixs = rois[:, dim*2]
# Assign each ROI to a level in the pyramid based on the ROI area.
if dim == 2:
y1, x1, y2, x2 = boxes.chunk(4, dim=1)
else:
y1, x1, y2, x2, z1, z2 = boxes.chunk(6, dim=1)
h = y2 - y1
w = x2 - x1
# Equation 1 in https://arxiv.org/abs/1612.03144. Account for
# the fact that our coordinates are normalized here.
# divide sqrt(h*w) by 1 instead image_area.
- roi_level = (4 + log2(torch.sqrt(h*w))).round().int().clamp(pyramid_levels[0], pyramid_levels[-1])
+ roi_level = (4 + torch.log2(torch.sqrt(h*w))).round().int().clamp(pyramid_levels[0], pyramid_levels[-1])
# if Pyramid contains additional level P6, adapt the roi_level assignment accordingly.
if len(pyramid_levels) == 5:
roi_level[h*w > 0.65] = 5
# Loop through levels and apply ROI pooling to each.
pooled = []
box_to_level = []
fmap_shapes = [f.shape for f in feature_maps]
for level_ix, level in enumerate(pyramid_levels):
ix = roi_level == level
if not ix.any():
continue
ix = torch.nonzero(ix)[:, 0]
level_boxes = boxes[ix, :]
# re-assign rois to feature map of original batch element.
ind = batch_ixs[ix].int()
# Keep track of which box is mapped to which level
box_to_level.append(ix)
# Stop gradient propogation to ROI proposals
level_boxes = level_boxes.detach()
if len(pool_size) == 2:
# remap to feature map coordinate system
y_exp, x_exp = fmap_shapes[level_ix][2:] # exp = expansion
level_boxes.mul_(torch.tensor([y_exp, x_exp, y_exp, x_exp], dtype=torch.float32).cuda())
pooled_features = roi_align.roi_align_2d(feature_maps[level_ix],
torch.cat((ind.unsqueeze(1).float(), level_boxes), dim=1),
pool_size)
else:
y_exp, x_exp, z_exp = fmap_shapes[level_ix][2:]
level_boxes.mul_(torch.tensor([y_exp, x_exp, y_exp, x_exp, z_exp, z_exp], dtype=torch.float32).cuda())
pooled_features = roi_align.roi_align_3d(feature_maps[level_ix],
torch.cat((ind.unsqueeze(1).float(), level_boxes), dim=1),
pool_size)
pooled.append(pooled_features)
# Pack pooled features into one tensor
pooled = torch.cat(pooled, dim=0)
# Pack box_to_level mapping into one array and add another
# column representing the order of pooled boxes
box_to_level = torch.cat(box_to_level, dim=0)
# Rearrange pooled features to match the order of the original boxes
_, box_to_level = torch.sort(box_to_level)
pooled = pooled[box_to_level, :, :]
return pooled
def refine_detections(cf, batch_ixs, rois, deltas, scores, regressions):
"""
Refine classified proposals (apply deltas to rpn rois), filter overlaps (nms) and return final detections.
:param rois: (n_proposals, 2 * dim) normalized boxes as proposed by RPN. n_proposals = batch_size * POST_NMS_ROIS
:param deltas: (n_proposals, n_classes, 2 * dim) box refinement deltas as predicted by mrcnn bbox regressor.
:param batch_ixs: (n_proposals) batch element assignment info for re-allocation.
:param scores: (n_proposals, n_classes) probabilities for all classes per roi as predicted by mrcnn classifier.
:param regressions: (n_proposals, n_classes, regression_features (+1 for uncertainty if predicted) regression vector
:return: result: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score, *regression vector features))
"""
# class IDs per ROI. Since scores of all classes are of interest (not just max class), all are kept at this point.
class_ids = []
fg_classes = cf.head_classes - 1
# repeat vectors to fill in predictions for all foreground classes.
for ii in range(1, fg_classes + 1):
class_ids += [ii] * rois.shape[0]
class_ids = torch.from_numpy(np.array(class_ids)).cuda()
batch_ixs = batch_ixs.repeat(fg_classes)
rois = rois.repeat(fg_classes, 1)
deltas = deltas.repeat(fg_classes, 1, 1)
scores = scores.repeat(fg_classes, 1)
regressions = regressions.repeat(fg_classes, 1, 1)
# get class-specific scores and bounding box deltas
idx = torch.arange(class_ids.size()[0]).long().cuda()
# using idx instead of slice [:,] squashes first dimension.
#len(class_ids)>scores.shape[1] --> probs is broadcasted by expansion from fg_classes-->len(class_ids)
batch_ixs = batch_ixs[idx]
deltas_specific = deltas[idx, class_ids]
class_scores = scores[idx, class_ids]
regressions = regressions[idx, class_ids]
# 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 = apply_box_deltas_2D(rois, deltas_specific * std_dev) * scale if cf.dim == 2 else \
apply_box_deltas_3D(rois, deltas_specific * std_dev) * scale
# round and cast to int since we're dealing with pixels now
refined_rois = clip_to_window(cf.window, refined_rois)
refined_rois = torch.round(refined_rois)
# filter out low confidence boxes
keep = idx
keep_bool = (class_scores >= cf.model_min_confidence)
if not 0 in torch.nonzero(keep_bool).size():
score_keep = torch.nonzero(keep_bool)[:, 0]
pre_nms_class_ids = class_ids[score_keep]
pre_nms_rois = refined_rois[score_keep]
pre_nms_scores = class_scores[score_keep]
pre_nms_batch_ixs = batch_ixs[score_keep]
for j, b in enumerate(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(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, :]
class_keep = nms.nms(ix_rois, ix_scores, cf.detection_nms_threshold)
# map indices back.
class_keep = keep[score_keep[bixs[ixs[order[class_keep]]]]]
# merge indices over classes for current batch element
b_keep = class_keep if i == 0 else unique1d(torch.cat((b_keep, class_keep)))
# only keep top-k boxes of current batch-element
top_ids = class_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 unique1d(torch.cat((batch_keep, b_keep)))
keep = batch_keep
else:
keep = torch.tensor([0]).long().cuda()
# arrange output
output = [refined_rois[keep], batch_ixs[keep].unsqueeze(1)]
output += [class_ids[keep].unsqueeze(1).float(), class_scores[keep].unsqueeze(1)]
output += [regressions[keep]]
result = torch.cat(output, dim=1)
# shape: (n_keeps, catted feats), catted feats: [0:dim*2] are box_coords, [dim*2] are batch_ics,
# [dim*2+1] are class_ids, [dim*2+2] are scores, [dim*2+3:] are regression vector features (incl uncertainty)
return result
def loss_example_mining(cf, batch_proposals, batch_gt_boxes, batch_gt_masks, batch_roi_scores,
batch_gt_class_ids, batch_gt_regressions):
"""
Subsamples proposals for mrcnn losses and generates targets. Sampling is done per batch element, seems to have positive
effects on training, as opposed to sampling over entire batch. Negatives are sampled via stochastic hard-example mining
(SHEM), where a number of negative proposals is drawn from larger pool of highest scoring proposals for stochasticity.
Scoring is obtained here as the max over all foreground probabilities as returned by mrcnn_classifier (worked better than
loss-based class-balancing methods like "online hard-example mining" or "focal loss".)
Classification-regression duality: regressions can be given along with classes (at least fg/bg, only class scores
are used for ranking).
:param batch_proposals: (n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ixs).
boxes as proposed by RPN. n_proposals here is determined by batch_size * POST_NMS_ROIS.
:param mrcnn_class_logits: (n_proposals, n_classes)
:param batch_gt_boxes: list over batch elements. Each element is a list over the corresponding roi target coordinates.
:param batch_gt_masks: list over batch elements. Each element is binary mask of shape (n_gt_rois, y, x, (z), c)
:param batch_gt_class_ids: list over batch elements. Each element is a list over the corresponding roi target labels.
if no classes predicted (only fg/bg from RPN): expected as pseudo classes [0, 1] for bg, fg.
:param batch_gt_regressions: list over b elements. Each element is a regression target vector. if None--> pseudo
:return: sample_indices: (n_sampled_rois) indices of sampled proposals to be used for loss functions.
:return: target_class_ids: (n_sampled_rois)containing target class labels of sampled proposals.
:return: target_deltas: (n_sampled_rois, 2 * dim) containing target deltas of sampled proposals for box refinement.
:return: target_masks: (n_sampled_rois, y, x, (z)) containing target masks of sampled proposals.
"""
# normalization of target coordinates
#global sample_regressions
if cf.dim == 2:
h, w = cf.patch_size
scale = torch.from_numpy(np.array([h, w, h, w])).float().cuda()
else:
h, w, z = cf.patch_size
scale = torch.from_numpy(np.array([h, w, h, w, z, z])).float().cuda()
positive_count = 0
negative_count = 0
sample_positive_indices = []
sample_negative_indices = []
sample_deltas = []
sample_masks = []
sample_class_ids = []
if batch_gt_regressions is not None:
sample_regressions = []
else:
target_regressions = torch.FloatTensor().cuda()
# loop over batch and get positive and negative sample rois.
for b in range(len(batch_gt_boxes)):
gt_masks = torch.from_numpy(batch_gt_masks[b]).float().cuda()
gt_class_ids = torch.from_numpy(batch_gt_class_ids[b]).int().cuda()
if batch_gt_regressions is not None:
gt_regressions = torch.from_numpy(batch_gt_regressions[b]).float().cuda()
#if np.any(batch_gt_class_ids[b] > 0): # skip roi selection for no gt images.
if np.any([len(coords)>0 for coords in batch_gt_boxes[b]]):
gt_boxes = torch.from_numpy(batch_gt_boxes[b]).float().cuda() / scale
else:
gt_boxes = torch.FloatTensor().cuda()
# get proposals and indices of current batch element.
proposals = batch_proposals[batch_proposals[:, -1] == b][:, :-1]
batch_element_indices = torch.nonzero(batch_proposals[:, -1] == b).squeeze(1)
# Compute overlaps matrix [proposals, gt_boxes]
if not 0 in gt_boxes.size():
if gt_boxes.shape[1] == 4:
assert cf.dim == 2, "gt_boxes shape {} doesnt match cf.dim{}".format(gt_boxes.shape, cf.dim)
overlaps = bbox_overlaps_2D(proposals, gt_boxes)
else:
assert cf.dim == 3, "gt_boxes shape {} doesnt match cf.dim{}".format(gt_boxes.shape, cf.dim)
overlaps = bbox_overlaps_3D(proposals, gt_boxes)
# Determine positive and negative ROIs
roi_iou_max = torch.max(overlaps, dim=1)[0]
# 1. Positive ROIs are those with >= 0.5 IoU with a GT box
positive_roi_bool = roi_iou_max >= (0.5 if cf.dim == 2 else 0.3)
# 2. Negative ROIs are those with < 0.1 with every GT box.
negative_roi_bool = roi_iou_max < (0.1 if cf.dim == 2 else 0.01)
else:
positive_roi_bool = torch.FloatTensor().cuda()
negative_roi_bool = torch.from_numpy(np.array([1]*proposals.shape[0])).cuda()
# Sample Positive ROIs
if not 0 in torch.nonzero(positive_roi_bool).size():
positive_indices = torch.nonzero(positive_roi_bool).squeeze(1)
positive_samples = int(cf.train_rois_per_image * cf.roi_positive_ratio)
rand_idx = torch.randperm(positive_indices.size()[0])
rand_idx = rand_idx[:positive_samples].cuda()
positive_indices = positive_indices[rand_idx]
positive_samples = positive_indices.size()[0]
positive_rois = proposals[positive_indices, :]
# Assign positive ROIs to GT boxes.
positive_overlaps = overlaps[positive_indices, :]
roi_gt_box_assignment = torch.max(positive_overlaps, dim=1)[1]
roi_gt_boxes = gt_boxes[roi_gt_box_assignment, :]
roi_gt_class_ids = gt_class_ids[roi_gt_box_assignment]
if batch_gt_regressions is not None:
roi_gt_regressions = gt_regressions[roi_gt_box_assignment]
# Compute bbox refinement targets for positive ROIs
deltas = box_refinement(positive_rois, roi_gt_boxes)
std_dev = torch.from_numpy(cf.bbox_std_dev).float().cuda()
deltas /= std_dev
roi_masks = gt_masks[roi_gt_box_assignment].unsqueeze(1) # .squeeze(-1)
assert roi_masks.shape[-1] == 1
# Compute mask targets
boxes = positive_rois
box_ids = torch.arange(roi_masks.shape[0]).cuda().unsqueeze(1).float()
if len(cf.mask_shape) == 2:
# todo what are the dims of roi_masks? (n_matched_boxes_with_gts, 1 (dummy channel dim), y,x, 1 (WHY?))
masks = roi_align.roi_align_2d(roi_masks,
torch.cat((box_ids, boxes), dim=1),
cf.mask_shape)
else:
masks = roi_align.roi_align_3d(roi_masks,
torch.cat((box_ids, boxes), dim=1),
cf.mask_shape)
masks = masks.squeeze(1)
# Threshold mask pixels at 0.5 to have GT masks be 0 or 1 to use with
# binary cross entropy loss.
masks = torch.round(masks)
sample_positive_indices.append(batch_element_indices[positive_indices])
sample_deltas.append(deltas)
sample_masks.append(masks)
sample_class_ids.append(roi_gt_class_ids)
if batch_gt_regressions is not None:
sample_regressions.append(roi_gt_regressions)
positive_count += positive_samples
else:
positive_samples = 0
# Sample negative ROIs. Add enough to maintain positive:negative ratio, but at least 1. Sample via SHEM.
if not 0 in torch.nonzero(negative_roi_bool).size():
negative_indices = torch.nonzero(negative_roi_bool).squeeze(1)
r = 1.0 / cf.roi_positive_ratio
b_neg_count = np.max((int(r * positive_samples - positive_samples), 1))
roi_scores_neg = batch_roi_scores[batch_element_indices[negative_indices]]
raw_sampled_indices = shem(roi_scores_neg, b_neg_count, cf.shem_poolsize)
sample_negative_indices.append(batch_element_indices[negative_indices[raw_sampled_indices]])
negative_count += raw_sampled_indices.size()[0]
if len(sample_positive_indices) > 0:
target_deltas = torch.cat(sample_deltas)
target_masks = torch.cat(sample_masks)
target_class_ids = torch.cat(sample_class_ids)
if batch_gt_regressions is not None:
target_regressions = torch.cat(sample_regressions)
# Pad target information with zeros for negative ROIs.
if positive_count > 0 and negative_count > 0:
sample_indices = torch.cat((torch.cat(sample_positive_indices), torch.cat(sample_negative_indices)), dim=0)
zeros = torch.zeros(negative_count, cf.dim * 2).cuda()
target_deltas = torch.cat([target_deltas, zeros], dim=0)
zeros = torch.zeros(negative_count, *cf.mask_shape).cuda()
target_masks = torch.cat([target_masks, zeros], dim=0)
zeros = torch.zeros(negative_count).int().cuda()
target_class_ids = torch.cat([target_class_ids, zeros], dim=0)
if batch_gt_regressions is not None:
# regression targets need to have 0 as background/negative with below practice
if 'regression_bin' in cf.prediction_tasks:
zeros = torch.zeros(negative_count, dtype=torch.float).cuda()
else:
zeros = torch.zeros(negative_count, cf.regression_n_features, dtype=torch.float).cuda()
target_regressions = torch.cat([target_regressions, zeros], dim=0)
elif positive_count > 0:
sample_indices = torch.cat(sample_positive_indices)
elif negative_count > 0:
sample_indices = torch.cat(sample_negative_indices)
target_deltas = torch.zeros(negative_count, cf.dim * 2).cuda()
target_masks = torch.zeros(negative_count, *cf.mask_shape).cuda()
target_class_ids = torch.zeros(negative_count).int().cuda()
if batch_gt_regressions is not None:
if 'regression_bin' in cf.prediction_tasks:
target_regressions = torch.zeros(negative_count, dtype=torch.float).cuda()
else:
target_regressions = torch.zeros(negative_count, cf.regression_n_features, dtype=torch.float).cuda()
else:
sample_indices = torch.LongTensor().cuda()
target_class_ids = torch.IntTensor().cuda()
target_deltas = torch.FloatTensor().cuda()
target_masks = torch.FloatTensor().cuda()
target_regressions = torch.FloatTensor().cuda()
return sample_indices, target_deltas, target_masks, target_class_ids, target_regressions
############################################################
# 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 , for conformity with retina nets: scale entries need to be list, 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
logger.info("anchor scales {} and feature map shapes {}".format(scales, feature_shapes))
expected_anchors = [np.prod(feature_shapes[level]) * len(ratios) * len(scales['xy'][level]) for level in pyramid_levels]
anchors = []
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))
elif len(feature_shapes[level]) == 3:
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))
else:
raise Exception("invalid feature_shapes[{}] size {}".format(level, feature_shapes[level]))
logger.info("level {}: expected anchors {}, built anchors {}.".format(level, expected_anchors[lix], anchors[-1].shape))
out_anchors = np.concatenate(anchors, axis=0)
logger.info("Total: expected anchors {}, built anchors {}.".format(np.sum(expected_anchors), out_anchors.shape))
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
from matplotlib import pyplot as plt
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]
#--> expects x1<x2 & y1<y2
- zeros = Variable(torch.zeros(y1.size()[0]), requires_grad=False)
+ zeros = 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
assert torch.all(iou<=1), "iou score>1 produced in bbox_overlaps_2D"
overlaps = iou.view(boxes2_repeat, boxes1_repeat) #--> per gt box: ious of all proposal boxes with that gt box
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)
+ zeros = 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
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_thresh_ixs = np.argwhere(anchor_iou_max >= anchor_matching_iou)
anchor_class_matches[above_thresh_ixs] = gt_class_ids[anchor_iou_argmax[above_thresh_ixs]]
# Subsample to balance positive anchors.
ids = np.where(anchor_class_matches > 0)[0]
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
# Leave all negative proposals negative for now and sample from them later 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]
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
############################################################
# Connected Componenent Analysis
############################################################
def get_coords(binary_mask, n_components, dim):
"""
loops over batch to perform connected component analysis on binary input mask. computes box coordinates around
n_components - biggest components (rois).
:param binary_mask: (b, y, x, (z)). binary mask for one specific foreground class.
:param n_components: int. number of components to extract per batch element and class.
:return: coords (b, n, (y1, x1, y2, x2 (,z1, z2))
:return: batch_components (b, n, (y1, x1, y2, x2, (z1), (z2))
"""
assert len(binary_mask.shape)==dim+1
binary_mask = binary_mask.astype('uint8')
batch_coords = []
batch_components = []
for ix,b in enumerate(binary_mask):
clusters, n_cands = lb(b) # performs connected component analysis.
uniques, counts = np.unique(clusters, return_counts=True)
keep_uniques = uniques[1:][np.argsort(counts[1:])[::-1]][:n_components] #only keep n_components largest components
p_components = np.array([(clusters == ii) * 1 for ii in keep_uniques]) # separate clusters and concat
p_coords = []
if p_components.shape[0] > 0:
for roi in p_components:
mask_ixs = np.argwhere(roi != 0)
# get coordinates around component.
roi_coords = [np.min(mask_ixs[:, 0]) - 1, np.min(mask_ixs[:, 1]) - 1, np.max(mask_ixs[:, 0]) + 1,
np.max(mask_ixs[:, 1]) + 1]
if dim == 3:
roi_coords += [np.min(mask_ixs[:, 2]), np.max(mask_ixs[:, 2])+1]
p_coords.append(roi_coords)
p_coords = np.array(p_coords)
#clip coords.
p_coords[p_coords < 0] = 0
p_coords[:, :4][p_coords[:, :4] > binary_mask.shape[-2]] = binary_mask.shape[-2]
if dim == 3:
p_coords[:, 4:][p_coords[:, 4:] > binary_mask.shape[-1]] = binary_mask.shape[-1]
batch_coords.append(p_coords)
batch_components.append(p_components)
return batch_coords, batch_components
# noinspection PyCallingNonCallable
def get_coords_gpu(binary_mask, n_components, dim):
"""
loops over batch to perform connected component analysis on binary input mask. computes box coordiantes around
n_components - biggest components (rois).
:param binary_mask: (b, y, x, (z)). binary mask for one specific foreground class.
:param n_components: int. number of components to extract per batch element and class.
:return: coords (b, n, (y1, x1, y2, x2 (,z1, z2))
:return: batch_components (b, n, (y1, x1, y2, x2, (z1), (z2))
"""
raise Exception("throws floating point exception")
assert len(binary_mask.shape)==dim+1
binary_mask = binary_mask.type(torch.uint8)
batch_coords = []
batch_components = []
for ix,b in enumerate(binary_mask):
clusters, n_cands = lb(b.cpu().data.numpy()) # peforms connected component analysis.
clusters = torch.from_numpy(clusters).cuda()
uniques = torch.unique(clusters)
counts = torch.stack([(clusters==unique).sum() for unique in uniques])
keep_uniques = uniques[1:][torch.sort(counts[1:])[1].flip(0)][:n_components] #only keep n_components largest components
p_components = torch.cat([(clusters == ii).unsqueeze(0) for ii in keep_uniques]).cuda() # separate clusters and concat
p_coords = []
if p_components.shape[0] > 0:
for roi in p_components:
mask_ixs = torch.nonzero(roi)
# get coordinates around component.
roi_coords = [torch.min(mask_ixs[:, 0]) - 1, torch.min(mask_ixs[:, 1]) - 1,
torch.max(mask_ixs[:, 0]) + 1,
torch.max(mask_ixs[:, 1]) + 1]
if dim == 3:
roi_coords += [torch.min(mask_ixs[:, 2]), torch.max(mask_ixs[:, 2])+1]
p_coords.append(roi_coords)
p_coords = torch.tensor(p_coords)
#clip coords.
p_coords[p_coords < 0] = 0
p_coords[:, :4][p_coords[:, :4] > binary_mask.shape[-2]] = binary_mask.shape[-2]
if dim == 3:
p_coords[:, 4:][p_coords[:, 4:] > binary_mask.shape[-1]] = binary_mask.shape[-1]
batch_coords.append(p_coords)
batch_components.append(p_components)
return batch_coords, batch_components
############################################################
# Pytorch Utility Functions
############################################################
def unique1d(tensor):
"""discard all elements of tensor that occur more than once; make tensor unique.
:param tensor:
:return:
"""
if tensor.size()[0] == 0 or tensor.size()[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]
-def log2(x):
- """Implementatin of Log2. Pytorch doesn't have a native implementation."""
- 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, 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*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 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((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()]
############################################################
# Weight Init
############################################################
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 [torch.nn.Conv2d, torch.nn.Conv3d,
torch.nn.ConvTranspose2d,
torch.nn.ConvTranspose3d,
torch.nn.Linear]]:
if init_type == 'xavier_uniform':
torch.nn.init.xavier_uniform_(m.weight.data)
if m.bias is not None:
m.bias.data.zero_()
elif init_type == 'xavier_normal':
torch.nn.init.xavier_normal_(m.weight.data)
if m.bias is not None:
m.bias.data.zero_()
elif init_type == "kaiming_uniform":
torch.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 = torch.nn.init._calculate_fan_in_and_fan_out(m.weight.data)
bound = 1 / np.sqrt(fan_out)
torch.nn.init.uniform_(m.bias, -bound, bound)
elif init_type == "kaiming_normal":
torch.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 = torch.nn.init._calculate_fan_in_and_fan_out(m.weight.data)
bound = 1 / np.sqrt(fan_out)
torch.nn.init.normal_(m.bias, -bound, bound)
net.logger.info("applied {} weight init.".format(init_type))
\ No newline at end of file