diff --git a/README.md b/README.md
index a7c4621..f3d0023 100644
--- a/README.md
+++ b/README.md
@@ -1,139 +1,145 @@
[
](https://join.slack.com/t/mdtoolkit/shared_invite/enQtNTQ3MjY2MzE0MDg2LWNjY2I2Njc5MTY0NmM0ZWIxNmQwZDRhYzk2MDdhM2QxYjliYTcwYzhkNTAxYmRkMDA0MjcyNDMyYjllNTZhY2M)
Copyright © German Cancer Research Center (DKFZ), Division of Medical Image Computing (MIC). Please make sure that your usage of this code is in compliance with the code license.
## Release Notes
-**v0.1.0**: Updates to python 3.7, torch 1.4.0, torchvision 0.5.0, entailing a change in custom extensions NMS and RoIAlign
- (now in C++ and CUDA). Scalar monitoring is changed to torch-included tensorboard. Added qualitative example
- plots for validation and testing. Default optimizer is changed to AdamW instead of Adam to account for
- fix in weight-decay handling, norms and biases can optionally be excluded from weight decay. Introduced
- optional dynamic learning-rate scheduling. A specific CUDA device can be selected via script argument. Added
- tutorial.\
+**v0.1.0**:
+- Updates to python 3.7, torch 1.4.0, torchvision 0.5.0, entailing a change in custom extensions NMS and RoIAlign
+ (now in C++ and CUDA).
+- Scalar monitoring is changed to torch-included tensorboard.
+- Added qualitative example plots for validation and testing.
+- Default optimizer is changed to AdamW instead of Adam to account for fix in weight-decay handling,
+norms and biases can optionally be excluded from weight decay.
+- Introduced optional dynamic learning-rate scheduling.
+- A specific CUDA device can now be selected via script argument.
+- Inside the models, GT class labels identification is changed from `'roi_labels'` to `'class_target'` to streamline naming scheme.
+- Added dataset [tutorial](experiments/tutorial.md).
+
**v0.0.2**: Small fixes mainly regarding server-env settings (cluster deployment).\
**v0.0.1**: Original framework as used for the corresponding paper, with Python 3.6 and torch 0.4.1 dependencies,
custom extensions NMS and RoIAlign in C and CUDA, scalar monitoring via plot files.
## Overview
This is a comprehensive framework for object detection featuring:
- 2D + 3D implementations of prevalent object detectors: e.g. Mask R-CNN [1], Retina Net [2], Retina U-Net [3].
- Modular and light-weight structure ensuring sharing of all processing steps (incl. backbone architecture) for comparability of models.
- training with bounding box and/or pixel-wise annotations.
- dynamic patching and tiling of 2D + 3D images (for training and inference).
- weighted consolidation of box predictions across patch-overlaps, ensembles, and dimensions [3].
- monitoring + evaluation simultaneously on object and patient level.
- 2D + 3D output visualizations.
- integration of COCO mean average precision metric [5].
- integration of MIC-DKFZ batch generators for extensive data augmentation [6].
- easy modification to evaluation of instance segmentation and/or semantic segmentation.
[1] He, Kaiming, et al. "Mask R-CNN" ICCV, 2017
[2] Lin, Tsung-Yi, et al. "Focal Loss for Dense Object Detection" TPAMI, 2018.
[3] Jaeger, Paul et al. "Retina U-Net: Embarrassingly Simple Exploitation
of Segmentation Supervision for Medical Object Detection" , 2018
[5] https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocotools/cocoeval.py
[6] https://github.com/MIC-DKFZ/batchgenerators
-A tutorial on how to add your own data set can be found under `experiments/tutorial.md`.
+A tutorial on how to add your own data set can be found under [`experiments/tutorial.md`](experiments/tutorial.md).
## How to cite this code
Please cite the original publication [3].
## Installation
Setup package in virtual environment
```
git clone https://github.com/MIC-DKFZ/medicaldetectiontoolkit.git.
cd medicaldetectiontoolkit
virtualenv -p python3.7 mdt
source mdt/bin/activate
python setup.py install
```
##### Custom Extensions
This framework uses two custom mixed C++/CUDA extensions: Non-maximum suppression (NMS) and RoIAlign. Both are adapted from the original pytorch extensions (under torchvision.ops.boxes and ops.roialign).
The extensions are automatically compiled from the provided source files under medicaldetectiontoolkit/custom_extensions with above setup.py.
However, the extensions need to be compiled specifically for certain GPU architectures. Hence, please ensure that the architectures you need are included in your shell's
environment variable ```TORCH_CUDA_ARCH_LIST``` before compilation.
Example: You want to use the modules with the new TITAN RTX GPU, which has
Compute Capability 7.5 (Turing Architecture), but sometimes you also want to use it with a TITAN Xp (6.1, Pascal). Before installation you need to
```export TORCH_CUDA_ARCH_LIST="6.1;7.5"```. A link list of GPU model names to Compute Capability can be found here: https://developer.nvidia.com/cuda-gpus.
Note: If you'd like to import the raw extensions (not the wrapper modules), be sure to import torch first.
## Prepare the Data
This framework is meant for you to be able to train models on your own data sets.
Two example data loaders are provided in medicaldetectiontoolkit/experiments including thorough documentation to ensure a quick start for your own project. The way I load Data is to have a preprocessing script, which after preprocessing saves the Data of whatever data type into numpy arrays (this is just run once). During training / testing, the data loader then loads these numpy arrays dynamically. (Please note the Data Input side is meant to be customized by you according to your own needs and the provided Data loaders are merely examples: LIDC has a powerful Dataloader that handles 2D/3D inputs and is optimized for patch-based training and inference. Toy-Experiments have a lightweight Dataloader, only handling 2D without patching. The latter makes sense if you want to get familiar with the framework.).
## Execute
1. Set I/O paths, model and training specifics in the configs file: medicaldetectiontoolkit/experiments/your_experiment/configs.py
2. Train the model:
```
python exec.py --mode train --exp_source experiments/my_experiment --exp_dir path/to/experiment/directory
```
This copies snapshots of configs and model to the specified exp_dir, where all outputs will be saved. By default, the data is split into 60% training and 20% validation and 20% testing data to perform a 5-fold cross validation (can be changed to hold-out test set in configs) and all folds will be trained iteratively. In order to train a single fold, specify it using the folds arg:
```
python exec.py --folds 0 1 2 .... # specify any combination of folds [0-4]
```
3. Run inference:
```
python exec.py --mode test --exp_dir path/to/experiment/directory
```
This runs the prediction pipeline and saves all results to exp_dir.
## Models
This framework features all models explored in [3] (implemented in 2D + 3D): The proposed Retina U-Net, a simple but effective Architecture fusing state-of-the-art semantic segmentation with object detection,
also implementations of prevalent object detectors, such as Mask R-CNN, Faster R-CNN+ (Faster R-CNN w\ RoIAlign), Retina Net, U-Faster R-CNN+ (the two stage counterpart of Retina U-Net: Faster R-CNN with auxiliary semantic segmentation), DetU-Net (a U-Net like segmentation architecture with heuristics for object detection.)
## Training annotations
This framework features training with pixelwise and/or bounding box annotations. To overcome the issue of box coordinates in
data augmentation, we feed the annotation masks through data augmentation (create a pseudo mask, if only bounding box annotations provided) and draw the boxes afterwards.
The framework further handles two types of pixel-wise annotations:
1. A label map with individual ROIs identified by increasing label values, accompanied by a vector containing in each position the class target for the lesion with the corresponding label (for this mode set get_rois_from_seg_flag = False when calling ConvertSegToBoundingBoxCoordinates in your Data Loader).
2. A binary label map. There is only one foreground class and single lesions are not identified. All lesions have the same class target (foreground). In this case the Dataloader runs a Connected Component Labelling algorithm to create processable lesion - class target pairs on the fly (for this mode set get_rois_from_seg_flag = True when calling ConvertSegToBoundingBoxCoordinates in your Data Loader).
## Prediction pipeline
This framework provides an inference module, which automatically handles patching of inputs, and tiling, ensembling, and weighted consolidation of output predictions:
## Consolidation of predictions (Weighted Box Clustering)
Multiple predictions of the same image (from test time augmentations, tested epochs and overlapping patches), result in a high amount of boxes (or cubes), which need to be consolidated. In semantic segmentation, the final output would typically be obtained by averaging every pixel over all predictions. As described in [3], **weighted box clustering** (WBC) does this for box predictions:
## Visualization / Monitoring
By default, loss functions and performance metrics are monitored:
Histograms of matched output predictions for training/validation/testing are plotted per foreground class:
Input images + ground truth annotations + output predictions of a sampled validation abtch are plotted after each epoch (here 2D sampled slice with +-3 neighbouring context slices in channels):
Zoomed into the last two lines of the plot:
## License
This framework is published under the [Apache License Version 2.0](LICENSE).
diff --git a/default_configs.py b/default_configs.py
index 4d90b82..9032b14 100644
--- a/default_configs.py
+++ b/default_configs.py
@@ -1,143 +1,143 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Default Configurations script. Avoids changing configs of all experiments if general settings are to be changed."""
import os
class DefaultConfigs:
- def __init__(self, model, server_env=None, dim=2):
+ def __init__(self, model, server_env=False, dim=2):
self.server_env = server_env
#########################
# I/O #
#########################
self.model = model
self.dim = dim
# int [0 < dataset_size]. select n patients from dataset for prototyping.
self.select_prototype_subset = None
# some default paths.
self.backbone_path = 'models/backbone.py'
self.source_dir = os.path.dirname(os.path.realpath(__file__)) #current dir.
self.input_df_name = 'info_df.pickle'
self.model_path = 'models/{}.py'.format(self.model)
if server_env:
self.source_dir = '/home/jaegerp/code/mamma_code/medicaldetectiontoolkit'
#########################
# Data Loader #
#########################
#random seed for fold_generator and batch_generator.
self.seed = 0
#number of threads for multithreaded batch generation.
- self.n_workers = os.cpu_count() - 1
+ self.n_workers = 16 if server_env else os.cpu_count()-1
# if True, segmentation losses learn all categories, else only foreground vs. background.
self.class_specific_seg_flag = False
#########################
# Architecture #
#########################
self.weight_decay = 0.0
# what weight or layer types to exclude from weight decay. options: ["bias", "norm"].
self.exclude_from_wd = ("norm",)
# 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.custom_init = False
# 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
#########################
# Schedule #
#########################
# number of folds in cross validation.
self.n_cv_splits = 5
# number of probabilistic samples in validation.
self.n_probabilistic_samples = None
#########################
# 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.hold_out_test_set = False
# if hold_out_test_set provided, ensemble predictions over models of all trained cv-folds.
# implications for hold-out test sets: if True, evaluate folds separately on the test set, aggregate only the
# evaluations. if False, aggregate the raw predictions across all folds, then evaluate.
self.ensemble_folds = False
# 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
# evaluates average precision per image and averages over images. instead computing one ap over 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
# monitor any value from training.
self.n_monitoring_figures = 1
# dict to assign specific plot_values to monitor_figures > 0. {1: ['class_loss'], 2: ['kl_loss', 'kl_sigmas']}
self.assign_values_to_extra_figure = {}
# save predictions to csv file in experiment dir.
self.save_preds_to_csv = True
# select a maximum number of patient cases to test. number or "all" for all
self.max_test_patients = "all"
#########################
# MRCNN #
#########################
# if True, mask loss is not applied. used for data sets, where no pixel-wise annotations are provided.
self.frcnn_mode = 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
# for probabilistic detection
self.n_latent_dims = 0
diff --git a/evaluator.py b/evaluator.py
index 1210f3a..dae3983 100644
--- a/evaluator.py
+++ b/evaluator.py
@@ -1,540 +1,541 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import os, time
from multiprocessing import Pool
import numpy as np
import pandas as pd
import torch
from sklearn.metrics import roc_auc_score, average_precision_score
from sklearn.metrics import roc_curve, precision_recall_curve
+import utils.exp_utils as utils
import utils.model_utils as mutils
import plotting
class Evaluator():
def __init__(self, cf, logger, mode='test'):
"""
:param mode: either 'val_sampling', 'val_patient' or 'test'. handles prediction lists of different forms.
"""
self.cf = cf
self.logger = logger
self.mode = mode
self.plot_dir = self.cf.test_dir if self.mode == "test" else self.cf.plot_dir
if self.cf.plot_prediction_histograms:
self.hist_dir = os.path.join(self.plot_dir, 'histograms')
os.makedirs(self.hist_dir, exist_ok=True)
if self.cf.plot_stat_curves:
self.curves_dir = os.path.join(self.plot_dir, 'stat_curves')
os.makedirs(self.curves_dir, exist_ok=True)
def eval_losses(self, batch_res_dicts):
if hasattr(self.cf, "losses_to_monitor"):
loss_names = self.cf.losses_to_monitor
else:
loss_names = {name for b_res_dict in batch_res_dicts for name in b_res_dict if 'loss' in name}
self.epoch_losses = {l_name: torch.tensor([b_res_dict[l_name] for b_res_dict in batch_res_dicts if l_name
in b_res_dict.keys()]).mean().item() for l_name in loss_names}
def eval_boxes(self, batch_res_dicts, pid_list):
""" """
df_list_preds = []
df_list_labels = []
df_list_class_preds = []
df_list_pids = []
df_list_type = []
df_list_match_iou = []
if self.mode == 'train' or self.mode=='val_sampling':
# one pid per batch element
# batch_size > 1, with varying patients across batch:
# [[[results_0, ...], [pid_0, ...]], [[results_n, ...], [pid_n, ...]], ...]
# -> [results_0, results_1, ..]
batch_inst_boxes = [b_res_dict['boxes'] for b_res_dict in batch_res_dicts] # len: nr of batches in epoch
batch_inst_boxes = [[b_inst_boxes] for whole_batch_boxes in batch_inst_boxes for b_inst_boxes in
whole_batch_boxes]
else:
# patient processing, one element per batch = one patient.
# [[results_0, pid_0], [results_1, pid_1], ...] -> [results_0, results_1, ..]
batch_inst_boxes = [b_res_dict['boxes'] for b_res_dict in batch_res_dicts]
assert len(batch_inst_boxes) == len(pid_list)
for match_iou in self.cf.ap_match_ious:
self.logger.info('evaluating with match_iou: {}'.format(match_iou))
for cl in list(self.cf.class_dict.keys()):
for pix, pid in enumerate(pid_list):
len_df_list_before_patient = len(df_list_pids)
# input of each batch element is a list of boxes, where each box is a dictionary.
for bix, b_boxes_list in enumerate(batch_inst_boxes[pix]):
b_tar_boxes = np.array([box['box_coords'] for box in b_boxes_list if
(box['box_type'] == 'gt' and box['box_label'] == cl)])
b_cand_boxes = np.array([box['box_coords'] for box in b_boxes_list if
(box['box_type'] == 'det' and
box['box_pred_class_id'] == cl)])
b_cand_scores = np.array([box['box_score'] for box in b_boxes_list if
(box['box_type'] == 'det' and
box['box_pred_class_id'] == cl)])
# check if predictions and ground truth boxes exist and match them according to match_iou.
if not 0 in b_cand_boxes.shape and not 0 in b_tar_boxes.shape:
overlaps = mutils.compute_overlaps(b_cand_boxes, b_tar_boxes)
match_cand_ixs = np.argwhere(np.max(overlaps, 1) > match_iou)[:, 0]
non_match_cand_ixs = np.argwhere(np.max(overlaps, 1) <= match_iou)[:, 0]
match_gt_ixs = np.argmax(overlaps[match_cand_ixs, :],
1) if not 0 in match_cand_ixs.shape else np.array([])
non_match_gt_ixs = np.array(
[ii for ii in np.arange(b_tar_boxes.shape[0]) if ii not in match_gt_ixs])
unique, counts = np.unique(match_gt_ixs, return_counts=True)
# check for double assignments, i.e. two predictions having been assigned to the same gt.
# according to the COCO-metrics, only one prediction counts as true positive, the rest counts as
# false positive. This case is supposed to be avoided by the model itself by,
# e.g. using a low enough NMS threshold.
if np.any(counts > 1):
double_match_gt_ixs = unique[np.argwhere(counts > 1)[:, 0]]
keep_max = []
double_match_list = []
for dg in double_match_gt_ixs:
double_match_cand_ixs = match_cand_ixs[np.argwhere(match_gt_ixs == dg)]
keep_max.append(double_match_cand_ixs[np.argmax(b_cand_scores[double_match_cand_ixs])])
double_match_list += [ii for ii in double_match_cand_ixs]
fp_ixs = np.array([ii for ii in match_cand_ixs if
(ii in double_match_list and ii not in keep_max)])
match_cand_ixs = np.array([ii for ii in match_cand_ixs if ii not in fp_ixs])
df_list_preds += [ii for ii in b_cand_scores[fp_ixs]]
df_list_labels += [0] * fp_ixs.shape[0]
df_list_class_preds += [cl] * fp_ixs.shape[0]
df_list_pids += [pid] * fp_ixs.shape[0]
df_list_type += ['det_fp'] * fp_ixs.shape[0]
# matched:
if not 0 in match_cand_ixs.shape:
df_list_preds += [ii for ii in b_cand_scores[match_cand_ixs]]
df_list_labels += [1] * match_cand_ixs.shape[0]
df_list_class_preds += [cl] * match_cand_ixs.shape[0]
df_list_pids += [pid] * match_cand_ixs.shape[0]
df_list_type += ['det_tp'] * match_cand_ixs.shape[0]
# rest fp:
if not 0 in non_match_cand_ixs.shape:
df_list_preds += [ii for ii in b_cand_scores[non_match_cand_ixs]]
df_list_labels += [0] * non_match_cand_ixs.shape[0]
df_list_class_preds += [cl] * non_match_cand_ixs.shape[0]
df_list_pids += [pid] * non_match_cand_ixs.shape[0]
df_list_type += ['det_fp'] * non_match_cand_ixs.shape[0]
# rest fn:
if not 0 in non_match_gt_ixs.shape:
df_list_preds += [0] * non_match_gt_ixs.shape[0]
df_list_labels += [1] * non_match_gt_ixs.shape[0]
df_list_class_preds += [cl] * non_match_gt_ixs.shape[0]
df_list_pids += [pid] * non_match_gt_ixs.shape[0]
df_list_type += ['det_fn'] * non_match_gt_ixs.shape[0]
# only fp:
if not 0 in b_cand_boxes.shape and 0 in b_tar_boxes.shape:
df_list_preds += [ii for ii in b_cand_scores]
df_list_labels += [0] * b_cand_scores.shape[0]
df_list_class_preds += [cl] * b_cand_scores.shape[0]
df_list_pids += [pid] * b_cand_scores.shape[0]
df_list_type += ['det_fp'] * b_cand_scores.shape[0]
# only fn:
if 0 in b_cand_boxes.shape and not 0 in b_tar_boxes.shape:
df_list_preds += [0] * b_tar_boxes.shape[0]
df_list_labels += [1] * b_tar_boxes.shape[0]
df_list_class_preds += [cl] * b_tar_boxes.shape[0]
df_list_pids += [pid] * b_tar_boxes.shape[0]
df_list_type += ['det_fn'] * b_tar_boxes.shape[0]
# empty patient with 0 detections needs patient dummy score, in order to not disappear from stats.
# filtered out for roi-level evaluation later. During training (and val_sampling),
# tn are assigned per sample independently of associated patients.
if len(df_list_pids) == len_df_list_before_patient:
df_list_preds += [0] * 1
df_list_labels += [0] * 1
df_list_class_preds += [cl] * 1
df_list_pids += [pid] * 1
df_list_type += ['patient_tn'] * 1 # true negative: no ground truth boxes, no detections.
df_list_match_iou += [match_iou] * (len(df_list_preds) - len(df_list_match_iou))
self.test_df = pd.DataFrame()
self.test_df['pred_score'] = df_list_preds
self.test_df['class_label'] = df_list_labels
self.test_df['pred_class'] = df_list_class_preds
self.test_df['pid'] = df_list_pids
self.test_df['det_type'] = df_list_type
self.test_df['fold'] = self.cf.fold
self.test_df['match_iou'] = df_list_match_iou
def evaluate_predictions(self, results_list, monitor_metrics=None):
"""
Performs the matching of predicted boxes and ground truth boxes. Loops over list of matching IoUs and foreground classes.
Resulting info of each prediction is stored as one line in an internal dataframe, with the keys:
det_type: 'tp' (true positive), 'fp' (false positive), 'fn' (false negative), 'tn' (true negative)
pred_class: foreground class which the object predicts.
pid: corresponding patient-id.
pred_score: confidence score [0, 1]
fold: corresponding fold of CV.
match_iou: utilized IoU for matching.
:param results_list: list of model predictions. Either from train/val_sampling (patch processing) for monitoring with form:
[[[results_0, ...], [pid_0, ...]], [[results_n, ...], [pid_n, ...]], ...]
Or from val_patient/testing (patient processing), with form: [[results_0, pid_0], [results_1, pid_1], ...])
:param monitor_metrics (optional): dict of dicts with all metrics of previous epochs.
:return monitor_metrics: if provided (during training), return monitor_metrics now including results of current epoch.
"""
self.logger.info('evaluating in mode {}'.format(self.mode))
batch_res_dicts = [batch[0] for batch in results_list] # len: nr of batches in epoch
if self.mode == 'train' or self.mode == 'val_sampling':
# one pid per batch element
# [[[results_0, ...], [pid_0, ...]], [[results_n, ...], [pid_n, ...]], ...]
# -> [pid_0, pid_1, ...]
# additional list wrapping to make conform with below per-patient batches, where one pid is linked to more than one batch instance
pid_list = [batch_instance_pid for batch in results_list for batch_instance_pid in batch[1]]
elif self.mode == "val_patient" or self.mode == "test":
# [[results_0, pid_0], [results_1, pid_1], ...] -> [pid_0, pid_1, ...]
# in patientbatchiterator there is only one pid per batch
pid_list = [np.unique(batch[1]) for batch in results_list]
assert np.all([len(pid) == 1 for pid in
pid_list]), "pid list in patient-eval mode, should only contain a single scalar per patient: {}".format(
pid_list)
pid_list = [pid[0] for pid in pid_list]
# todo remove assert
pid_list_orig = [item[1] for item in results_list]
assert np.all(pid_list == pid_list_orig)
else:
raise Exception("undefined run mode encountered")
self.eval_losses(batch_res_dicts)
self.eval_boxes(batch_res_dicts, pid_list)
if monitor_metrics is not None:
# return all_stats, updated monitor_metrics
return self.return_metrics(monitor_metrics)
def return_metrics(self, monitor_metrics=None):
"""
calculates AP/AUC scores for internal dataframe. called directly from evaluate_predictions during training for monitoring,
or from score_test_df during inference (for single folds or aggregated test set). Loops over foreground classes
and score_levels (typically 'roi' and 'patient'), gets scores and stores them. Optionally creates plots of
prediction histograms and roc/prc curves.
:param monitor_metrics: dict of dicts with all metrics of previous epochs.
this function adds metrics for current epoch and returns the same object.
:return: all_stats: list. Contains dicts with resulting scores for each combination of foreground class and
score_level.
:return: monitor_metrics
"""
# -------------- monitoring independent of class, score level ------------
if monitor_metrics is not None:
for l_name in self.epoch_losses:
monitor_metrics[l_name] = [self.epoch_losses[l_name]]
df = self.test_df
all_stats = []
for cl in list(self.cf.class_dict.keys()):
cl_df = df[df.pred_class == cl]
for score_level in self.cf.report_score_level:
stats_dict = {}
stats_dict['name'] = 'fold_{} {} cl_{}'.format(self.cf.fold, score_level, cl)
if score_level == 'rois':
# kick out dummy entries for true negative patients. not needed on roi-level.
spec_df = cl_df[cl_df.det_type != 'patient_tn']
stats_dict['ap'] = get_roi_ap_from_df([spec_df, self.cf.min_det_thresh, self.cf.per_patient_ap])
# AUC not sensible on roi-level, since true negative box predictions do not exist. Would reward
# higher amounts of low confidence false positives.
stats_dict['auc'] = np.nan
stats_dict['roc'] = np.nan
stats_dict['prc'] = np.nan
# for the aggregated test set case, additionally get the scores for averaging over fold results.
if len(df.fold.unique()) > 1:
aps = []
for fold in df.fold.unique():
fold_df = spec_df[spec_df.fold == fold]
aps.append(get_roi_ap_from_df([fold_df, self.cf.min_det_thresh, self.cf.per_patient_ap]))
stats_dict['mean_ap'] = np.mean(aps)
stats_dict['mean_auc'] = 0
# on patient level, aggregate predictions per patient (pid): The patient predicted score is the highest
# confidence prediction for this class. The patient class label is 1 if roi of this class exists in patient, else 0.
if score_level == 'patient':
spec_df = cl_df.groupby(['pid'], as_index=False).agg({'class_label': 'max', 'pred_score': 'max', 'fold': 'first'})
if len(spec_df.class_label.unique()) > 1:
stats_dict['auc'] = roc_auc_score(spec_df.class_label.tolist(), spec_df.pred_score.tolist())
stats_dict['roc'] = roc_curve(spec_df.class_label.tolist(), spec_df.pred_score.tolist())
else:
stats_dict['auc'] = np.nan
stats_dict['roc'] = np.nan
if (spec_df.class_label == 1).any():
stats_dict['ap'] = average_precision_score(spec_df.class_label.tolist(), spec_df.pred_score.tolist())
stats_dict['prc'] = precision_recall_curve(spec_df.class_label.tolist(), spec_df.pred_score.tolist())
else:
stats_dict['ap'] = np.nan
stats_dict['prc'] = np.nan
# for the aggregated test set case, additionally get the scores for averaging over fold results.
if len(df.fold.unique()) > 1:
aucs = []
aps = []
for fold in df.fold.unique():
fold_df = spec_df[spec_df.fold == fold]
if len(fold_df.class_label.unique()) > 1:
aucs.append(roc_auc_score(fold_df.class_label.tolist(), fold_df.pred_score.tolist()))
if (fold_df.class_label == 1).any():
aps.append(average_precision_score(fold_df.class_label.tolist(), fold_df.pred_score.tolist()))
stats_dict['mean_auc'] = np.mean(aucs)
stats_dict['mean_ap'] = np.mean(aps)
# fill new results into monitor_metrics dict. for simplicity, only one class (of interest) is monitored on patient level.
if monitor_metrics is not None and not (score_level == 'patient' and cl != self.cf.patient_class_of_interest):
score_level_name = 'patient' if score_level == 'patient' else self.cf.class_dict[cl]
monitor_metrics[score_level_name + '_ap'].append(stats_dict['ap'] if stats_dict['ap'] > 0 else np.nan)
if score_level == 'patient':
monitor_metrics[score_level_name + '_auc'].append(
stats_dict['auc'] if stats_dict['auc'] > 0 else np.nan)
if self.cf.plot_prediction_histograms:
out_filename = os.path.join(self.hist_dir, 'pred_hist_{}_{}_{}_cl{}'.format(
self.cf.fold, 'val' if 'val' in self.mode else self.mode, score_level, cl))
type_list = None if score_level == 'patient' else spec_df.det_type.tolist()
- plotting.plot_prediction_hist(spec_df.class_label.tolist(), spec_df.pred_score.tolist(), type_list, out_filename)
+ utils.split_off_process(plotting.plot_prediction_hist, spec_df.class_label.tolist(),
+ spec_df.pred_score.tolist(), type_list, out_filename)
all_stats.append(stats_dict)
# analysis of the hyper-parameter cf.min_det_thresh, for optimization on validation set.
if self.cf.scan_det_thresh:
conf_threshs = list(np.arange(0.9, 1, 0.01))
pool = Pool(processes=10)
mp_inputs = [[spec_df, ii, self.cf.per_patient_ap] for ii in conf_threshs]
aps = pool.map(get_roi_ap_from_df, mp_inputs, chunksize=1)
pool.close()
pool.join()
self.logger.info('results from scanning over det_threshs:', [[i, j] for i, j in zip(conf_threshs, aps)])
if self.cf.plot_stat_curves:
out_filename = os.path.join(self.curves_dir, '{}_{}_stat_curves'.format(self.cf.fold, self.mode))
- plotting.plot_stat_curves(all_stats, out_filename)
-
+ utils.split_off_process(plotting.plot_stat_curves, all_stats, out_filename)
# get average stats over foreground classes on roi level.
avg_ap = np.mean([d['ap'] for d in all_stats if 'rois' in d['name']])
all_stats.append({'name': 'average_foreground_roi', 'auc': 0, 'ap': avg_ap})
if len(df.fold.unique()) > 1:
avg_mean_ap = np.mean([d['mean_ap'] for d in all_stats if 'rois' in d['name']])
all_stats[-1]['mean_ap'] = avg_mean_ap
all_stats[-1]['mean_auc'] = 0
# in small data sets, values of model_selection_criterion can be identical across epochs, wich breaks the
# ranking of model_selector. Thus, pertube identical values by a neglectibale random term.
for sc in self.cf.model_selection_criteria:
if 'val' in self.mode and monitor_metrics[sc].count(monitor_metrics[sc][-1]) > 1 and monitor_metrics[sc][-1] is not None:
monitor_metrics[sc][-1] += 1e-6 * np.random.rand()
return all_stats, monitor_metrics
def write_to_results_table(self, stats, metrics_to_score, out_path):
"""Write overall results to a common inter-experiment table.
:param metrics_to_score:
:return:
"""
with open(out_path, 'a') as handle:
# ---column headers---
handle.write('\n{},'.format("Experiment Name"))
handle.write('{},'.format("Time Stamp"))
handle.write('{},'.format("Samples Seen"))
handle.write('{},'.format("Spatial Dim"))
handle.write('{},'.format("Patch Size"))
handle.write('{},'.format("CV Folds"))
handle.write('{},'.format("WBC IoU"))
handle.write('{},'.format("Merge-2D-to-3D IoU"))
for s in stats:
#if self.cf.class_dict[self.cf.patient_class_of_interest] in s['name'] or "average" in s["name"]:
for metric in metrics_to_score:
if metric in s.keys() and not np.isnan(s[metric]):
if metric == 'ap':
handle.write('{} : {}_{},'.format(s['name'], metric.upper(),
"_".join((np.array(self.cf.ap_match_ious) * 100)
.astype("int").astype("str"))))
else:
handle.write('{} : {},'.format(s['name'], metric.upper()))
else:
print("WARNING: skipped metric {} since not avail".format(metric))
handle.write('\n')
# --- columns content---
handle.write('{},'.format(self.cf.exp_dir.split(os.sep)[-1]))
handle.write('{},'.format(time.strftime("%d%b%y %H:%M:%S")))
handle.write('{},'.format(self.cf.num_epochs * self.cf.num_train_batches * self.cf.batch_size))
handle.write('{}D,'.format(self.cf.dim))
handle.write('{},'.format("x".join([str(self.cf.patch_size[i]) for i in range(self.cf.dim)])))
handle.write('{},'.format(str(self.test_df.fold.unique().tolist()).replace(",", "")))
handle.write('{},'.format(self.cf.wcs_iou))
handle.write('{},'.format(self.cf.merge_3D_iou if self.cf.merge_2D_to_3D_preds else str("N/A")))
for s in stats:
#if self.cf.class_dict[self.cf.patient_class_of_interest] in s['name'] or "mean" in s["name"]:
for metric in metrics_to_score:
if metric in s.keys() and not np.isnan(s[metric]):
handle.write('{:0.3f}, '.format(s[metric]))
handle.write('\n')
def score_test_df(self, internal_df=True):
"""
Writes out resulting scores to text files: First checks for class-internal-df (typically current) fold,
gets resulting scores, writes them to a text file and pickles data frame. Also checks if data-frame pickles of
all folds of cross-validation exist in exp_dir. If true, loads all dataframes, aggregates test sets over folds,
and calculates and writes out overall metrics.
"""
if internal_df:
self.test_df.to_pickle(os.path.join(self.cf.test_dir, '{}_test_df.pickle'.format(self.cf.fold)))
stats, _ = self.return_metrics()
with open(os.path.join(self.cf.test_dir, 'results.txt'), 'a') as handle:
handle.write('\n****************************\n')
handle.write('\nresults for fold {} \n'.format(self.cf.fold))
handle.write('\n****************************\n')
handle.write('\nfold df shape {}\n \n'.format(self.test_df.shape))
for s in stats:
handle.write('AUC {:0.4f} AP {:0.4f} {} \n'.format(s['auc'], s['ap'], s['name']))
fold_df_paths = [ii for ii in os.listdir(self.cf.test_dir) if ('test_df.pickle' in ii and not 'overall' in ii)]
- if len(fold_df_paths) == self.cf.n_cv_splits and self.cf.fold == self.cf.n_cv_splits - 1:
+ if len(fold_df_paths) == self.cf.n_cv_splits:
results_table_path = os.path.join((os.sep).join(self.cf.exp_dir.split(os.sep)[:-1]), 'results_table.csv')
if not self.cf.hold_out_test_set or not self.cf.ensemble_folds:
with open(os.path.join(self.cf.test_dir, 'results.txt'), 'a') as handle:
self.cf.fold = 'overall'
dfs_list = [pd.read_pickle(os.path.join(self.cf.test_dir, ii)) for ii in fold_df_paths]
for ix, df in enumerate(dfs_list):
df['fold'] = ix
self.test_df = pd.concat(dfs_list)
stats, _ = self.return_metrics()
handle.write('\n****************************\n')
handle.write('\nOVERALL RESULTS \n')
handle.write('\n****************************\n')
handle.write('\ndf shape \n \n'.format(self.test_df.shape))
for s in stats:
handle.write('\nAUC {:0.4f} (mu {:0.4f}) AP {:0.4f} (mu {:0.4f}) {}\n '
.format(s['auc'], s['mean_auc'], s['ap'], s['mean_ap'], s['name']))
metrics_to_score = ["auc", "mean_auc", "ap", "mean_ap"]
self.write_to_results_table(stats, metrics_to_score, out_path=results_table_path)
else:
metrics_to_score = ["auc", "ap"]
self.write_to_results_table(stats, metrics_to_score, out_path=results_table_path)
def get_roi_ap_from_df(inputs):
'''
:param df: data frame.
:param det_thresh: min_threshold for filtering out low confidence predictions.
:param per_patient_ap: boolean flag. evaluate average precision per image and average over images,
instead of computing one ap over data set.
:return: average_precision (float)
'''
df, det_thresh, per_patient_ap = inputs
if per_patient_ap:
pids_list = df.pid.unique()
aps = []
for match_iou in df.match_iou.unique():
iou_df = df[df.match_iou == match_iou]
for pid in pids_list:
pid_df = iou_df[iou_df.pid == pid]
all_p = len(pid_df[pid_df.class_label == 1])
pid_df = pid_df[(pid_df.det_type == 'det_fp') | (pid_df.det_type == 'det_tp')].sort_values('pred_score', ascending=False)
pid_df = pid_df[pid_df.pred_score > det_thresh]
if (len(pid_df) ==0 and all_p == 0):
pass
elif (len(pid_df) > 0 and all_p == 0):
aps.append(0)
else:
aps.append(compute_roi_ap(pid_df, all_p))
return np.mean(aps)
else:
aps = []
for match_iou in df.match_iou.unique():
iou_df = df[df.match_iou == match_iou]
all_p = len(iou_df[iou_df.class_label == 1])
iou_df = iou_df[(iou_df.det_type == 'det_fp') | (iou_df.det_type == 'det_tp')].sort_values('pred_score', ascending=False)
iou_df = iou_df[iou_df.pred_score > det_thresh]
if all_p > 0:
aps.append(compute_roi_ap(iou_df, all_p))
return np.mean(aps)
def compute_roi_ap(df, all_p):
"""
adapted from: https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocotools/cocoeval.py
:param df: dataframe containing class labels of predictions sorted in descending manner by their prediction score.
:param all_p: number of all ground truth objects. (for denominator of recall.)
:return:
"""
tp = df.class_label.values
fp = (tp == 0) * 1
#recall thresholds, where precision will be measured
R = np.linspace(.0, 1, 101, endpoint=True)
tp_sum = np.cumsum(tp)
fp_sum = np.cumsum(fp)
nd = len(tp)
rc = tp_sum / all_p
pr = tp_sum / (fp_sum + tp_sum)
# initialize precision array over recall steps.
q = np.zeros((len(R),))
# numpy is slow without cython optimization for accessing elements
# use python array gets significant speed improvement
pr = pr.tolist()
q = q.tolist()
for i in range(nd - 1, 0, -1):
if pr[i] > pr[i - 1]:
pr[i - 1] = pr[i]
#discretize empiric recall steps with given bins.
inds = np.searchsorted(rc, R, side='left')
try:
for ri, pi in enumerate(inds):
q[ri] = pr[pi]
except:
pass
return np.mean(q)
\ No newline at end of file
diff --git a/exec.py b/exec.py
index 7ce003a..8131603 100644
--- a/exec.py
+++ b/exec.py
@@ -1,294 +1,303 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""execution script."""
import argparse
import os, warnings
import time
import torch
import utils.exp_utils as utils
from evaluator import Evaluator
from predictor import Predictor
from plotting import plot_batch_prediction
for msg in ["Attempting to set identical bottom==top results",
"This figure includes Axes that are not compatible with tight_layout",
"Data has no positive values, and therefore cannot be log-scaled.",
".*invalid value encountered in double_scalars.*",
".*Mean of empty slice.*"]:
warnings.filterwarnings("ignore", msg)
def train(logger):
"""
perform the training routine for a given fold. saves plots and selected parameters to the experiment dir
specified in the configs.
"""
logger.info('performing training in {}D over fold {} on experiment {} with model {}'.format(
cf.dim, cf.fold, cf.exp_dir, cf.model))
net = model.net(cf, logger).cuda()
- optimizer = torch.optim.AdamW(utils.parse_params_for_optim(net, weight_decay=cf.weight_decay,
- exclude_from_wd=cf.exclude_from_wd),
- lr=cf.learning_rate[0])
+ if hasattr(cf, "optimizer") and cf.optimizer.lower() == "adam":
+ logger.info("Using Adam optimizer.")
+ optimizer = torch.optim.Adam(utils.parse_params_for_optim(net, weight_decay=cf.weight_decay,
+ exclude_from_wd=cf.exclude_from_wd),
+ lr=cf.learning_rate[0])
+ else:
+ logger.info("Using AdamW optimizer.")
+ optimizer = torch.optim.AdamW(utils.parse_params_for_optim(net, weight_decay=cf.weight_decay,
+ exclude_from_wd=cf.exclude_from_wd),
+ lr=cf.learning_rate[0])
+
+
if cf.dynamic_lr_scheduling:
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode=cf.scheduling_mode, factor=cf.lr_decay_factor,
patience=cf.scheduling_patience)
model_selector = utils.ModelSelector(cf, logger)
train_evaluator = Evaluator(cf, logger, mode='train')
val_evaluator = Evaluator(cf, logger, mode=cf.val_mode)
starting_epoch = 1
# prepare monitoring
monitor_metrics = utils.prepare_monitoring(cf)
if cf.resume:
checkpoint_path = os.path.join(cf.fold_dir, "last_checkpoint")
starting_epoch, net, optimizer, monitor_metrics = \
utils.load_checkpoint(checkpoint_path, net, optimizer)
logger.info('resumed from checkpoint {} to epoch {}'.format(checkpoint_path, starting_epoch))
logger.info('loading dataset and initializing batch generators...')
batch_gen = data_loader.get_train_generators(cf, logger)
for epoch in range(starting_epoch, cf.num_epochs + 1):
logger.info('starting training epoch {}'.format(epoch))
start_time = time.time()
net.train()
train_results_list = []
for bix in range(cf.num_train_batches):
batch = next(batch_gen['train'])
tic_fw = time.time()
results_dict = net.train_forward(batch)
tic_bw = time.time()
optimizer.zero_grad()
results_dict['torch_loss'].backward()
optimizer.step()
print('\rtr. batch {0}/{1} (ep. {2}) fw {3:.2f}s / bw {4:.2f} s / total {5:.2f} s || '.format(
bix + 1, cf.num_train_batches, epoch, tic_bw - tic_fw, time.time() - tic_bw,
time.time() - tic_fw) + results_dict['logger_string'], flush=True, end="")
train_results_list.append(({k:v for k,v in results_dict.items() if k != "seg_preds"}, batch["pid"]))
print()
_, monitor_metrics['train'] = train_evaluator.evaluate_predictions(train_results_list, monitor_metrics['train'])
logger.info('generating training example plot.')
utils.split_off_process(plot_batch_prediction, batch, results_dict, cf, outfile=os.path.join(
cf.plot_dir, 'pred_example_{}_train.png'.format(cf.fold)))
train_time = time.time() - start_time
logger.info('starting validation in mode {}.'.format(cf.val_mode))
with torch.no_grad():
net.eval()
if cf.do_validation:
val_results_list = []
val_predictor = Predictor(cf, net, logger, mode='val')
for _ in range(batch_gen['n_val']):
batch = next(batch_gen[cf.val_mode])
if cf.val_mode == 'val_patient':
results_dict = val_predictor.predict_patient(batch)
elif cf.val_mode == 'val_sampling':
results_dict = net.train_forward(batch, is_validation=True)
#val_results_list.append([results_dict['boxes'], batch['pid']])
val_results_list.append(({k:v for k,v in results_dict.items() if k != "seg_preds"}, batch["pid"]))
_, monitor_metrics['val'] = val_evaluator.evaluate_predictions(val_results_list, monitor_metrics['val'])
model_selector.run_model_selection(net, optimizer, monitor_metrics, epoch)
# update monitoring and prediction plots
monitor_metrics.update({"lr":
{str(g): group['lr'] for (g, group) in enumerate(optimizer.param_groups)}})
logger.metrics2tboard(monitor_metrics, global_step=epoch)
epoch_time = time.time() - start_time
- logger.info('trained epoch {}: took {:.2f} s ({:.2f} s train / {:.2f} s val)'.format(
- epoch, epoch_time, train_time, epoch_time-train_time))
+ logger.info('trained epoch {}: took {} ({} train / {} val)'.format(
+ epoch, utils.get_formatted_duration(epoch_time, "ms"), utils.get_formatted_duration(train_time, "ms"),
+ utils.get_formatted_duration(epoch_time-train_time, "ms")))
batch = next(batch_gen['val_sampling'])
results_dict = net.train_forward(batch, is_validation=True)
logger.info('generating validation-sampling example plot.')
utils.split_off_process(plot_batch_prediction, batch, results_dict, cf, outfile=os.path.join(
cf.plot_dir, 'pred_example_{}_val.png'.format(cf.fold)))
# -------------- scheduling -----------------
if cf.dynamic_lr_scheduling:
scheduler.step(monitor_metrics["val"][cf.scheduling_criterion][-1])
else:
for param_group in optimizer.param_groups:
param_group['lr'] = cf.learning_rate[epoch-1]
def test(logger):
"""
perform testing for a given fold (or hold out set). save stats in evaluator.
"""
logger.info('starting testing model of fold {} in exp {}'.format(cf.fold, cf.exp_dir))
net = model.net(cf, logger).cuda()
test_predictor = Predictor(cf, net, logger, mode='test')
test_evaluator = Evaluator(cf, logger, mode='test')
batch_gen = data_loader.get_test_generator(cf, logger)
test_results_list = test_predictor.predict_test_set(batch_gen, return_results=True)
test_evaluator.evaluate_predictions(test_results_list)
test_evaluator.score_test_df()
if __name__ == '__main__':
stime = time.time()
parser = argparse.ArgumentParser()
parser.add_argument('-m', '--mode', type=str, default='train_test',
help='one out of: train / test / train_test / analysis / create_exp')
parser.add_argument('-f','--folds', nargs='+', type=int, default=None,
help='None runs over all folds in CV. otherwise specify list of folds.')
parser.add_argument('--exp_dir', type=str, default='/path/to/experiment/directory',
help='path to experiment dir. will be created if non existent.')
parser.add_argument('--server_env', default=False, action='store_true',
help='change IO settings to deploy models on a cluster.')
parser.add_argument('--data_dest', type=str, default=None, help="path to final data folder if different from config.")
parser.add_argument('--use_stored_settings', default=False, action='store_true',
help='load configs from existing exp_dir instead of source dir. always done for testing, '
'but can be set to true to do the same for training. useful in job scheduler environment, '
'where source code might change before the job actually runs.')
parser.add_argument('--resume', action="store_true", default=False,
help='if given, resume from checkpoint(s) of the specified folds.')
parser.add_argument('--exp_source', type=str, default='experiments/toy_exp',
help='specifies, from which source experiment to load configs and data_loader.')
parser.add_argument('--no_benchmark', action='store_true', help="Do not use cudnn.benchmark.")
parser.add_argument('--cuda_device', type=int, default=0, help="Index of CUDA device to use.")
parser.add_argument('-d', '--dev', default=False, action='store_true', help="development mode: shorten everything")
args = parser.parse_args()
folds = args.folds
torch.backends.cudnn.benchmark = not args.no_benchmark
if args.mode == 'train' or args.mode == 'train_test':
cf = utils.prep_exp(args.exp_source, args.exp_dir, args.server_env, args.use_stored_settings)
if args.dev:
folds = [0,1]
cf.batch_size, cf.num_epochs, cf.min_save_thresh, cf.save_n_models = 3 if cf.dim==2 else 1, 1, 0, 2
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 = 2
cf.data_dest = args.data_dest
logger = utils.get_logger(cf.exp_dir, cf.server_env)
logger.info("cudnn benchmark: {}, deterministic: {}.".format(torch.backends.cudnn.benchmark,
torch.backends.cudnn.deterministic))
logger.info("sending tensors to CUDA device: {}.".format(torch.cuda.get_device_name(args.cuda_device)))
data_loader = utils.import_module('dl', os.path.join(args.exp_source, 'data_loader.py'))
model = utils.import_module('model', cf.model_path)
logger.info("loaded model from {}".format(cf.model_path))
if folds is None:
folds = range(cf.n_cv_splits)
with torch.cuda.device(args.cuda_device):
for fold in folds:
cf.fold_dir = os.path.join(cf.exp_dir, 'fold_{}'.format(fold))
cf.fold = fold
cf.resume = args.resume
if not os.path.exists(cf.fold_dir):
os.mkdir(cf.fold_dir)
logger.set_logfile(fold=fold)
train(logger)
cf.resume = False
if args.mode == 'train_test':
test(logger)
elif args.mode == 'test':
cf = utils.prep_exp(args.exp_source, args.exp_dir, args.server_env, is_training=False, use_stored_settings=True)
if args.dev:
folds = [0,1]
cf.test_n_epochs = 2; cf.max_test_patients = 2
cf.data_dest = args.data_dest
logger = utils.get_logger(cf.exp_dir, cf.server_env)
data_loader = utils.import_module('dl', os.path.join(args.exp_source, 'data_loader.py'))
model = utils.import_module('model', cf.model_path)
logger.info("loaded model from {}".format(cf.model_path))
if folds is None:
folds = range(cf.n_cv_splits)
with torch.cuda.device(args.cuda_device):
for fold in folds:
cf.fold_dir = os.path.join(cf.exp_dir, 'fold_{}'.format(fold))
cf.fold = fold
logger.set_logfile(fold=fold)
test(logger)
# load raw predictions saved by predictor during testing, run aggregation algorithms and evaluation.
elif args.mode == 'analysis':
cf = utils.prep_exp(args.exp_source, args.exp_dir, args.server_env, is_training=False, use_stored_settings=True)
logger = utils.get_logger(cf.exp_dir, cf.server_env)
if args.dev:
cf.test_n_epochs = 2
if cf.hold_out_test_set and cf.ensemble_folds:
# create and save (unevaluated) predictions across all folds
predictor = Predictor(cf, net=None, logger=logger, mode='analysis')
results_list = predictor.load_saved_predictions(apply_wbc=True)
utils.create_csv_output([(res_dict["boxes"], pid) for res_dict, pid in results_list], cf, logger)
logger.info('starting evaluation...')
cf.fold = 'overall_hold_out'
evaluator = Evaluator(cf, logger, mode='test')
evaluator.evaluate_predictions(results_list)
evaluator.score_test_df()
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 folds is None:
folds = range(cf.n_cv_splits)
for fold in folds:
cf.fold_dir = os.path.join(cf.exp_dir, 'fold_{}'.format(fold))
cf.fold = fold
logger.set_logfile(fold=fold)
if cf.fold_dir in fold_dirs:
predictor = Predictor(cf, net=None, logger=logger, mode='analysis')
results_list = predictor.load_saved_predictions(apply_wbc=True)
logger.info('starting evaluation...')
evaluator = Evaluator(cf, logger, mode='test')
evaluator.evaluate_predictions(results_list)
evaluator.score_test_df()
else:
logger.info("Skipping fold {} since no model parameters found.".format(fold))
# create experiment folder and copy scripts without starting job.
# useful for cloud deployment where configs might change before job actually runs.
elif args.mode == 'create_exp':
cf = utils.prep_exp(args.exp_source, args.exp_dir, args.server_env, use_stored_settings=False)
logger = utils.get_logger(cf.exp_dir)
logger.info('created experiment directory at {}'.format(cf.exp_dir))
else:
raise RuntimeError('mode specified in args is not implemented...')
- mins, secs = divmod((time.time() - stime), 60)
- h, mins = divmod(mins, 60)
- t = "{:d}h:{:02d}m:{:02d}s".format(int(h), int(mins), int(secs))
+
+ t = utils.get_formatted_duration(time.time() - stime)
logger.info("{} total runtime: {}".format(os.path.split(__file__)[1], t))
del logger
\ No newline at end of file
diff --git a/experiments/lidc_exp/configs.py b/experiments/lidc_exp/configs.py
index 6d791f8..c527d79 100644
--- a/experiments/lidc_exp/configs.py
+++ b/experiments/lidc_exp/configs.py
@@ -1,341 +1,344 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import sys
import os
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
import numpy as np
from default_configs import DefaultConfigs
class configs(DefaultConfigs):
- def __init__(self, server_env=None):
+ def __init__(self, server_env=False):
#########################
# Preprocessing #
#########################
self.root_dir = '/home/gregor/networkdrives/E130-Personal/Goetz/Datenkollektive/Lungendaten/Nodules_LIDC_IDRI'
self.raw_data_dir = '{}/new_nrrd'.format(self.root_dir)
self.pp_dir = '/media/gregor/HDD2TB/data/lidc/lidc_mdt'
self.target_spacing = (0.7, 0.7, 1.25)
#########################
# I/O #
#########################
# one out of [2, 3]. dimension the model operates in.
- self.dim = 2
+ self.dim = 3
# one out of ['mrcnn', 'retina_net', 'retina_unet', 'detection_unet', 'ufrcnn'].
- self.model = 'mrcnn'
+ self.model = 'retina_unet'
DefaultConfigs.__init__(self, self.model, server_env, self.dim)
# int [0 < dataset_size]. select n patients from dataset for prototyping. If None, all data is used.
self.select_prototype_subset = None
# path to preprocessed data.
self.pp_name = 'lidc_mdt'
self.input_df_name = 'info_df.pickle'
self.pp_data_path = '/media/gregor/HDD2TB/data/lidc/{}'.format(self.pp_name)
self.pp_test_data_path = self.pp_data_path #change if test_data in separate folder.
# settings for deployment in cloud.
if server_env:
# path to preprocessed data.
self.pp_name = 'lidc_mdt_npz'
self.crop_name = 'pp_fg_slices_packed'
self.pp_data_path = '/datasets/datasets_ramien/lidc_exp/data/{}'.format(self.pp_name)
self.pp_test_data_path = self.pp_data_path
self.select_prototype_subset = None
#########################
# Data Loader #
#########################
# select modalities from preprocessed data
self.channels = [0]
self.n_channels = len(self.channels)
# patch_size to be used for training. pre_crop_size is the patch_size before data augmentation.
self.pre_crop_size_2D = [300, 300]
self.patch_size_2D = [288, 288]
self.pre_crop_size_3D = [156, 156, 96]
self.patch_size_3D = [128, 128, 64]
self.patch_size = self.patch_size_2D if self.dim == 2 else self.patch_size_3D
self.pre_crop_size = self.pre_crop_size_2D if self.dim == 2 else self.pre_crop_size_3D
# ratio of free sampled batch elements before class balancing is triggered
# (>0 to include "empty"/background patches.)
self.batch_sample_slack = 0.2
# set 2D network to operate in 3D images.
self.merge_2D_to_3D_preds = self.dim == 2
# feed +/- n neighbouring slices into channel dimension. set to None for no context.
self.n_3D_context = None
if self.n_3D_context is not None and self.dim == 2:
self.n_channels *= (self.n_3D_context * 2 + 1)
#########################
# Architecture #
#########################
self.start_filts = 48 if self.dim == 2 else 18
self.end_filts = self.start_filts * 4 if self.dim == 2 else self.start_filts * 2
self.res_architecture = 'resnet50' # 'resnet101' , 'resnet50'
self.norm = None # one of None, 'instance_norm', 'batch_norm'
- self.weight_decay = 1e-8
+ # 0 for no weight decay
+ self.weight_decay = 0
# one of 'xavier_uniform', 'xavier_normal', or 'kaiming_normal', None (=default = 'kaiming_uniform')
self.weight_init = None
#########################
# Schedule / Selection #
#########################
- self.num_epochs = 80
- self.num_train_batches = 200 if self.dim == 2 else 300
+ self.num_epochs = 100
+ self.num_train_batches = 200 if self.dim == 2 else 200
self.batch_size = 20 if self.dim == 2 else 8
self.do_validation = True
# decide whether to validate on entire patient volumes (like testing) or sampled patches (like training)
# the former is more accurate, while the latter is faster (depending on volume size)
self.val_mode = 'val_sampling' # one of 'val_sampling' , 'val_patient'
if self.val_mode == 'val_patient':
self.max_val_patients = 50 # if 'None' iterates over entire val_set once.
if self.val_mode == 'val_sampling':
self.num_val_batches = 50
+ self.optimizer = "Adam"
+
# set dynamic_lr_scheduling to True to apply LR scheduling with below settings.
self.dynamic_lr_scheduling = False
self.lr_decay_factor = 0.25
self.scheduling_patience = np.ceil(16000 / (self.num_train_batches * self.batch_size))
self.scheduling_criterion = 'malignant_ap'
self.scheduling_mode = 'min' if "loss" in self.scheduling_criterion else 'max'
#########################
# Testing / Plotting #
#########################
# set the top-n-epochs to be saved for temporal averaging in testing.
- self.save_n_models = 4
- self.test_n_epochs = 4
+ self.save_n_models = 5
+ self.test_n_epochs = 5
# set a minimum epoch number for saving in case of instabilities in the first phase of training.
- self.min_save_thresh = 1 if self.dim == 2 else 1
+ self.min_save_thresh = 0 if self.dim == 2 else 0
self.report_score_level = ['patient', 'rois'] # choose list from 'patient', 'rois'
self.class_dict = {1: 'benign', 2: 'malignant'} # 0 is background.
self.patient_class_of_interest = 2 # patient metrics are only plotted for one class.
self.ap_match_ious = [0.1] # list of ious to be evaluated for ap-scoring.
self.model_selection_criteria = ['malignant_ap', 'benign_ap'] # criteria to average over for saving epochs.
self.min_det_thresh = 0.1 # minimum confidence value to select predictions for evaluation.
# threshold for clustering predictions together (wcs = weighted cluster scoring).
# needs to be >= the expected overlap of predictions coming from one model (typically NMS threshold).
# if too high, preds of the same object are separate clusters.
self.wcs_iou = 1e-5
self.plot_prediction_histograms = True
self.plot_stat_curves = False
#########################
# Data Augmentation #
#########################
self.da_kwargs={
'do_elastic_deform': True,
'alpha':(0., 1500.),
'sigma':(30., 50.),
'do_rotation':True,
'angle_x': (0., 2 * np.pi),
'angle_y': (0., 0),
'angle_z': (0., 0),
'do_scale': True,
'scale':(0.8, 1.1),
'random_crop':False,
'rand_crop_dist': (self.patch_size[0] / 2. - 3, self.patch_size[1] / 2. - 3),
'border_mode_data': 'constant',
'border_cval_data': 0,
'order_data': 1
}
if self.dim == 3:
self.da_kwargs['do_elastic_deform'] = False
self.da_kwargs['angle_x'] = (0, 0.0)
self.da_kwargs['angle_y'] = (0, 0.0) #must be 0!!
self.da_kwargs['angle_z'] = (0., 2 * np.pi)
#########################
# Add model specifics #
#########################
{'detection_unet': self.add_det_unet_configs,
'mrcnn': self.add_mrcnn_configs,
'ufrcnn': self.add_mrcnn_configs,
'retina_net': self.add_mrcnn_configs,
'retina_unet': self.add_mrcnn_configs,
}[self.model]()
def add_det_unet_configs(self):
self.learning_rate = [1e-4] * self.num_epochs
# aggregation from pixel perdiction to object scores (connected component). One of ['max', 'median']
self.aggregation_operation = 'max'
# max number of roi candidates to identify per batch element and class.
self.n_roi_candidates = 10 if self.dim == 2 else 30
# loss mode: either weighted cross entropy ('wce'), batch-wise dice loss ('dice), or the sum of both ('dice_wce')
self.seg_loss_mode = 'dice_wce'
# if <1, false positive predictions in foreground are penalized less.
self.fp_dice_weight = 1 if self.dim == 2 else 1
self.wce_weights = [0.3, 1, 1]
self.detection_min_confidence = self.min_det_thresh
# if 'True', loss distinguishes all classes, else only foreground vs. background (class agnostic).
self.class_specific_seg_flag = True
self.num_seg_classes = 3 if self.class_specific_seg_flag else 2
self.head_classes = self.num_seg_classes
def add_mrcnn_configs(self):
# learning rate is a list with one entry per epoch.
- self.learning_rate = [3e-4] * self.num_epochs
+ self.learning_rate = [1e-4] * self.num_epochs
# disable the re-sampling of mask proposals to original size for speed-up.
# since evaluation is detection-driven (box-matching) and not instance segmentation-driven (iou-matching),
# mask-outputs are optional.
self.return_masks_in_val = True
self.return_masks_in_test = False
# set number of proposal boxes to plot after each epoch.
self.n_plot_rpn_props = 5 if self.dim == 2 else 30
# number of classes for head networks: n_foreground_classes + 1 (background)
self.head_classes = 3
# seg_classes hier refers to the first stage classifier (RPN)
self.num_seg_classes = 2 # foreground vs. background
# feature map strides per pyramid level are inferred from architecture.
self.backbone_strides = {'xy': [4, 8, 16, 32], 'z': [1, 2, 4, 8]}
# anchor scales are chosen according to expected object sizes in data set. Default uses only one anchor scale
# per pyramid level. (outer list are pyramid levels (corresponding to BACKBONE_STRIDES), inner list are scales per level.)
self.rpn_anchor_scales = {'xy': [[8], [16], [32], [64]], 'z': [[2], [4], [8], [16]]}
# choose which pyramid levels to extract features from: P2: 0, P3: 1, P4: 2, P5: 3.
self.pyramid_levels = [0, 1, 2, 3]
# number of feature maps in rpn. typically lowered in 3D to save gpu-memory.
self.n_rpn_features = 512 if self.dim == 2 else 128
# anchor ratios and strides per position in feature maps.
self.rpn_anchor_ratios = [0.5, 1, 2]
self.rpn_anchor_stride = 1
# Threshold for first stage (RPN) non-maximum suppression (NMS): LOWER == HARDER SELECTION
self.rpn_nms_threshold = 0.7 if self.dim == 2 else 0.7
# loss sampling settings.
- self.rpn_train_anchors_per_image = 64 #per batch element
+ self.rpn_train_anchors_per_image = 6 #per batch element
self.train_rois_per_image = 6 #per batch element
self.roi_positive_ratio = 0.5
self.anchor_matching_iou = 0.7
# factor of top-k candidates to draw from per negative sample (stochastic-hard-example-mining).
# poolsize to draw top-k candidates from will be shem_poolsize * n_negative_samples.
self.shem_poolsize = 10
self.pool_size = (7, 7) if self.dim == 2 else (7, 7, 3)
self.mask_pool_size = (14, 14) if self.dim == 2 else (14, 14, 5)
self.mask_shape = (28, 28) if self.dim == 2 else (28, 28, 10)
self.rpn_bbox_std_dev = np.array([0.1, 0.1, 0.1, 0.2, 0.2, 0.2])
self.bbox_std_dev = np.array([0.1, 0.1, 0.1, 0.2, 0.2, 0.2])
self.window = np.array([0, 0, self.patch_size[0], self.patch_size[1], 0, self.patch_size_3D[2]])
self.scale = np.array([self.patch_size[0], self.patch_size[1], self.patch_size[0], self.patch_size[1],
self.patch_size_3D[2], self.patch_size_3D[2]])
if self.dim == 2:
self.rpn_bbox_std_dev = self.rpn_bbox_std_dev[:4]
self.bbox_std_dev = self.bbox_std_dev[:4]
self.window = self.window[:4]
self.scale = self.scale[:4]
# pre-selection in proposal-layer (stage 1) for NMS-speedup. applied per batch element.
self.pre_nms_limit = 3000 if self.dim == 2 else 6000
# n_proposals to be selected after NMS per batch element. too high numbers blow up memory if "detect_while_training" is True,
# since proposals of the entire batch are forwarded through second stage in as one "batch".
self.roi_chunk_size = 2500 if self.dim == 2 else 600
self.post_nms_rois_training = 500 if self.dim == 2 else 75
self.post_nms_rois_inference = 500
# Final selection of detections (refine_detections)
self.model_max_instances_per_batch_element = 10 if self.dim == 2 else 30 # per batch element and class.
self.detection_nms_threshold = 1e-5 # needs to be > 0, otherwise all predictions are one cluster.
self.model_min_confidence = 0.1
if self.dim == 2:
self.backbone_shapes = np.array(
[[int(np.ceil(self.patch_size[0] / stride)),
int(np.ceil(self.patch_size[1] / stride))]
for stride in self.backbone_strides['xy']])
else:
self.backbone_shapes = np.array(
[[int(np.ceil(self.patch_size[0] / stride)),
int(np.ceil(self.patch_size[1] / stride)),
int(np.ceil(self.patch_size[2] / stride_z))]
for stride, stride_z in zip(self.backbone_strides['xy'], self.backbone_strides['z']
)])
if self.model == 'ufrcnn':
self.operate_stride1 = True
self.class_specific_seg_flag = True
self.num_seg_classes = 3 if self.class_specific_seg_flag else 2
self.frcnn_mode = True
- if self.model == 'retina_net' or self.model == 'retina_unet' or self.model == 'prob_detector':
+ if self.model == 'retina_net' or self.model == 'retina_unet':
# implement extra anchor-scales according to retina-net publication.
self.rpn_anchor_scales['xy'] = [[ii[0], ii[0] * (2 ** (1 / 3)), ii[0] * (2 ** (2 / 3))] for ii in
self.rpn_anchor_scales['xy']]
self.rpn_anchor_scales['z'] = [[ii[0], ii[0] * (2 ** (1 / 3)), ii[0] * (2 ** (2 / 3))] for ii in
self.rpn_anchor_scales['z']]
self.n_anchors_per_pos = len(self.rpn_anchor_ratios) * 3
self.n_rpn_features = 256 if self.dim == 2 else 64
# pre-selection of detections for NMS-speedup. per entire batch.
self.pre_nms_limit = 10000 if self.dim == 2 else 50000
# anchor matching iou is lower than in Mask R-CNN according to https://arxiv.org/abs/1708.02002
self.anchor_matching_iou = 0.5
# if 'True', seg loss distinguishes all classes, else only foreground vs. background (class agnostic).
self.num_seg_classes = 3 if self.class_specific_seg_flag else 2
if self.model == 'retina_unet':
self.operate_stride1 = True
diff --git a/experiments/lidc_exp/data_loader.py b/experiments/lidc_exp/data_loader.py
index 3078ac6..2fe9dde 100644
--- a/experiments/lidc_exp/data_loader.py
+++ b/experiments/lidc_exp/data_loader.py
@@ -1,488 +1,489 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
'''
Example Data Loader for the LIDC data set. This dataloader expects preprocessed data in .npy or .npz files per patient and
a pandas dataframe in the same directory containing the meta-info e.g. file paths, labels, foregound slice-ids.
'''
import numpy as np
import os
from collections import OrderedDict
import pandas as pd
import pickle
import time
import subprocess
-import utils.dataloader_utils as dutils
# batch generator tools from https://github.com/MIC-DKFZ/batchgenerators
from batchgenerators.dataloading.data_loader import SlimDataLoaderBase
from batchgenerators.transforms.spatial_transforms import MirrorTransform as Mirror
from batchgenerators.transforms.abstract_transforms import Compose
from batchgenerators.dataloading.multi_threaded_augmenter import MultiThreadedAugmenter
from batchgenerators.dataloading import SingleThreadedAugmenter
from batchgenerators.transforms.spatial_transforms import SpatialTransform
from batchgenerators.transforms.crop_and_pad_transforms import CenterCropTransform
from batchgenerators.transforms.utility_transforms import ConvertSegToBoundingBoxCoordinates
-
+import utils.dataloader_utils as dutils
+import utils.exp_utils as utils
def get_train_generators(cf, logger):
"""
wrapper function for creating the training batch generator pipeline. returns the train/val generators.
selects patients according to cv folds (generated by first run/fold of experiment):
splits the data into n-folds, where 1 split is used for val, 1 split for testing and the rest for training. (inner loop test set)
If cf.hold_out_test_set is True, adds the test split to the training data.
"""
all_data = load_dataset(cf, logger)
all_pids_list = np.unique([v['pid'] for (k, v) in all_data.items()])
splits_file = os.path.join(cf.exp_dir, 'fold_ids.pickle')
if not os.path.exists(splits_file) and not cf.created_fold_id_pickle:
fg = dutils.fold_generator(seed=cf.seed, n_splits=cf.n_cv_splits, len_data=len(all_pids_list)).get_fold_names()
with open(splits_file, 'wb') as handle:
pickle.dump(fg, handle)
cf.created_fold_id_pickle = True
else:
with open(splits_file, 'rb') as handle:
fg = pickle.load(handle)
train_ix, val_ix, test_ix, _ = fg[cf.fold]
train_pids = [all_pids_list[ix] for ix in train_ix]
val_pids = [all_pids_list[ix] for ix in val_ix]
if cf.hold_out_test_set:
train_pids += [all_pids_list[ix] for ix in test_ix]
train_data = {k: v for (k, v) in all_data.items() if any(p == v['pid'] for p in train_pids)}
val_data = {k: v for (k, v) in all_data.items() if any(p == v['pid'] for p in val_pids)}
logger.info("data set loaded with: {} train / {} val / {} test patients".format(len(train_ix), len(val_ix), len(test_ix)))
batch_gen = {}
batch_gen['train'] = create_data_gen_pipeline(train_data, cf=cf, is_training=True)
batch_gen['val_sampling'] = create_data_gen_pipeline(val_data, cf=cf, is_training=False)
if cf.val_mode == 'val_patient':
batch_gen['val_patient'] = PatientBatchIterator(val_data, cf=cf)
batch_gen['n_val'] = len(val_ix) if cf.max_val_patients is None else min(len(val_ix), cf.max_val_patients)
else:
batch_gen['n_val'] = cf.num_val_batches
return batch_gen
def get_test_generator(cf, logger):
"""
wrapper function for creating the test batch generator pipeline.
selects patients according to cv folds (generated by first run/fold of experiment)
If cf.hold_out_test_set is True, gets the data from an external folder instead.
"""
if cf.hold_out_test_set:
pp_name = cf.pp_test_name
test_ix = None
else:
pp_name = None
with open(os.path.join(cf.exp_dir, 'fold_ids.pickle'), 'rb') as handle:
fold_list = pickle.load(handle)
_, _, test_ix, _ = fold_list[cf.fold]
# warnings.warn('WARNING: using validation set for testing!!!')
test_data = load_dataset(cf, logger, test_ix, pp_data_path=cf.pp_test_data_path, pp_name=pp_name)
logger.info("data set loaded with: {} test patients".format(len(test_ix)))
batch_gen = {}
batch_gen['test'] = PatientBatchIterator(test_data, cf=cf)
batch_gen['n_test'] = len(test_ix) if cf.max_test_patients=="all" else \
min(cf.max_test_patients, len(test_ix))
return batch_gen
def load_dataset(cf, logger, subset_ixs=None, pp_data_path=None, pp_name=None):
"""
loads the dataset. if deployed in cloud also copies and unpacks the data to the working directory.
:param subset_ixs: subset indices to be loaded from the dataset. used e.g. for testing to only load the test folds.
:return: data: dictionary with one entry per patient (in this case per patient-breast, since they are treated as
individual images for training) each entry is a dictionary containing respective meta-info as well as paths to the preprocessed
numpy arrays to be loaded during batch-generation
"""
if pp_data_path is None:
pp_data_path = cf.pp_data_path
if pp_name is None:
pp_name = cf.pp_name
if cf.server_env:
copy_data = True
target_dir = os.path.join(cf.data_dest, pp_name)
if not os.path.exists(target_dir):
cf.data_source_dir = pp_data_path
os.makedirs(target_dir)
subprocess.call('rsync -av {} {}'.format(
os.path.join(cf.data_source_dir, cf.input_df_name), os.path.join(target_dir, cf.input_df_name)), shell=True)
logger.info('created target dir and info df at {}'.format(os.path.join(target_dir, cf.input_df_name)))
elif subset_ixs is None:
copy_data = False
pp_data_path = target_dir
p_df = pd.read_pickle(os.path.join(pp_data_path, cf.input_df_name))
if cf.select_prototype_subset is not None:
prototype_pids = p_df.pid.tolist()[:cf.select_prototype_subset]
p_df = p_df[p_df.pid.isin(prototype_pids)]
logger.warning('WARNING: using prototyping data subset!!!')
if subset_ixs is not None:
subset_pids = [np.unique(p_df.pid.tolist())[ix] for ix in subset_ixs]
p_df = p_df[p_df.pid.isin(subset_pids)]
logger.info('subset: selected {} instances from df'.format(len(p_df)))
if cf.server_env:
if copy_data:
copy_and_unpack_data(logger, p_df.pid.tolist(), cf.fold_dir, cf.data_source_dir, target_dir)
class_targets = p_df['class_target'].tolist()
pids = p_df.pid.tolist()
imgs = [os.path.join(pp_data_path, '{}_img.npy'.format(pid)) for pid in pids]
segs = [os.path.join(pp_data_path,'{}_rois.npy'.format(pid)) for pid in pids]
data = OrderedDict()
for ix, pid in enumerate(pids):
# for the experiment conducted here, malignancy scores are binarized: (benign: 1-2, malignant: 3-5)
targets = [1 if ii >= 3 else 0 for ii in class_targets[ix]]
data[pid] = {'data': imgs[ix], 'seg': segs[ix], 'pid': pid, 'class_target': targets}
data[pid]['fg_slices'] = p_df['fg_slices'].tolist()[ix]
return data
def create_data_gen_pipeline(patient_data, cf, is_training=True):
"""
create mutli-threaded train/val/test batch generation and augmentation pipeline.
:param patient_data: dictionary containing one dictionary per patient in the train/test subset.
:param is_training: (optional) whether to perform data augmentation (training) or not (validation/testing)
:return: multithreaded_generator
"""
# create instance of batch generator as first element in pipeline.
data_gen = BatchGenerator(patient_data, batch_size=cf.batch_size, cf=cf)
# add transformations to pipeline.
my_transforms = []
if is_training:
mirror_transform = Mirror(axes=np.arange(cf.dim))
my_transforms.append(mirror_transform)
spatial_transform = SpatialTransform(patch_size=cf.patch_size[:cf.dim],
patch_center_dist_from_border=cf.da_kwargs['rand_crop_dist'],
do_elastic_deform=cf.da_kwargs['do_elastic_deform'],
alpha=cf.da_kwargs['alpha'], sigma=cf.da_kwargs['sigma'],
do_rotation=cf.da_kwargs['do_rotation'], angle_x=cf.da_kwargs['angle_x'],
angle_y=cf.da_kwargs['angle_y'], angle_z=cf.da_kwargs['angle_z'],
do_scale=cf.da_kwargs['do_scale'], scale=cf.da_kwargs['scale'],
random_crop=cf.da_kwargs['random_crop'])
my_transforms.append(spatial_transform)
else:
my_transforms.append(CenterCropTransform(crop_size=cf.patch_size[:cf.dim]))
my_transforms.append(ConvertSegToBoundingBoxCoordinates(cf.dim, get_rois_from_seg_flag=False, class_specific_seg_flag=cf.class_specific_seg_flag))
all_transforms = Compose(my_transforms)
# multithreaded_generator = SingleThreadedAugmenter(data_gen, all_transforms)
multithreaded_generator = MultiThreadedAugmenter(data_gen, all_transforms, num_processes=cf.n_workers, seeds=range(cf.n_workers))
return multithreaded_generator
class BatchGenerator(SlimDataLoaderBase):
"""
creates the training/validation batch generator. Samples n_batch_size patients (draws a slice from each patient if 2D)
from the data set while maintaining foreground-class balance. Returned patches are cropped/padded to pre_crop_size.
Actual patch_size is obtained after data augmentation.
:param data: data dictionary as provided by 'load_dataset'.
:param batch_size: number of patients to sample for the batch
- :return dictionary containing the batch data (b, c, x, y, (z)) / seg (b, 1, x, y, (z)) / pids / class_target
+ :return dictionary containing the batch data (b, c, y, x(, z)) / seg (b, 1, y, x(, z)) / pids / class_target
"""
def __init__(self, data, batch_size, cf):
super(BatchGenerator, self).__init__(data, batch_size)
self.cf = cf
self.crop_margin = np.array(self.cf.patch_size)/8. #min distance of ROI center to edge of cropped_patch.
self.p_fg = 0.5
def generate_train_batch(self):
batch_data, batch_segs, batch_pids, batch_targets, batch_patient_labels = [], [], [], [], []
class_targets_list = [v['class_target'] for (k, v) in self._data.items()]
if self.cf.head_classes > 2:
# samples patients towards equilibrium of foreground classes on a roi-level (after randomly sampling the ratio "batch_sample_slack).
batch_ixs = dutils.get_class_balanced_patients(
class_targets_list, self.batch_size, self.cf.head_classes - 1, slack_factor=self.cf.batch_sample_slack)
else:
batch_ixs = np.random.choice(len(class_targets_list), self.batch_size)
patients = list(self._data.items())
for b in batch_ixs:
patient = patients[b][1]
- data = np.transpose(np.load(patient['data'], mmap_mode='r'), axes=(1, 2, 0))[np.newaxis] # (c, y, x, z)
+ # data shape: from (z,y,x) to (c, y, x, z).
+ data = np.transpose(np.load(patient['data'], mmap_mode='r'), axes=(1, 2, 0))[np.newaxis]
seg = np.transpose(np.load(patient['seg'], mmap_mode='r'), axes=(1, 2, 0))
batch_pids.append(patient['pid'])
batch_targets.append(patient['class_target'])
if self.cf.dim == 2:
# draw random slice from patient while oversampling slices containing foreground objects with p_fg.
if len(patient['fg_slices']) > 0:
fg_prob = self.p_fg / len(patient['fg_slices'])
bg_prob = (1 - self.p_fg) / (data.shape[3] - len(patient['fg_slices']))
slices_prob = [fg_prob if ix in patient['fg_slices'] else bg_prob for ix in range(data.shape[3])]
slice_id = np.random.choice(data.shape[3], p=slices_prob)
else:
slice_id = np.random.choice(data.shape[3])
# if set to not None, add neighbouring slices to each selected slice in channel dimension.
if self.cf.n_3D_context is not None:
padded_data = dutils.pad_nd_image(data[0], [(data.shape[-1] + (self.cf.n_3D_context*2))], mode='constant')
padded_slice_id = slice_id + self.cf.n_3D_context
data = (np.concatenate([padded_data[..., ii][np.newaxis] for ii in range(
padded_slice_id - self.cf.n_3D_context, padded_slice_id + self.cf.n_3D_context + 1)], axis=0))
else:
data = data[..., slice_id]
seg = seg[..., slice_id]
# pad data if smaller than pre_crop_size.
if np.any([data.shape[dim + 1] < ps for dim, ps in enumerate(self.cf.pre_crop_size)]):
new_shape = [np.max([data.shape[dim + 1], ps]) for dim, ps in enumerate(self.cf.pre_crop_size)]
data = dutils.pad_nd_image(data, new_shape, mode='constant')
seg = dutils.pad_nd_image(seg, new_shape, mode='constant')
# crop patches of size pre_crop_size, while sampling patches containing foreground with p_fg.
crop_dims = [dim for dim, ps in enumerate(self.cf.pre_crop_size) if data.shape[dim + 1] > ps]
if len(crop_dims) > 0:
fg_prob_sample = np.random.rand(1)
# with p_fg: sample random pixel from random ROI and shift center by random value.
if fg_prob_sample < self.p_fg and np.sum(seg) > 0:
seg_ixs = np.argwhere(seg == np.random.choice(np.unique(seg)[1:], 1))
roi_anchor_pixel = seg_ixs[np.random.choice(seg_ixs.shape[0], 1)][0]
assert seg[tuple(roi_anchor_pixel)] > 0
# sample the patch center coords. constrained by edges of images - pre_crop_size /2. And by
# distance to the desired ROI < patch_size /2.
# (here final patch size to account for center_crop after data augmentation).
sample_seg_center = {}
for ii in crop_dims:
low = np.max((self.cf.pre_crop_size[ii]//2, roi_anchor_pixel[ii] - (self.cf.patch_size[ii]//2 - self.crop_margin[ii])))
high = np.min((data.shape[ii + 1] - self.cf.pre_crop_size[ii]//2,
roi_anchor_pixel[ii] + (self.cf.patch_size[ii]//2 - self.crop_margin[ii])))
# happens if lesion on the edge of the image. dont care about roi anymore,
# just make sure pre-crop is inside image.
if low >= high:
low = data.shape[ii + 1] // 2 - (data.shape[ii + 1] // 2 - self.cf.pre_crop_size[ii] // 2)
high = data.shape[ii + 1] // 2 + (data.shape[ii + 1] // 2 - self.cf.pre_crop_size[ii] // 2)
sample_seg_center[ii] = np.random.randint(low=low, high=high)
else:
# not guaranteed to be empty. probability of emptiness depends on the data.
sample_seg_center = {ii: np.random.randint(low=self.cf.pre_crop_size[ii]//2,
high=data.shape[ii + 1] - self.cf.pre_crop_size[ii]//2) for ii in crop_dims}
for ii in crop_dims:
min_crop = int(sample_seg_center[ii] - self.cf.pre_crop_size[ii] // 2)
max_crop = int(sample_seg_center[ii] + self.cf.pre_crop_size[ii] // 2)
data = np.take(data, indices=range(min_crop, max_crop), axis=ii + 1)
seg = np.take(seg, indices=range(min_crop, max_crop), axis=ii)
batch_data.append(data)
batch_segs.append(seg[np.newaxis])
data = np.array(batch_data)
seg = np.array(batch_segs).astype(np.uint8)
class_target = np.array(batch_targets)
return {'data': data, 'seg': seg, 'pid': batch_pids, 'class_target': class_target}
class PatientBatchIterator(SlimDataLoaderBase):
"""
creates a test generator that iterates over entire given dataset returning 1 patient per batch.
Can be used for monitoring if cf.val_mode = 'patient_val' for a monitoring closer to actualy evaluation (done in 3D),
if willing to accept speed-loss during training.
:return: out_batch: dictionary containing one patient with batch_size = n_3D_patches in 3D or
batch_size = n_2D_patches in 2D .
"""
def __init__(self, data, cf): #threads in augmenter
super(PatientBatchIterator, self).__init__(data, 0)
self.cf = cf
self.patient_ix = 0
self.dataset_pids = [v['pid'] for (k, v) in data.items()]
self.patch_size = cf.patch_size
if len(self.patch_size) == 2:
self.patch_size = self.patch_size + [1]
def generate_train_batch(self):
pid = self.dataset_pids[self.patient_ix]
patient = self._data[pid]
data = np.transpose(np.load(patient['data'], mmap_mode='r'), axes=(1, 2, 0))[np.newaxis] # (c, y, x, z)
seg = np.transpose(np.load(patient['seg'], mmap_mode='r'), axes=(1, 2, 0))
batch_class_targets = np.array([patient['class_target']])
# pad data if smaller than patch_size seen during training.
if np.any([data.shape[dim + 1] < ps for dim, ps in enumerate(self.patch_size)]):
new_shape = [data.shape[0]] + [np.max([data.shape[dim + 1], self.patch_size[dim]]) for dim, ps in enumerate(self.patch_size)]
data = dutils.pad_nd_image(data, new_shape) # use 'return_slicer' to crop image back to original shape.
seg = dutils.pad_nd_image(seg, new_shape)
# get 3D targets for evaluation, even if network operates in 2D. 2D predictions will be merged to 3D in predictor.
if self.cf.dim == 3 or self.cf.merge_2D_to_3D_preds:
out_data = data[np.newaxis]
out_seg = seg[np.newaxis, np.newaxis]
out_targets = batch_class_targets
batch_3D = {'data': out_data, 'seg': out_seg, 'class_target': out_targets, 'pid': pid}
converter = ConvertSegToBoundingBoxCoordinates(dim=3, get_rois_from_seg_flag=False, class_specific_seg_flag=self.cf.class_specific_seg_flag)
batch_3D = converter(**batch_3D)
batch_3D.update({'patient_bb_target': batch_3D['bb_target'],
- 'patient_roi_labels': batch_3D['roi_labels'],
+ 'patient_roi_labels': batch_3D['class_target'],
'original_img_shape': out_data.shape})
if self.cf.dim == 2:
out_data = np.transpose(data, axes=(3, 0, 1, 2)) # (z, c, y, x )
out_seg = np.transpose(seg, axes=(2, 0, 1))[:, np.newaxis]
out_targets = np.array(np.repeat(batch_class_targets, out_data.shape[0], axis=0))
# if set to not None, add neighbouring slices to each selected slice in channel dimension.
if self.cf.n_3D_context is not None:
slice_range = range(self.cf.n_3D_context, out_data.shape[0] + self.cf.n_3D_context)
out_data = np.pad(out_data, ((self.cf.n_3D_context, self.cf.n_3D_context), (0, 0), (0, 0), (0, 0)), 'constant', constant_values=0)
out_data = np.array(
[np.concatenate([out_data[ii] for ii in range(
slice_id - self.cf.n_3D_context, slice_id + self.cf.n_3D_context + 1)], axis=0) for slice_id in
slice_range])
batch_2D = {'data': out_data, 'seg': out_seg, 'class_target': out_targets, 'pid': pid}
converter = ConvertSegToBoundingBoxCoordinates(dim=2, get_rois_from_seg_flag=False, class_specific_seg_flag=self.cf.class_specific_seg_flag)
batch_2D = converter(**batch_2D)
if self.cf.merge_2D_to_3D_preds:
batch_2D.update({'patient_bb_target': batch_3D['patient_bb_target'],
'patient_roi_labels': batch_3D['patient_roi_labels'],
'original_img_shape': out_data.shape})
else:
batch_2D.update({'patient_bb_target': batch_2D['bb_target'],
- 'patient_roi_labels': batch_2D['roi_labels'],
+ 'patient_roi_labels': batch_2D['class_target'],
'original_img_shape': out_data.shape})
out_batch = batch_3D if self.cf.dim == 3 else batch_2D
patient_batch = out_batch
# crop patient-volume to patches of patch_size used during training. stack patches up in batch dimension.
# in this case, 2D is treated as a special case of 3D with patch_size[z] = 1.
if np.any([data.shape[dim + 1] > self.patch_size[dim] for dim in range(3)]):
patch_crop_coords_list = dutils.get_patch_crop_coords(data[0], self.patch_size)
new_img_batch, new_seg_batch, new_class_targets_batch = [], [], []
for cix, c in enumerate(patch_crop_coords_list):
seg_patch = seg[c[0]:c[1], c[2]: c[3], c[4]:c[5]]
new_seg_batch.append(seg_patch)
# if set to not None, add neighbouring slices to each selected slice in channel dimension.
# correct patch_crop coordinates by added slices of 3D context.
if self.cf.dim == 2 and self.cf.n_3D_context is not None:
tmp_c_5 = c[5] + (self.cf.n_3D_context * 2)
if cix == 0:
data = np.pad(data, ((0, 0), (0, 0), (0, 0), (self.cf.n_3D_context, self.cf.n_3D_context)), 'constant', constant_values=0)
else:
tmp_c_5 = c[5]
new_img_batch.append(data[:, c[0]:c[1], c[2]:c[3], c[4]:tmp_c_5])
data = np.array(new_img_batch) # (n_patches, c, x, y, z)
seg = np.array(new_seg_batch)[:, np.newaxis] # (n_patches, 1, x, y, z)
batch_class_targets = np.repeat(batch_class_targets, len(patch_crop_coords_list), axis=0)
if self.cf.dim == 2:
if self.cf.n_3D_context is not None:
data = np.transpose(data[:, 0], axes=(0, 3, 1, 2))
else:
# all patches have z dimension 1 (slices). discard dimension
data = data[..., 0]
seg = seg[..., 0]
patch_batch = {'data': data, 'seg': seg, 'class_target': batch_class_targets, 'pid': pid}
patch_batch['patch_crop_coords'] = np.array(patch_crop_coords_list)
patch_batch['patient_bb_target'] = patient_batch['patient_bb_target']
patch_batch['patient_roi_labels'] = patient_batch['patient_roi_labels']
patch_batch['original_img_shape'] = patient_batch['original_img_shape']
converter = ConvertSegToBoundingBoxCoordinates(self.cf.dim, get_rois_from_seg_flag=False, class_specific_seg_flag=self.cf.class_specific_seg_flag)
patch_batch = converter(**patch_batch)
out_batch = patch_batch
self.patient_ix += 1
if self.patient_ix == len(self.dataset_pids):
self.patient_ix = 0
return out_batch
def copy_and_unpack_data(logger, pids, fold_dir, source_dir, target_dir):
-
start_time = time.time()
with open(os.path.join(fold_dir, 'file_list.txt'), 'w') as handle:
for pid in pids:
handle.write('{}_img.npz\n'.format(pid))
handle.write('{}_rois.npz\n'.format(pid))
subprocess.call('rsync -av --files-from {} {} {}'.format(os.path.join(fold_dir, 'file_list.txt'),
source_dir, target_dir), shell=True)
- dutils.unpack_dataset(target_dir)
+ n_threads = 8
+ dutils.unpack_dataset(target_dir, threads=n_threads)
copied_files = os.listdir(target_dir)
- logger.info("copying and unpacking data set finished : {} files in target dir: {}. took {} sec".format(
- len(copied_files), target_dir, np.round(time.time() - start_time, 0)))
+ t = utils.get_formatted_duration(time.time() - start_time)
+ logger.info("\ncopying and unpacking data set finished using {} threads.\n{} files in target dir: {}. Took {}\n"
+ .format(n_threads, len(copied_files), target_dir, t))
if __name__=="__main__":
- import utils.exp_utils as utils
total_stime = time.time()
cf_file = utils.import_module("cf", "configs.py")
cf = cf_file.configs()
cf.created_fold_id_pickle = False
cf.exp_dir = "dev/"
cf.plot_dir = cf.exp_dir + "plots"
os.makedirs(cf.exp_dir, exist_ok=True)
cf.fold = 0
logger = utils.get_logger(cf.exp_dir)
#batch_gen = get_train_generators(cf, logger)
#train_batch = next(batch_gen["train"])
test_gen = get_test_generator(cf, logger)
test_batch = next(test_gen["test"])
mins, secs = divmod((time.time() - total_stime), 60)
h, mins = divmod(mins, 60)
t = "{:d}h:{:02d}m:{:02d}s".format(int(h), int(mins), int(secs))
print("{} total runtime: {}".format(os.path.split(__file__)[1], t))
\ No newline at end of file
diff --git a/experiments/lidc_exp/preprocessing.py b/experiments/lidc_exp/preprocessing.py
index 975d474..a99eea3 100644
--- a/experiments/lidc_exp/preprocessing.py
+++ b/experiments/lidc_exp/preprocessing.py
@@ -1,149 +1,150 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
'''
This preprocessing script loads nrrd files obtained by the data conversion tool: https://github.com/MIC-DKFZ/LIDC-IDRI-processing/tree/v1.0.1
After applying preprocessing, images are saved as numpy arrays and the meta information for the corresponding patient is stored
as a line in the dataframe saved as info_df.pickle.
'''
import os, sys
from pathlib import Path
import SimpleITK as sitk
import numpy as np
from multiprocessing import Pool
import pandas as pd
import numpy.testing as npt
from skimage.transform import resize
import subprocess
import pickle
PROJECT_ROOT = Path(__file__).absolute().parent.parent.parent
sys.path.append(str(PROJECT_ROOT))
import utils.exp_utils as utils
def resample_array(src_imgs, src_spacing, target_spacing):
src_spacing = np.round(src_spacing, 3)
target_shape = [int(src_imgs.shape[ix] * src_spacing[::-1][ix] / target_spacing[::-1][ix]) for ix in range(len(src_imgs.shape))]
for i in range(len(target_shape)):
try:
assert target_shape[i] > 0
except:
raise AssertionError("AssertionError:", src_imgs.shape, src_spacing, target_spacing)
img = src_imgs.astype(float)
resampled_img = resize(img, target_shape, order=1, clip=True, mode='edge').astype('float32')
return resampled_img
def pp_patient(inputs):
ix, path = inputs
pid = path.split('/')[-1]
img = sitk.ReadImage(os.path.join(path, '{}_ct_scan.nrrd'.format(pid)))
+ # sitk.GetArray switches an image with shape (x,y,z) to (z,y,x)
img_arr = sitk.GetArrayFromImage(img)
print('processing {}'.format(pid), img.GetSpacing(), img_arr.shape)
img_arr = resample_array(img_arr, img.GetSpacing(), cf.target_spacing)
img_arr = np.clip(img_arr, -1200, 600)
#img_arr = (1200 + img_arr) / (600 + 1200) * 255 # a+x / (b-a) * (c-d) (c, d = new)
img_arr = img_arr.astype(np.float32)
img_arr = (img_arr - np.mean(img_arr)) / np.std(img_arr).astype(np.float16)
df = pd.read_csv(os.path.join(cf.root_dir, 'characteristics.csv'), sep=';')
df = df[df.PatientID == pid]
final_rois = np.zeros_like(img_arr, dtype=np.uint8)
mal_labels = []
roi_ids = set([ii.split('.')[0].split('_')[-1] for ii in os.listdir(path) if '.nii.gz' in ii])
rix = 1
for rid in roi_ids:
roi_id_paths = [ii for ii in os.listdir(path) if '{}.nii'.format(rid) in ii]
nodule_ids = [ii.split('_')[2].lstrip("0") for ii in roi_id_paths]
rater_labels = [df[df.NoduleID == int(ii)].Malignancy.values[0] for ii in nodule_ids]
rater_labels.extend([0] * (4-len(rater_labels)))
mal_label = np.mean([ii for ii in rater_labels if ii > -1])
roi_rater_list = []
for rp in roi_id_paths:
roi = sitk.ReadImage(os.path.join(cf.raw_data_dir, pid, rp))
roi_arr = sitk.GetArrayFromImage(roi).astype(np.uint8)
roi_arr = resample_array(roi_arr, roi.GetSpacing(), cf.target_spacing)
assert roi_arr.shape == img_arr.shape, [roi_arr.shape, img_arr.shape, pid, roi.GetSpacing()]
for ix in range(len(img_arr.shape)):
npt.assert_almost_equal(roi.GetSpacing()[ix], img.GetSpacing()[ix])
roi_rater_list.append(roi_arr)
roi_rater_list.extend([np.zeros_like(roi_rater_list[-1])]*(4-len(roi_id_paths)))
roi_raters = np.array(roi_rater_list)
roi_raters = np.mean(roi_raters, axis=0)
roi_raters[roi_raters < 0.5] = 0
if np.sum(roi_raters) > 0:
mal_labels.append(mal_label)
final_rois[roi_raters >= 0.5] = rix
rix += 1
else:
# indicate rois suppressed by majority voting of raters
print('suppressed roi!', roi_id_paths)
with open(os.path.join(cf.pp_dir, 'suppressed_rois.txt'), 'a') as handle:
handle.write(" ".join(roi_id_paths))
fg_slices = [ii for ii in np.unique(np.argwhere(final_rois != 0)[:, 0])]
mal_labels = np.array(mal_labels)
assert len(mal_labels) + 1 == len(np.unique(final_rois)), [len(mal_labels), np.unique(final_rois), pid]
np.save(os.path.join(cf.pp_dir, '{}_rois.npy'.format(pid)), final_rois)
np.save(os.path.join(cf.pp_dir, '{}_img.npy'.format(pid)), img_arr)
with open(os.path.join(cf.pp_dir, 'meta_info_{}.pickle'.format(pid)), 'wb') as handle:
meta_info_dict = {'pid': pid, 'class_target': mal_labels, 'spacing': img.GetSpacing(), 'fg_slices': fg_slices}
pickle.dump(meta_info_dict, handle)
def aggregate_meta_info(exp_dir):
files = [os.path.join(exp_dir, f) for f in os.listdir(exp_dir) if 'meta_info' in f]
df = pd.DataFrame(columns=['pid', 'class_target', 'spacing', 'fg_slices'])
for f in files:
with open(f, 'rb') as handle:
df.loc[len(df)] = pickle.load(handle)
df.to_pickle(os.path.join(exp_dir, 'info_df.pickle'))
print ("aggregated meta info to df with length", len(df))
if __name__ == "__main__":
cf_file = utils.import_module("cf", "configs.py")
cf = cf_file.configs()
paths = [os.path.join(cf.raw_data_dir, ii) for ii in os.listdir(cf.raw_data_dir)]
if not os.path.exists(cf.pp_dir):
os.mkdir(cf.pp_dir)
pool = Pool(processes=os.cpu_count())
p1 = pool.map(pp_patient, enumerate(paths))
pool.close()
pool.join()
# for i in enumerate(paths):
# pp_patient(i)
aggregate_meta_info(cf.pp_dir)
subprocess.call('cp {} {}'.format(os.path.join(cf.pp_dir, 'info_df.pickle'), os.path.join(cf.pp_dir, 'info_df_bk.pickle')), shell=True)
diff --git a/experiments/toy_exp/configs.py b/experiments/toy_exp/configs.py
index 0376e0c..38fe39f 100644
--- a/experiments/toy_exp/configs.py
+++ b/experiments/toy_exp/configs.py
@@ -1,351 +1,351 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import sys
import os
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
import numpy as np
from default_configs import DefaultConfigs
class configs(DefaultConfigs):
- def __init__(self, server_env=None):
+ def __init__(self, server_env=False):
#########################
# Preprocessing #
#########################
self.root_dir = '/home/gregor/datasets/toy_mdt'
#########################
# I/O #
#########################
# one out of [2, 3]. dimension the model operates in.
self.dim = 2
# one out of ['mrcnn', 'retina_net', 'retina_unet', 'detection_unet', 'ufrcnn'].
- self.model = 'ufrcnn'
+ self.model = 'retina_unet'
DefaultConfigs.__init__(self, self.model, server_env, self.dim)
# int [0 < dataset_size]. select n patients from dataset for prototyping.
self.select_prototype_subset = None
self.hold_out_test_set = True
# including val set. will be 3/4 train, 1/4 val.
self.n_train_val_data = 2500
# choose one of the 3 toy experiments described in https://arxiv.org/pdf/1811.08661.pdf
# one of ['donuts_shape', 'donuts_pattern', 'circles_scale'].
toy_mode = 'donuts_shape_noise'
# path to preprocessed data.
self.input_df_name = 'info_df.pickle'
self.pp_name = os.path.join(toy_mode, 'train')
self.pp_data_path = os.path.join(self.root_dir, self.pp_name)
self.pp_test_name = os.path.join(toy_mode, 'test')
self.pp_test_data_path = os.path.join(self.root_dir, self.pp_test_name)
# settings for deployment in cloud.
if server_env:
# path to preprocessed data.
pp_root_dir = '/datasets/datasets_ramien/toy_exp/data'
self.pp_name = os.path.join(toy_mode, 'train')
self.pp_data_path = os.path.join(pp_root_dir, self.pp_name)
self.pp_test_name = os.path.join(toy_mode, 'test')
self.pp_test_data_path = os.path.join(pp_root_dir, self.pp_test_name)
self.select_prototype_subset = None
#########################
# Data Loader #
#########################
# select modalities from preprocessed data
self.channels = [0]
self.n_channels = len(self.channels)
# patch_size to be used for training. pre_crop_size is the patch_size before data augmentation.
self.pre_crop_size_2D = [320, 320]
self.patch_size_2D = [320, 320]
self.patch_size = self.patch_size_2D if self.dim == 2 else self.patch_size_3D
self.pre_crop_size = self.pre_crop_size_2D if self.dim == 2 else self.pre_crop_size_3D
# ratio of free sampled batch elements before class balancing is triggered
# (>0 to include "empty"/background patches.)
self.batch_sample_slack = 0.2
# set 2D network to operate in 3D images.
self.merge_2D_to_3D_preds = False
# feed +/- n neighbouring slices into channel dimension. set to None for no context.
self.n_3D_context = None
if self.n_3D_context is not None and self.dim == 2:
self.n_channels *= (self.n_3D_context * 2 + 1)
#########################
# Architecture #
#########################
self.start_filts = 48 if self.dim == 2 else 18
self.end_filts = self.start_filts * 4 if self.dim == 2 else self.start_filts * 2
- self.res_architecture = 'resnet50' # 'resnet101' , 'resnet50'
+ self.res_architecture = 'resnet50' # 'resnet101', 'resnet50'
self.norm = None # one of None, 'instance_norm', 'batch_norm'
- self.weight_decay = 1e-6
+ # 0 for no weight decay
+ self.weight_decay = 3e-6
+ # which weights to exclude from weight decay, options: ["norm", "bias"].
self.exclude_from_wd = ("norm",)
- # one of 'xavier_uniform', 'xavier_normal', or 'kaiming_normal', None (=default = 'kaiming_uniform')
+ # one of 'xavier_uniform', 'xavier_normal', or 'kaiming_normal', None (= default = 'kaiming_uniform')
self.weight_init = None
#########################
# Schedule / Selection #
#########################
self.num_epochs = 28
self.num_train_batches = 100 if self.dim == 2 else 200
self.batch_size = 16 if self.dim == 2 else 8
self.do_validation = True
# decide whether to validate on entire patient volumes (like testing) or sampled patches (like training)
- # the former is morge accurate, while the latter is faster (depending on volume size)
+ # the former is more accurate, while the latter is faster (depending on volume size)
self.val_mode = 'val_patient' # one of 'val_sampling' , 'val_patient'
if self.val_mode == 'val_patient':
self.max_val_patients = None # if 'None' iterates over entire val_set once.
if self.val_mode == 'val_sampling':
self.num_val_batches = 50
# set dynamic_lr_scheduling to True to apply LR scheduling with below settings.
self.dynamic_lr_scheduling = True
self.lr_decay_factor = 0.5
self.scheduling_patience = np.ceil(7200 / (self.num_train_batches * self.batch_size))
self.scheduling_criterion = 'malignant_ap'
self.scheduling_mode = 'min' if "loss" in self.scheduling_criterion else 'max'
#########################
# Testing / Plotting #
#########################
- # set the top-n-epochs to be saved for temporal averaging in testing.
+ # set the top-n epochs to be saved for temporal averaging in testing.
self.save_n_models = 5
self.test_n_epochs = 5
# set a minimum epoch number for saving in case of instabilities in the first phase of training.
self.min_save_thresh = 0 if self.dim == 2 else 0
self.report_score_level = ['patient', 'rois'] # choose list from 'patient', 'rois'
self.class_dict = {1: 'benign', 2: 'malignant'} # 0 is background.
self.patient_class_of_interest = 2 # patient metrics are only plotted for one class.
self.ap_match_ious = [0.1] # list of ious to be evaluated for ap-scoring.
self.model_selection_criteria = ['benign_ap', 'malignant_ap'] # criteria to average over for saving epochs.
self.min_det_thresh = 0.1 # minimum confidence value to select predictions for evaluation.
# threshold for clustering predictions together (wcs = weighted cluster scoring).
# needs to be >= the expected overlap of predictions coming from one model (typically NMS threshold).
# if too high, preds of the same object are separate clusters.
self.wcs_iou = 1e-5
self.plot_prediction_histograms = True
self.plot_stat_curves = False
#########################
# Data Augmentation #
#########################
self.da_kwargs={
'do_elastic_deform': True,
'alpha':(0., 1500.),
'sigma':(30., 50.),
'do_rotation':True,
'angle_x': (0., 2 * np.pi),
'angle_y': (0., 0),
'angle_z': (0., 0),
'do_scale': True,
'scale':(0.8, 1.1),
'random_crop':False,
'rand_crop_dist': (self.patch_size[0] / 2. - 3, self.patch_size[1] / 2. - 3),
'border_mode_data': 'constant',
'border_cval_data': 0,
'order_data': 1
}
if self.dim == 3:
self.da_kwargs['do_elastic_deform'] = False
self.da_kwargs['angle_x'] = (0, 0.0)
self.da_kwargs['angle_y'] = (0, 0.0) #must be 0!!
self.da_kwargs['angle_z'] = (0., 2 * np.pi)
#########################
# Add model specifics #
#########################
{'detection_unet': self.add_det_unet_configs,
'mrcnn': self.add_mrcnn_configs,
'ufrcnn': self.add_mrcnn_configs,
- 'ufrcnn_surrounding': self.add_mrcnn_configs,
'retina_net': self.add_mrcnn_configs,
'retina_unet': self.add_mrcnn_configs,
- 'prob_detector': self.add_mrcnn_configs,
}[self.model]()
def add_det_unet_configs(self):
self.learning_rate = [1e-4] * self.num_epochs
# aggregation from pixel perdiction to object scores (connected component). One of ['max', 'median']
self.aggregation_operation = 'max'
# max number of roi candidates to identify per image (slice in 2D, volume in 3D)
self.n_roi_candidates = 3 if self.dim == 2 else 8
# loss mode: either weighted cross entropy ('wce'), batch-wise dice loss ('dice), or the sum of both ('dice_wce')
self.seg_loss_mode = 'dice_wce'
# if <1, false positive predictions in foreground are penalized less.
self.fp_dice_weight = 1 if self.dim == 2 else 1
self.wce_weights = [0.3, 1, 1]
self.detection_min_confidence = self.min_det_thresh
# if 'True', loss distinguishes all classes, else only foreground vs. background (class agnostic).
self.class_specific_seg_flag = True
self.num_seg_classes = 3 if self.class_specific_seg_flag else 2
self.head_classes = self.num_seg_classes
def add_mrcnn_configs(self):
# learning rate is a list with one entry per epoch.
self.learning_rate = [3e-4] * self.num_epochs
# disable mask head loss. (e.g. if no pixelwise annotations available)
self.frcnn_mode = False
# disable the re-sampling of mask proposals to original size for speed-up.
# since evaluation is detection-driven (box-matching) and not instance segmentation-driven (iou-matching),
# mask-outputs are optional.
self.return_masks_in_val = True
self.return_masks_in_test = False
# set number of proposal boxes to plot after each epoch.
self.n_plot_rpn_props = 0 if self.dim == 2 else 0
# number of classes for head networks: n_foreground_classes + 1 (background)
self.head_classes = 3
# seg_classes hier refers to the first stage classifier (RPN)
self.num_seg_classes = 2 # foreground vs. background
# feature map strides per pyramid level are inferred from architecture.
self.backbone_strides = {'xy': [4, 8, 16, 32], 'z': [1, 2, 4, 8]}
# anchor scales are chosen according to expected object sizes in data set. Default uses only one anchor scale
# per pyramid level. (outer list are pyramid levels (corresponding to BACKBONE_STRIDES), inner list are scales per level.)
self.rpn_anchor_scales = {'xy': [[8], [16], [32], [64]], 'z': [[2], [4], [8], [16]]}
# choose which pyramid levels to extract features from: P2: 0, P3: 1, P4: 2, P5: 3.
self.pyramid_levels = [0, 1, 2, 3]
# number of feature maps in rpn. typically lowered in 3D to save gpu-memory.
self.n_rpn_features = 512 if self.dim == 2 else 128
# anchor ratios and strides per position in feature maps.
self.rpn_anchor_ratios = [0.5, 1., 2.]
self.rpn_anchor_stride = 1
# Threshold for first stage (RPN) non-maximum suppression (NMS): LOWER == HARDER SELECTION
self.rpn_nms_threshold = 0.7 if self.dim == 2 else 0.7
# loss sampling settings.
self.rpn_train_anchors_per_image = 64 #per batch element
self.train_rois_per_image = 2 #per batch element
self.roi_positive_ratio = 0.5
self.anchor_matching_iou = 0.7
# factor of top-k candidates to draw from per negative sample (stochastic-hard-example-mining).
# poolsize to draw top-k candidates from will be shem_poolsize * n_negative_samples.
self.shem_poolsize = 10
self.pool_size = (7, 7) if self.dim == 2 else (7, 7, 3)
self.mask_pool_size = (14, 14) if self.dim == 2 else (14, 14, 5)
self.mask_shape = (28, 28) if self.dim == 2 else (28, 28, 10)
self.rpn_bbox_std_dev = np.array([0.1, 0.1, 0.1, 0.2, 0.2, 0.2])
self.bbox_std_dev = np.array([0.1, 0.1, 0.1, 0.2, 0.2, 0.2])
self.window = np.array([0, 0, self.patch_size[0], self.patch_size[1]])
self.scale = np.array([self.patch_size[0], self.patch_size[1], self.patch_size[0], self.patch_size[1]])
if self.dim == 2:
self.rpn_bbox_std_dev = self.rpn_bbox_std_dev[:4]
self.bbox_std_dev = self.bbox_std_dev[:4]
self.window = self.window[:4]
self.scale = self.scale[:4]
# pre-selection in proposal-layer (stage 1) for NMS-speedup. applied per batch element.
self.pre_nms_limit = 3000 if self.dim == 2 else 6000
# n_proposals to be selected after NMS per batch element. too high numbers blow up memory if "detect_while_training" is True,
# since proposals of the entire batch are forwarded through second stage in as one "batch".
self.roi_chunk_size = 800 if self.dim == 2 else 600
self.post_nms_rois_training = 500 if self.dim == 2 else 75
self.post_nms_rois_inference = 500
# Final selection of detections (refine_detections)
self.model_max_instances_per_batch_element = 10 if self.dim == 2 else 30 # per batch element and class.
self.detection_nms_threshold = 1e-5 # needs to be > 0, otherwise all predictions are one cluster.
self.model_min_confidence = 0.1
if self.dim == 2:
self.backbone_shapes = np.array(
[[int(np.ceil(self.patch_size[0] / stride)),
int(np.ceil(self.patch_size[1] / stride))]
for stride in self.backbone_strides['xy']])
else:
self.backbone_shapes = np.array(
[[int(np.ceil(self.patch_size[0] / stride)),
int(np.ceil(self.patch_size[1] / stride)),
int(np.ceil(self.patch_size[2] / stride_z))]
for stride, stride_z in zip(self.backbone_strides['xy'], self.backbone_strides['z']
)])
if self.model == 'ufrcnn':
self.operate_stride1 = True
self.class_specific_seg_flag = True
self.num_seg_classes = 3 if self.class_specific_seg_flag else 2
self.frcnn_mode = True
- if self.model == 'retina_net' or self.model == 'retina_unet' or self.model == 'prob_detector':
+ if self.model == 'retina_net' or self.model == 'retina_unet':
# implement extra anchor-scales according to retina-net publication.
self.rpn_anchor_scales['xy'] = [[ii[0], ii[0] * (2 ** (1 / 3)), ii[0] * (2 ** (2 / 3))] for ii in
self.rpn_anchor_scales['xy']]
self.rpn_anchor_scales['z'] = [[ii[0], ii[0] * (2 ** (1 / 3)), ii[0] * (2 ** (2 / 3))] for ii in
self.rpn_anchor_scales['z']]
self.n_anchors_per_pos = len(self.rpn_anchor_ratios) * 3
self.n_rpn_features = 256 if self.dim == 2 else 64
# pre-selection of detections for NMS-speedup. per entire batch.
self.pre_nms_limit = 10000 if self.dim == 2 else 50000
# anchor matching iou is lower than in Mask R-CNN according to https://arxiv.org/abs/1708.02002
self.anchor_matching_iou = 0.5
# if 'True', seg loss distinguishes all classes, else only foreground vs. background (class agnostic).
self.num_seg_classes = 3 if self.class_specific_seg_flag else 2
if self.model == 'retina_unet':
self.operate_stride1 = True
diff --git a/experiments/toy_exp/data_loader.py b/experiments/toy_exp/data_loader.py
index c123011..384f506 100644
--- a/experiments/toy_exp/data_loader.py
+++ b/experiments/toy_exp/data_loader.py
@@ -1,312 +1,312 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import numpy as np
import os
from collections import OrderedDict
import pandas as pd
import pickle
import time
import subprocess
import utils.dataloader_utils as dutils
# batch generator tools from https://github.com/MIC-DKFZ/batchgenerators
from batchgenerators.dataloading.data_loader import SlimDataLoaderBase
from batchgenerators.transforms.spatial_transforms import MirrorTransform as Mirror
from batchgenerators.transforms.abstract_transforms import Compose
from batchgenerators.dataloading.multi_threaded_augmenter import MultiThreadedAugmenter
from batchgenerators.dataloading import SingleThreadedAugmenter
from batchgenerators.transforms.spatial_transforms import SpatialTransform
from batchgenerators.transforms.crop_and_pad_transforms import CenterCropTransform
from batchgenerators.transforms.utility_transforms import ConvertSegToBoundingBoxCoordinates
def get_train_generators(cf, logger):
"""
wrapper function for creating the training batch generator pipeline. returns the train/val generators.
selects patients according to cv folds (generated by first run/fold of experiment):
splits the data into n-folds, where 1 split is used for val, 1 split for testing and the rest for training. (inner loop test set)
If cf.hold_out_test_set is True, adds the test split to the training data.
"""
all_data = load_dataset(cf, logger)
all_pids_list = np.unique([v['pid'] for (k, v) in all_data.items()])
assert cf.n_train_val_data <= len(all_pids_list), \
"requested {} train val samples, but dataset only has {} train val samples.".format(
cf.n_train_val_data, len(all_pids_list))
train_pids = all_pids_list[:int(2*cf.n_train_val_data//3)]
val_pids = all_pids_list[int(np.ceil(2*cf.n_train_val_data//3)):cf.n_train_val_data]
train_data = {k: v for (k, v) in all_data.items() if any(p == v['pid'] for p in train_pids)}
val_data = {k: v for (k, v) in all_data.items() if any(p == v['pid'] for p in val_pids)}
logger.info("data set loaded with: {} train / {} val patients".format(len(train_pids), len(val_pids)))
batch_gen = {}
batch_gen['train'] = create_data_gen_pipeline(train_data, cf=cf, do_aug=False)
batch_gen['val_sampling'] = create_data_gen_pipeline(val_data, cf=cf, do_aug=False)
if cf.val_mode == 'val_patient':
batch_gen['val_patient'] = PatientBatchIterator(val_data, cf=cf)
batch_gen['n_val'] = len(val_pids) if cf.max_val_patients is None else min(len(val_pids), cf.max_val_patients)
else:
batch_gen['n_val'] = cf.num_val_batches
return batch_gen
def get_test_generator(cf, logger):
"""
wrapper function for creating the test batch generator pipeline.
selects patients according to cv folds (generated by first run/fold of experiment)
If cf.hold_out_test_set is True, gets the data from an external folder instead.
"""
if cf.hold_out_test_set:
pp_name = cf.pp_test_name
test_ix = None
else:
pp_name = None
with open(os.path.join(cf.exp_dir, 'fold_ids.pickle'), 'rb') as handle:
fold_list = pickle.load(handle)
_, _, test_ix, _ = fold_list[cf.fold]
# warnings.warn('WARNING: using validation set for testing!!!')
test_data = load_dataset(cf, logger, test_ix, pp_data_path=cf.pp_test_data_path, pp_name=pp_name)
logger.info("data set loaded with: {} test patients from {}".format(len(test_data.keys()), cf.pp_test_data_path))
batch_gen = {}
batch_gen['test'] = PatientBatchIterator(test_data, cf=cf)
batch_gen['n_test'] = len(test_data.keys()) if cf.max_test_patients=="all" else \
min(cf.max_test_patients, len(test_data.keys()))
return batch_gen
def load_dataset(cf, logger, subset_ixs=None, pp_data_path=None, pp_name=None):
"""
loads the dataset. if deployed in cloud also copies and unpacks the data to the working directory.
:param subset_ixs: subset indices to be loaded from the dataset. used e.g. for testing to only load the test folds.
:return: data: dictionary with one entry per patient (in this case per patient-breast, since they are treated as
individual images for training) each entry is a dictionary containing respective meta-info as well as paths to the preprocessed
numpy arrays to be loaded during batch-generation
"""
if pp_data_path is None:
pp_data_path = cf.pp_data_path
if pp_name is None:
pp_name = cf.pp_name
if cf.server_env:
copy_data = True
target_dir = os.path.join(cf.data_dest, pp_name)
if not os.path.exists(target_dir):
cf.data_source_dir = pp_data_path
os.makedirs(target_dir)
subprocess.call('rsync -av {} {}'.format(
os.path.join(cf.data_source_dir, cf.input_df_name), os.path.join(target_dir, cf.input_df_name)), shell=True)
logger.info('created target dir and info df at {}'.format(os.path.join(target_dir, cf.input_df_name)))
elif subset_ixs is None:
copy_data = False
pp_data_path = target_dir
p_df = pd.read_pickle(os.path.join(pp_data_path, cf.input_df_name))
if subset_ixs is not None:
subset_pids = [np.unique(p_df.pid.tolist())[ix] for ix in subset_ixs]
p_df = p_df[p_df.pid.isin(subset_pids)]
logger.info('subset: selected {} instances from df'.format(len(p_df)))
if cf.server_env:
if copy_data:
copy_and_unpack_data(logger, p_df.pid.tolist(), cf.fold_dir, cf.data_source_dir, target_dir)
class_targets = p_df['class_id'].tolist()
pids = p_df.pid.tolist()
imgs = [os.path.join(pp_data_path, '{}.npy'.format(pid)) for pid in pids]
segs = [os.path.join(pp_data_path,'{}.npy'.format(pid)) for pid in pids]
data = OrderedDict()
for ix, pid in enumerate(pids):
data[pid] = {'data': imgs[ix], 'seg': segs[ix], 'pid': pid, 'class_target': [class_targets[ix]]}
return data
def create_data_gen_pipeline(patient_data, cf, do_aug=True):
"""
create mutli-threaded train/val/test batch generation and augmentation pipeline.
:param patient_data: dictionary containing one dictionary per patient in the train/test subset.
:param is_training: (optional) whether to perform data augmentation (training) or not (validation/testing)
:return: multithreaded_generator
"""
# create instance of batch generator as first element in pipeline.
data_gen = BatchGenerator(patient_data, batch_size=cf.batch_size, cf=cf)
# add transformations to pipeline.
my_transforms = []
if do_aug:
mirror_transform = Mirror(axes=np.arange(2, cf.dim+2, 1))
my_transforms.append(mirror_transform)
spatial_transform = SpatialTransform(patch_size=cf.patch_size[:cf.dim],
patch_center_dist_from_border=cf.da_kwargs['rand_crop_dist'],
do_elastic_deform=cf.da_kwargs['do_elastic_deform'],
alpha=cf.da_kwargs['alpha'], sigma=cf.da_kwargs['sigma'],
do_rotation=cf.da_kwargs['do_rotation'], angle_x=cf.da_kwargs['angle_x'],
angle_y=cf.da_kwargs['angle_y'], angle_z=cf.da_kwargs['angle_z'],
do_scale=cf.da_kwargs['do_scale'], scale=cf.da_kwargs['scale'],
random_crop=cf.da_kwargs['random_crop'])
my_transforms.append(spatial_transform)
else:
my_transforms.append(CenterCropTransform(crop_size=cf.patch_size[:cf.dim]))
my_transforms.append(ConvertSegToBoundingBoxCoordinates(cf.dim, get_rois_from_seg_flag=False, class_specific_seg_flag=cf.class_specific_seg_flag))
all_transforms = Compose(my_transforms)
# multithreaded_generator = SingleThreadedAugmenter(data_gen, all_transforms)
multithreaded_generator = MultiThreadedAugmenter(data_gen, all_transforms, num_processes=cf.n_workers, seeds=range(cf.n_workers))
return multithreaded_generator
class BatchGenerator(SlimDataLoaderBase):
"""
creates the training/validation batch generator. Samples n_batch_size patients (draws a slice from each patient if 2D)
from the data set while maintaining foreground-class balance. Returned patches are cropped/padded to pre_crop_size.
Actual patch_size is obtained after data augmentation.
:param data: data dictionary as provided by 'load_dataset'.
:param batch_size: number of patients to sample for the batch
- :return dictionary containing the batch data (b, c, x, y, (z)) / seg (b, 1, x, y, (z)) / pids / class_target
+ :return dictionary containing the batch data (b, c, y, x, (z)) / seg (b, 1, y, x, (z)) / pids / class_target
"""
def __init__(self, data, batch_size, cf):
super(BatchGenerator, self).__init__(data, batch_size)
self.cf = cf
def generate_train_batch(self):
batch_data, batch_segs, batch_pids, batch_targets = [], [], [], []
class_targets_list = [v['class_target'] for (k, v) in self._data.items()]
#samples patients towards equilibrium of foreground classes on a roi-level (after randomly sampling the ratio "batch_sample_slack).
batch_ixs = dutils.get_class_balanced_patients(
class_targets_list, self.batch_size, self.cf.head_classes - 1, slack_factor=self.cf.batch_sample_slack)
patients = list(self._data.items())
for b in batch_ixs:
patient = patients[b][1]
all_data = np.load(patient['data'], mmap_mode='r')
data = all_data[0]
seg = all_data[1].astype('uint8')
batch_pids.append(patient['pid'])
batch_targets.append(patient['class_target'])
batch_data.append(data[np.newaxis])
batch_segs.append(seg[np.newaxis])
data = np.array(batch_data)
seg = np.array(batch_segs).astype(np.uint8)
class_target = np.array(batch_targets)
return {'data': data, 'seg': seg, 'pid': batch_pids, 'class_target': class_target}
class PatientBatchIterator(SlimDataLoaderBase):
"""
creates a test generator that iterates over entire given dataset returning 1 patient per batch.
Can be used for monitoring if cf.val_mode = 'patient_val' for a monitoring closer to actualy evaluation (done in 3D),
if willing to accept speed-loss during training.
:return: out_batch: dictionary containing one patient with batch_size = n_3D_patches in 3D or
batch_size = n_2D_patches in 2D .
"""
def __init__(self, data, cf): #threads in augmenter
super(PatientBatchIterator, self).__init__(data, 0)
self.cf = cf
self.patient_ix = 0
self.dataset_pids = [v['pid'] for (k, v) in data.items()]
self.patch_size = cf.patch_size
if len(self.patch_size) == 2:
self.patch_size = self.patch_size + [1]
def generate_train_batch(self):
pid = self.dataset_pids[self.patient_ix]
patient = self._data[pid]
all_data = np.load(patient['data'], mmap_mode='r')
data = all_data[0]
seg = all_data[1].astype('uint8')
batch_class_targets = np.array([patient['class_target']])
out_data = data[None, None]
out_seg = seg[None, None]
#print('check patient data loader', out_data.shape, out_seg.shape)
batch_2D = {'data': out_data, 'seg': out_seg, 'class_target': batch_class_targets, 'pid': pid}
converter = ConvertSegToBoundingBoxCoordinates(dim=2, get_rois_from_seg_flag=False, class_specific_seg_flag=self.cf.class_specific_seg_flag)
batch_2D = converter(**batch_2D)
batch_2D.update({'patient_bb_target': batch_2D['bb_target'],
- 'patient_roi_labels': batch_2D['roi_labels'],
+ 'patient_roi_labels': batch_2D['class_target'],
'original_img_shape': out_data.shape})
self.patient_ix += 1
if self.patient_ix == len(self.dataset_pids):
self.patient_ix = 0
return batch_2D
def copy_and_unpack_data(logger, pids, fold_dir, source_dir, target_dir):
start_time = time.time()
with open(os.path.join(fold_dir, 'file_list.txt'), 'w') as handle:
for pid in pids:
handle.write('{}.npy\n'.format(pid))
subprocess.call('rsync -ahv --files-from {} {} {}'.format(os.path.join(fold_dir, 'file_list.txt'),
source_dir, target_dir), shell=True)
# dutils.unpack_dataset(target_dir)
copied_files = os.listdir(target_dir)
logger.info("copying data set finished : {} files in target dir: {}. took {} sec".format(
len(copied_files), target_dir, np.round(time.time() - start_time, 0)))
if __name__=="__main__":
import utils.exp_utils as utils
total_stime = time.time()
cf_file = utils.import_module("cf", "configs.py")
cf = cf_file.configs()
logger = utils.get_logger("dev")
batch_gen = get_train_generators(cf, logger)
train_batch = next(batch_gen["train"])
pids = []
total = 100
for i in range(total):
print("\r producing batch {}/{}.".format(i, total), end="", flush=True)
train_batch = next(batch_gen["train"])
pids.append(train_batch["pid"])
print()
mins, secs = divmod((time.time() - total_stime), 60)
h, mins = divmod(mins, 60)
t = "{:d}h:{:02d}m:{:02d}s".format(int(h), int(mins), int(secs))
print("{} total runtime: {}".format(os.path.split(__file__)[1], t))
\ No newline at end of file
diff --git a/experiments/tutorial.md b/experiments/tutorial.md
index 669ab06..5f114b1 100644
--- a/experiments/tutorial.md
+++ b/experiments/tutorial.md
@@ -1,100 +1,99 @@
# Tutorial
##### for Including a Dataset into the Framework
## Introduction
This tutorial aims at providing a muster routine for including a new dataset into the framework in order to
use the included models and algorithms with it.\
The tutorial and toy dataset (under `toy_exp`) are in 2D, yet the switch to 3D is simply made by providing 3D data and proceeding
analogically, as can be seen from the provided LIDC scripts (under `lidc_exp`).
Datasets in the framework are set up under `medicaldetectiontoolkit/experiments/` and
require three fundamental scripts:
1. A **preprocessing** script that performs one-time routines on your raw data bringing it into a suitable, easily usable
format.
2. A **data-loading** script (required name `data_loader.py`) that efficiently assembles the preprocessed data into
network-processable batches.
3. A **configs** file (`configs.py`) which specifies all settings, from data loading to network architecture.
This file is automatically complemented by `default_settings.py` which holds default and dataset-independent settings.
## Preprocessing
This script (`generate_toys.py` in case of the provided toy dataset, `preprocessing.py` in case of LIDC) is required
to bring your raw data into an easily usable format. We recommend, you put all one-time processes (like normalization,
resampling, cropping, type conversions) into this script in order to avoid the need for repetitive actions during
data loading.\
For the framework usage, we follow a simple workload separation scheme, where network computations
are performed on the GPU while data loading and augmentations are performed on the CPU. Hence, the framework requires
numpy arrays (`.npy`) as input to the networks, therefore your preprocessed data (images and segmentations) should
already be in that format. In terms of data dimensions, we follow the scheme: (y, x (,z)), meaning coronal, sagittal,
and axial dimensions, respectively.
Class labels for the Regions of Interest (RoIs) need to be provided as lists per data sample.
If you have segmenation data, you may use the [batchgenerators](https://github.com/MIC-DKFZ/batchgenerators) transform
ConvertSegToBoundingBoxCoordinates to generate bounding boxes from your segmentations. In that case, the order of the
class labels in the list needs to correspond to the RoI labels in the segmentation.\
Example: An image (2D or 3D) has two RoIs, one of class 1, the
other of class 0. In your segmentation, every pixel is 0 (bg), except for the area marking class 1, which has value 1,
and the area of class 0, which has value 2. Your list of class labels for this sample should be `[1, 0]`. I.e.,
the index of the RoI's class label in the sample's label list corresponds to its marking in the segmentation shifted
by -1.\
If you do not have segmentations (only models Faster R-CNN and RetinaNet can be used), you can directly provide bounding
boxes. In that case, RoIs are simply identified by their indices in the lists: class label list `[cl_of_roi_a, cl_of_roi_b]`
corresponds to bbox list `[coords_of_roi_a, coords_of_roi_b]`.
Please store all your light-weight information (patient id, class targets, (relative) paths or identifiers for data and seg) about the
preprocessed data set in a pandas dataframe, say `info_df.pkl`.
## Data Loading
The goal of `data_loader.py` is to sample or iterate, load into CPU RAM, assemble, and eventually augment the preprocessed data.\
The framework requires the data loader to provide at least a function `get_train_generators`, which yields a dict
holding a train-data loader under key `"train"` and validation loader under `"val_sampling"` or `"val_patient"`,
analogically for `get_test_generator` with `"test"`.\
We recommend you closely follow our structure as in the provided datasets, which includes a data loader suitable for
sampling single patches or parts of the whole patient data with focus on class equilibrium (BatchGenerator,
used in training and optionally validation) and a PatientIterator which is intended for test and optionally valdiation and
iterates through all patients one by one, not discarding
any parts of the patient image. In detail, the structure is as follows.
Data loading is performed with the help of the batchgenerators package. Starting from farthest to closest to the
preprocessed data, the data loader contains:
1. Method `get_train_generators` which is called by the execution script and in the end provides train and val data loaders.
Same goes for `get_test_generator` for the test loader.
2. Method `load_dataset` which reads the `info_df.pkl` and provides a dictionary holding, per patient id, paths
to images and segmentations, and light-weight info like class targets.
3. Method `create_data_gen_pipeline` which initiates the train data loader (instance of class BatchGenerator),
assembles the chosen data-augmentation procedures and passes the BatchGenerator into a MultiThreadedAugmenter (MTA). The MTA
is a wrapper that manages multi-threaded loading (and augmentation).
4. Class BatchGenerator. This data loader is used for sampling, e.g., according to the scheme described in
`utils/dataloader_utils.get_class_balanced_patients`. It needs to implement a `__next__` method providing the batch;
the batch is a dictionary with (at least) keys: `"data"`, `"pid"`, `"class_target"` (as well as `"seg"` if using segmentations).
- `"data"` needs to hold your image (2D or 3D) as a numpy array with dimensions: (b, c, y, x(, z)), where b is the
batch dimension (b = batch size), c the channel dimension (if you have multi-modal data c > 1), y, x, z are
- the spatial dimensions; z is omitted in case of 2D data.
+ the spatial dimensions; z is omitted in case of 2D data. Since the batchgenerators package uses shape convention (x,y,z),
+ please make sure you switch augmentation settings explicitly affecting x and y (like rotation angle) accordingly.
- `"seg"` has the same format as `"data"`, except that its channel dimension has always size c = 1.
- `"pid"` is a list of patient or sample identifiers, one per sample, i.e., shape (b,).
- `"class_target"` which holds, as mentioned in preprocessing, class labels for the RoIs. It's a list of length b, holding
itself lists of varying lengths n_rois(sample).
**Note**: the above description only applies if you use ConvertSegToBoundingBoxCoordinates. Class targets after batch
generation need to make room for a background class (network heads need to be able to predict class 0 = bg). Since,
in preprocessing, we started classes at id 0, we now need to shift them by +1. This is done automatically inside
- ConvertSegToBoundingBoxCoordinates. That transform also renames `"class_target"` to `"roi_labels"`, which is the label
- required by the rest of the framework. In case you do not use that transform, please shift and rename the labels
- in your BatchGenerator.
+ ConvertSegToBoundingBoxCoordinates. In case you do not use that transform, please shift the labels in your BatchGenerator.
5. Class PatientIterator. This data loader is intended for testing and validation. It needs to provide the same output as
above BatchGenerator, however, initial batch size is always limited to one (one patient). Output batch size may vary
if patching is applied. Please refer to the LIDC PatientIterator
-to see how to include patching. Note that this Iterator is not supposed to go through the MTA, transforms (mainly
+to see how to include patching. Note that this Iterator is _not_ supposed to go through the MTA, transforms (mainly
ConvertSegToBoundingBoxCoordinates) therefore need to be applied within this class directly.
## Configs
The current work flow is intended for running multiple experiments with the same dataset but different configs. This is
done by setting the desired values in `configs.py` in the data set's source directory, then creating an experiment
via the execution script (`exec.py`, modes "create_exp" or "train" or "train_test"), which copies a snapshot of configs,
data loader, default configs, and selected model to the provided experiment directory.
`configs.py` introduces class `configs`, which, when instantiated, inherits settings in `default_configs.py` and adds
model-specific settings to itself. Aside from setting all the right input/output paths, you can tune almost anything, from
network architecture to data-loading settings to train and test routine settings.\
Furthermore, throughout the whole framework, you have the option to include server-environment specific settings by passing
argument `--server_env` to the exec script. E.g., in the configs, we use this flag, to overwrite local paths by the
paths we use on our GPU cluster.
\ No newline at end of file
diff --git a/models/detection_unet.py b/models/detection_unet.py
index db1025a..678c65c 100644
--- a/models/detection_unet.py
+++ b/models/detection_unet.py
@@ -1,214 +1,215 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""
Unet-like Backbone architecture, with non-parametric heuristics for box detection on semantic segmentation outputs.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from scipy.ndimage.measurements import label as lb
import numpy as np
import utils.exp_utils as utils
import utils.model_utils as mutils
class net(nn.Module):
def __init__(self, cf, logger):
super(net, self).__init__()
self.cf = cf
self.logger = logger
backbone = utils.import_module('bbone', cf.backbone_path)
conv = mutils.NDConvGenerator(cf.dim)
# set operate_stride1=True to generate a unet-like FPN.)
self.fpn = backbone.FPN(cf, conv, operate_stride1=True).cuda()
self.conv_final = conv(cf.end_filts, cf.num_seg_classes, ks=1, pad=0, norm=cf.norm, relu=None)
if self.cf.weight_init is not None:
logger.info("using pytorch weight init of type {}".format(self.cf.weight_init))
mutils.initialize_weights(self)
else:
logger.info("using default pytorch weight init")
def forward(self, x):
"""
forward pass of network.
:param x: input image. shape (b, c, y, x, (z))
:return: seg_logits: shape (b, n_classes, y, x, (z))
:return: out_box_coords: list over n_classes. elements are arrays(b, n_rois, (y1, x1, y2, x2, (z1), (z2)))
:return: out_max_scores: list over n_classes. elements are arrays(b, n_rois)
"""
out_features = self.fpn(x)[0]
seg_logits = self.conv_final(out_features)
out_box_coords, out_max_scores = [], []
smax = F.softmax(seg_logits, dim=1).detach().cpu().data.numpy()
for cl in range(1, len(self.cf.class_dict.keys()) + 1):
max_scores = [[] for _ in range(x.shape[0])]
hard_mask = np.copy(smax).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 = 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.
for bix, broi in enumerate(rois):
for nix, nroi in enumerate(broi):
component_score = np.max(smax[bix, cl][nroi > 0]) if self.cf.aggregation_operation == 'max' \
else np.median(smax[bix, cl][nroi > 0])
max_scores[bix].append(component_score)
out_box_coords.append(box_coords)
out_max_scores.append(max_scores)
return seg_logits, out_box_coords, out_max_scores
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]
'monitor_values': dict of values to be monitored.
"""
img = batch['data']
seg = batch['seg']
var_img = torch.FloatTensor(img).cuda()
var_seg = torch.FloatTensor(seg).cuda().long()
var_seg_ohe = torch.FloatTensor(mutils.get_one_hot_encoding(seg, self.cf.num_seg_classes)).cuda()
results_dict = {}
seg_logits, box_coords, max_scores = self.forward(var_img)
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(max_scores[cix][bix])):
if max_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': max_scores[cix][bix][rix],
'box_pred_class_id': cix + 1, # add 0 for background.
'box_type': 'det'})
-
-
+ if "roi_labels" in batch.keys():
+ raise Exception("Key for roi-wise class targets changed in v0.1.0 from 'roi_labels' to 'class_target'.\n"
+ "If you use DKFZ's batchgenerators, please make sure you run version >= 0.20.1.")
for bix in range(img.shape[0]):
for tix in range(len(batch['bb_target'][bix])):
results_dict['boxes'][bix].append({'box_coords': batch['bb_target'][bix][tix],
- 'box_label': batch['roi_labels'][bix][tix],
+ 'box_label': batch['class_target'][bix][tix],
'box_type': 'gt'})
# compute segmentation loss as either weighted cross entropy, dice loss, or the sum of both.
loss = torch.FloatTensor([0]).cuda()
if self.cf.seg_loss_mode == 'dice' or self.cf.seg_loss_mode == 'dice_wce':
loss += 1 - mutils.batch_dice(F.softmax(seg_logits, dim=1), var_seg_ohe,
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, var_seg[:, 0], weight=torch.tensor(self.cf.wce_weights).float().cuda())
results_dict['seg_preds'] = np.argmax(F.softmax(seg_logits, 1).cpu().data.numpy(), 1)[:, np.newaxis]
results_dict['torch_loss'] = loss
results_dict['monitor_values'] = {'loss': loss.item()}
results_dict['logger_string'] = "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 = batch['data']
var_img = torch.FloatTensor(img).cuda()
seg_logits, box_coords, max_scores = self.forward(var_img)
results_dict = {}
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(max_scores[cix][bix])):
if max_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': max_scores[cix][bix][rix],
'box_pred_class_id': cix + 1, # add 0 for background.
'box_type': 'det'})
results_dict['seg_preds'] = np.argmax(F.softmax(seg_logits, 1).cpu().data.numpy(), 1)[:, np.newaxis].astype('uint8')
return results_dict
def get_coords(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))
"""
binary_mask = binary_mask.astype('uint8')
batch_coords = []
batch_components = []
for ix, b in enumerate(binary_mask):
clusters, n_cands = lb(b) # peforms connected component analysis.
uniques, counts = np.unique(clusters, return_counts=True)
# only keep n_components largest components.
keep_uniques = uniques[1:][np.argsort(counts[1:])[::-1]][:n_components]
# separate clusters and concat.
p_components = np.array([(clusters == ii) * 1 for ii in keep_uniques])
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
diff --git a/models/mrcnn.py b/models/mrcnn.py
index 1dc4434..02490cf 100644
--- a/models/mrcnn.py
+++ b/models/mrcnn.py
@@ -1,1181 +1,1184 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""
Parts are based on https://github.com/multimodallearning/pytorch-mask-rcnn
published under MIT license.
"""
import sys
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils
sys.path.append("..")
import utils.model_utils as mutils
import utils.exp_utils as utils
from custom_extensions.nms import nms
from custom_extensions.roi_align import roi_align
############################################################
# Networks on top of backbone
############################################################
class RPN(nn.Module):
"""
Region Proposal Network.
"""
def __init__(self, cf, conv):
super(RPN, self).__init__()
self.dim = conv.dim
self.conv_shared = conv(cf.end_filts, cf.n_rpn_features, ks=3, stride=cf.rpn_anchor_stride, pad=1, relu=cf.relu)
self.conv_class = conv(cf.n_rpn_features, 2 * len(cf.rpn_anchor_ratios), ks=1, stride=1, relu=None)
self.conv_bbox = conv(cf.n_rpn_features, 2 * self.dim * len(cf.rpn_anchor_ratios), ks=1, stride=1, relu=None)
def forward(self, x):
"""
:param x: input feature maps (b, in_channels, y, x, (z))
:return: rpn_class_logits (b, 2, n_anchors)
:return: rpn_probs_logits (b, 2, n_anchors)
:return: rpn_bbox (b, 2 * dim, n_anchors)
"""
# Shared convolutional base of the RPN.
x = self.conv_shared(x)
# Anchor Score. (batch, anchors per location * 2, y, x, (z)).
rpn_class_logits = self.conv_class(x)
# Reshape to (batch, 2, anchors)
axes = (0, 2, 3, 1) if self.dim == 2 else (0, 2, 3, 4, 1)
rpn_class_logits = rpn_class_logits.permute(*axes)
rpn_class_logits = rpn_class_logits.contiguous()
rpn_class_logits = rpn_class_logits.view(x.size()[0], -1, 2)
# Softmax on last dimension (fg vs. bg).
rpn_probs = F.softmax(rpn_class_logits, dim=2)
# Bounding box refinement. (batch, anchors_per_location * (y, x, (z), log(h), log(w), (log(d)), y, x, (z))
rpn_bbox = self.conv_bbox(x)
# Reshape to (batch, 2*dim, anchors)
rpn_bbox = rpn_bbox.permute(*axes)
rpn_bbox = rpn_bbox.contiguous()
rpn_bbox = rpn_bbox.view(x.size()[0], -1, self.dim * 2)
return [rpn_class_logits, rpn_probs, rpn_bbox]
class Classifier(nn.Module):
"""
Head network for classification and bounding box refinement. Performs RoiAlign, processes resulting features through a
shared convolutional base and finally branches off the classifier- and regression head.
"""
def __init__(self, cf, conv):
super(Classifier, self).__init__()
self.dim = conv.dim
self.in_channels = cf.end_filts
self.pool_size = cf.pool_size
self.pyramid_levels = cf.pyramid_levels
# instance_norm does not work with spatial dims (1, 1, (1))
norm = cf.norm if cf.norm != 'instance_norm' else None
self.conv1 = conv(cf.end_filts, cf.end_filts * 4, ks=self.pool_size, stride=1, norm=norm, relu=cf.relu)
self.conv2 = conv(cf.end_filts * 4, cf.end_filts * 4, ks=1, stride=1, norm=norm, relu=cf.relu)
self.linear_class = nn.Linear(cf.end_filts * 4, cf.head_classes)
self.linear_bbox = nn.Linear(cf.end_filts * 4, cf.head_classes * 2 * self.dim)
def forward(self, x, rois):
"""
:param x: input feature maps (b, in_channels, y, x, (z))
:param rois: normalized box coordinates as proposed by the RPN to be forwarded through
the second stage (n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ix). Proposals of all batch elements
have been merged to one vector, while the origin info has been stored for re-allocation.
:return: mrcnn_class_logits (n_proposals, n_head_classes)
:return: mrcnn_bbox (n_proposals, n_head_classes, 2 * dim) predicted corrections to be applied to proposals for refinement.
"""
x = pyramid_roi_align(x, rois, self.pool_size, self.pyramid_levels, self.dim)
x = self.conv1(x)
x = self.conv2(x)
x = x.view(-1, self.in_channels * 4)
mrcnn_class_logits = self.linear_class(x)
mrcnn_bbox = self.linear_bbox(x)
mrcnn_bbox = mrcnn_bbox.view(mrcnn_bbox.size()[0], -1, self.dim * 2)
return [mrcnn_class_logits, mrcnn_bbox]
class Mask(nn.Module):
"""
Head network for proposal-based mask segmentation. Performs RoiAlign, some convolutions and applies sigmoid on the
output logits to allow for overlapping classes.
"""
def __init__(self, cf, conv):
super(Mask, self).__init__()
self.pool_size = cf.mask_pool_size
self.pyramid_levels = cf.pyramid_levels
self.dim = conv.dim
self.conv1 = conv(cf.end_filts, cf.end_filts, ks=3, stride=1, pad=1, norm=cf.norm, relu=cf.relu)
self.conv2 = conv(cf.end_filts, cf.end_filts, ks=3, stride=1, pad=1, norm=cf.norm, relu=cf.relu)
self.conv3 = conv(cf.end_filts, cf.end_filts, ks=3, stride=1, pad=1, norm=cf.norm, relu=cf.relu)
self.conv4 = conv(cf.end_filts, cf.end_filts, ks=3, stride=1, pad=1, norm=cf.norm, relu=cf.relu)
if conv.dim == 2:
self.deconv = nn.ConvTranspose2d(cf.end_filts, cf.end_filts, kernel_size=2, stride=2)
else:
self.deconv = nn.ConvTranspose3d(cf.end_filts, cf.end_filts, kernel_size=2, stride=2)
self.relu = nn.ReLU(inplace=True) if cf.relu == 'relu' else nn.LeakyReLU(inplace=True)
self.conv5 = conv(cf.end_filts, cf.head_classes, ks=1, stride=1, relu=None)
self.sigmoid = nn.Sigmoid()
def forward(self, x, rois):
"""
:param x: input feature maps (b, in_channels, y, x, (z))
:param rois: normalized box coordinates as proposed by the RPN to be forwarded through
the second stage (n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ix). Proposals of all batch elements
have been merged to one vector, while the origin info has been stored for re-allocation.
:return: x: masks (n_sampled_proposals (n_detections in inference), n_classes, y, x, (z))
"""
x = pyramid_roi_align(x, rois, self.pool_size, self.pyramid_levels, self.dim)
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)
x = self.conv4(x)
x = self.relu(self.deconv(x))
x = self.conv5(x)
x = self.sigmoid(x)
return x
############################################################
# Loss Functions
############################################################
def compute_rpn_class_loss(rpn_match, rpn_class_logits, shem_poolsize):
"""
:param rpn_match: (n_anchors). [-1, 0, 1] for negative, neutral, and positive matched anchors.
:param rpn_class_logits: (n_anchors, 2). logits from RPN classifier.
:param shem_poolsize: int. factor of top-k candidates to draw from per negative sample
(stochastic-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.
"""
# filter out neutral anchors.
pos_indices = torch.nonzero(rpn_match == 1)
neg_indices = torch.nonzero(rpn_match == -1)
# loss for positive samples
if 0 not in pos_indices.size():
pos_indices = pos_indices.squeeze(1)
roi_logits_pos = rpn_class_logits[pos_indices]
pos_loss = F.cross_entropy(roi_logits_pos, torch.LongTensor([1] * pos_indices.shape[0]).cuda())
else:
pos_loss = torch.FloatTensor([0]).cuda()
# loss for negative samples: draw hard negative examples (SHEM)
# that match the number of positive samples, but at least 1.
if 0 not in neg_indices.size():
neg_indices = neg_indices.squeeze(1)
roi_logits_neg = rpn_class_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())
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_rpn_bbox_loss(rpn_target_deltas, rpn_pred_deltas, rpn_match):
"""
:param rpn_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 rpn_pred_deltas: predicted deltas from RPN. (b, n_anchors, (dy, dx, (dz), log(dh), log(dw), (log(dd))))
:param rpn_match: (n_anchors). [-1, 0, 1] for negative, neutral, and positive matched anchors.
:return: loss: torch 1D tensor.
"""
if 0 not in torch.nonzero(rpn_match == 1).size():
indices = torch.nonzero(rpn_match == 1).squeeze(1)
# Pick bbox deltas that contribute to the loss
rpn_pred_deltas = rpn_pred_deltas[indices]
# Trim target bounding box deltas to the same length as rpn_bbox.
target_deltas = rpn_target_deltas[:rpn_pred_deltas.size()[0], :]
# Smooth L1 loss
loss = F.smooth_l1_loss(rpn_pred_deltas, target_deltas)
else:
loss = torch.FloatTensor([0]).cuda()
return loss
def compute_mrcnn_class_loss(target_class_ids, pred_class_logits):
"""
:param target_class_ids: (n_sampled_rois) batch dimension was merged into roi dimension.
:param pred_class_logits: (n_sampled_rois, n_classes)
:return: loss: torch 1D tensor.
"""
if 0 not in target_class_ids.size():
loss = F.cross_entropy(pred_class_logits, target_class_ids.long())
else:
loss = torch.FloatTensor([0.]).cuda()
return loss
def compute_mrcnn_bbox_loss(mrcnn_target_deltas, mrcnn_pred_deltas, target_class_ids):
"""
:param mrcnn_target_deltas: (n_sampled_rois, (dy, dx, (dz), log(dh), log(dw), (log(dh)))
:param mrcnn_pred_deltas: (n_sampled_rois, n_classes, (dy, dx, (dz), log(dh), log(dw), (log(dh)))
:param target_class_ids: (n_sampled_rois)
:return: loss: torch 1D tensor.
"""
if 0 not in torch.nonzero(target_class_ids > 0).size():
positive_roi_ix = torch.nonzero(target_class_ids > 0)[:, 0]
positive_roi_class_ids = target_class_ids[positive_roi_ix].long()
target_bbox = mrcnn_target_deltas[positive_roi_ix, :].detach()
pred_bbox = mrcnn_pred_deltas[positive_roi_ix, positive_roi_class_ids, :]
loss = F.smooth_l1_loss(pred_bbox, target_bbox)
else:
loss = torch.FloatTensor([0]).cuda()
return loss
def compute_mrcnn_mask_loss(target_masks, pred_masks, target_class_ids):
"""
:param target_masks: (n_sampled_rois, y, x, (z)) A float32 tensor of values 0 or 1. Uses zero padding to fill array.
:param pred_masks: (n_sampled_rois, n_classes, y, x, (z)) float32 tensor with values between [0, 1].
:param target_class_ids: (n_sampled_rois)
:return: loss: torch 1D tensor.
"""
if 0 not in torch.nonzero(target_class_ids > 0).size():
# Only positive ROIs contribute to the loss. And only
# the class specific mask of each ROI.
positive_ix = torch.nonzero(target_class_ids > 0)[:, 0]
positive_class_ids = target_class_ids[positive_ix].long()
y_true = target_masks[positive_ix, :, :].detach()
y_pred = pred_masks[positive_ix, positive_class_ids, :, :]
loss = F.binary_cross_entropy(y_pred, y_true)
else:
loss = torch.FloatTensor([0]).cuda()
return loss
############################################################
# Helper Layers
############################################################
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 = mutils.apply_box_deltas_2D(anchors[order, :], deltas)
boxes = mutils.clip_boxes_2D(boxes, cf.window)
else:
boxes = mutils.apply_box_deltas_3D(anchors[order, :], deltas)
boxes = mutils.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(normalized_boxes.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 + 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 detection_target_layer(batch_proposals, batch_mrcnn_class_scores, batch_gt_class_ids, batch_gt_boxes, batch_gt_masks, cf):
"""
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 are 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".)
: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 batch_mrcnn_class_scores: (n_proposals, n_classes)
:param batch_gt_class_ids: list over batch elements. Each element is a list over the corresponding roi target labels.
: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)
: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
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 = []
std_dev = torch.from_numpy(cf.bbox_std_dev).float().cuda()
# loop over batch and get positive and negative sample rois.
for b in range(len(batch_gt_class_ids)):
gt_class_ids = torch.from_numpy(batch_gt_class_ids[b]).int().cuda()
gt_masks = torch.from_numpy(batch_gt_masks[b]).float().cuda()
if np.any(batch_gt_class_ids[b] > 0): # skip roi selection for no gt images.
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 0 not 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 = mutils.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 = mutils.bbox_overlaps_3D(proposals, gt_boxes)
# Determine postive 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 0 not 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]
# Compute bbox refinement targets for positive ROIs
deltas = mutils.box_refinement(positive_rois, roi_gt_boxes)
deltas /= std_dev
# Assign positive ROIs to GT masks
roi_masks = gt_masks[roi_gt_box_assignment]
assert roi_masks.shape[1] == 1, "desired to have more than one channel in gt masks?"
# Compute mask targets
boxes = positive_rois
box_ids = torch.arange(roi_masks.shape[0]).cuda().unsqueeze(1).float()
if len(cf.mask_shape) == 2:
# need to remap normalized box coordinates to unnormalized mask coordinates.
y_exp, x_exp = roi_masks.shape[2:] # exp = expansion
boxes.mul_(torch.tensor([y_exp, x_exp, y_exp, x_exp], dtype=torch.float32).cuda())
masks = roi_align.roi_align_2d(roi_masks, torch.cat((box_ids, boxes), dim=1), cf.mask_shape)
else:
y_exp, x_exp, z_exp = roi_masks.shape[2:] # exp = expansion
boxes.mul_(torch.tensor([y_exp, x_exp, y_exp, x_exp, z_exp, z_exp], dtype=torch.float32).cuda())
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)
positive_count += positive_samples
else:
positive_samples = 0
# Negative ROIs. Add enough to maintain positive:negative ratio, but at least 1. Sample via SHEM.
if 0 not 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_probs_neg = batch_mrcnn_class_scores[batch_element_indices[negative_indices]]
raw_sampled_indices = mutils.shem(roi_probs_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)
# 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).int().cuda()
target_class_ids = torch.cat([target_class_ids, zeros], 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)
elif positive_count > 0:
sample_indices = torch.cat(sample_positive_indices)
elif negative_count > 0:
sample_indices = torch.cat(sample_negative_indices)
zeros = torch.zeros(negative_count).int().cuda()
target_class_ids = zeros
zeros = torch.zeros(negative_count, cf.dim * 2).cuda()
target_deltas = zeros
zeros = torch.zeros(negative_count, *cf.mask_shape).cuda()
target_masks = zeros
else:
sample_indices = torch.LongTensor().cuda()
target_class_ids = torch.IntTensor().cuda()
target_deltas = torch.FloatTensor().cuda()
target_masks = torch.FloatTensor().cuda()
return sample_indices, target_class_ids, target_deltas, target_masks
############################################################
# Output Handler
############################################################
# def refine_detections(rois, probs, deltas, batch_ixs, cf):
# """
# Refine classified proposals, filter overlaps and return final detections.
#
# :param rois: (n_proposals, 2 * dim) normalized boxes as proposed by RPN. n_proposals = batch_size * POST_NMS_ROIS
# :param probs: (n_proposals, n_classes) softmax probabilities for all rois as predicted by mrcnn classifier.
# :param deltas: (n_proposals, n_classes, 2 * dim) box refinement deltas as predicted by mrcnn bbox regressor.
# :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))
# """
# # 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()
#
# rois = rois.repeat(fg_classes, 1)
# probs = probs.repeat(fg_classes, 1)
# deltas = deltas.repeat(fg_classes, 1, 1)
# batch_ixs = batch_ixs.repeat(fg_classes)
#
# # get class-specific scores and bounding box deltas
# idx = torch.arange(class_ids.size()[0]).long().cuda()
# class_scores = probs[idx, class_ids]
# deltas_specific = deltas[idx, class_ids]
# batch_ixs = batch_ixs[idx]
#
# # 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(rois, deltas_specific * std_dev) * scale if cf.dim == 2 else \
# mutils.apply_box_deltas_3D(rois, deltas_specific * 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)
# refined_rois = torch.round(refined_rois)
#
# # filter out low confidence boxes
# keep = idx
# keep_bool = (class_scores >= cf.model_min_confidence)
# if 0 not 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(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, :]
#
# 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)
#
# # 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 mutils.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 mutils.unique1d(torch.cat((batch_keep, b_keep)))
#
# keep = batch_keep
#
# else:
# keep = torch.tensor([0]).long().cuda()
#
# # arrange output
# result = torch.cat((refined_rois[keep],
# batch_ixs[keep].unsqueeze(1),
# class_ids[keep].unsqueeze(1).float(),
# class_scores[keep].unsqueeze(1)), dim=1)
#
# return result
def refine_detections(cf, batch_ixs, rois, deltas, scores):
"""
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.
: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)
# 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]
# 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(rois, deltas_specific * std_dev) * scale if cf.dim == 2 else \
mutils.apply_box_deltas_3D(rois, deltas_specific * std_dev) * scale
# round and cast to int since we're dealing with pixels now
refined_rois = mutils.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(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, :]
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 mutils.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 mutils.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)]
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 get_results(cf, img_shape, detections, detection_masks, box_results_list=None, return_masks=True):
"""
Restores batch dimension of merged detections, unmolds detections, creates and fills results dict.
:param img_shape:
:param detections: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score)
:param detection_masks: (n_final_detections, n_classes, y, x, (z)) raw molded masks as returned by mask-head.
:param box_results_list: None or list of output boxes for monitoring/plotting.
each element is a list of boxes per batch element.
:param return_masks: boolean. If True, full resolution masks are returned for all proposals (speed trade-off).
: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()
if cf.dim == 2:
detection_masks = detection_masks.permute(0, 2, 3, 1).cpu().data.numpy()
else:
detection_masks = detection_masks.permute(0, 2, 3, 4, 1).cpu().data.numpy()
# restore batch dimension of merged detections using the batch_ix info.
batch_ixs = detections[:, cf.dim*2]
detections = [detections[batch_ixs == ix] for ix in range(img_shape[0])]
mrcnn_mask = [detection_masks[batch_ixs == ix] for ix in range(img_shape[0])]
# for test_forward, where no previous list exists.
if box_results_list is None:
box_results_list = [[] for _ in range(img_shape[0])]
seg_preds = []
# loop over batch and unmold detections.
for ix in range(img_shape[0]):
if 0 not in detections[ix].shape:
boxes = detections[ix][:, :2 * cf.dim].astype(np.int32)
class_ids = detections[ix][:, 2 * cf.dim + 1].astype(np.int32)
scores = detections[ix][:, 2 * cf.dim + 2]
masks = mrcnn_mask[ix][np.arange(boxes.shape[0]), ..., class_ids]
# Filter out detections with zero area. Often only happens in early
# stages of training when the network weights are still a bit random.
if cf.dim == 2:
exclude_ix = np.where((boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) <= 0)[0]
else:
exclude_ix = np.where(
(boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 5] - boxes[:, 4]) <= 0)[0]
if exclude_ix.shape[0] > 0:
boxes = np.delete(boxes, exclude_ix, axis=0)
class_ids = np.delete(class_ids, exclude_ix, axis=0)
scores = np.delete(scores, exclude_ix, axis=0)
masks = np.delete(masks, exclude_ix, axis=0)
# Resize masks to original image size and set boundary threshold.
full_masks = []
permuted_image_shape = list(img_shape[2:]) + [img_shape[1]]
if return_masks:
for i in range(masks.shape[0]):
# Convert neural network mask to full size mask.
full_masks.append(mutils.unmold_mask_2D(masks[i], boxes[i], permuted_image_shape)
if cf.dim == 2 else mutils.unmold_mask_3D(masks[i], boxes[i], permuted_image_shape))
# if masks are returned, take max over binary full masks of all predictions in this image.
# right now only binary masks for plotting/monitoring. for instance segmentation return all proposal masks.
final_masks = np.max(np.array(full_masks), 0) if len(full_masks) > 0 else np.zeros(
(*permuted_image_shape[:-1],))
# add final predictions to results.
if 0 not in boxes.shape:
for ix2, score in enumerate(scores):
box_results_list[ix].append({'box_coords': boxes[ix2], 'box_score': score,
'box_type': 'det', 'box_pred_class_id': class_ids[ix2]})
else:
# pad with zero dummy masks.
final_masks = np.zeros(img_shape[2:])
seg_preds.append(final_masks)
# create and fill results dictionary.
results_dict = {'boxes': box_results_list,
'seg_preds': np.round(np.array(seg_preds))[:, np.newaxis].astype('uint8')}
return results_dict
############################################################
# Mask R-CNN Class
############################################################
class net(nn.Module):
def __init__(self, cf, logger):
super(net, self).__init__()
self.cf = cf
self.logger = logger
self.build()
if self.cf.weight_init is not None:
logger.info("using pytorch weight init of type {}".format(self.cf.weight_init))
mutils.initialize_weights(self)
else:
logger.info("using default pytorch weight init")
def build(self):
"""Build Mask R-CNN architecture."""
# Image size must be dividable by 2 multiple times.
h, w = self.cf.patch_size[:2]
if h / 2**5 != int(h / 2**5) or w / 2**5 != int(w / 2**5):
raise Exception("Image size must be dividable by 2 at least 5 times "
"to avoid fractions when downscaling and upscaling."
"For example, use 256, 320, 384, 448, 512, ... etc. ")
if len(self.cf.patch_size) == 3:
d = self.cf.patch_size[2]
if d / 2**3 != int(d / 2**3):
raise Exception("Image z dimension must be dividable by 2 at least 3 times "
"to avoid fractions when downscaling and upscaling.")
# instanciate abstract multi dimensional conv class and backbone class.
conv = mutils.NDConvGenerator(self.cf.dim)
backbone = utils.import_module('bbone', self.cf.backbone_path)
# build Anchors, FPN, RPN, Classifier / Bbox-Regressor -head, Mask-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)
self.rpn = RPN(self.cf, conv)
self.classifier = Classifier(self.cf, conv)
self.mask = Mask(self.cf, conv)
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.
- data_dict['roi_masks']: (b, n(b), 1, h(n), w(n) (z(n))) list like batch['roi_labels'] but with
+ data_dict['roi_masks']: (b, n(b), 1, h(n), w(n) (z(n))) list like batch['class_target'] but with
arrays (masks) inplace of integers. n == number of rois per this batch element.
:return: results_dict: dictionary with keys:
'boxes': list over batch elements. each batch element is a list of boxes. each box is a dictionary:
[[{box_0}, ... {box_n}], [{box_0}, ... {box_n}], ...]
'seg_preds': pixel-wise class predictions (b, 1, y, x, (z)) with values [0, n_classes].
'monitor_values': dict of values to be monitored.
"""
img = batch['data']
- gt_class_ids = batch['roi_labels']
+ if "roi_labels" in batch.keys():
+ raise Exception("Key for roi-wise class targets changed in v0.1.0 from 'roi_labels' to 'class_target'.\n"
+ "If you use DKFZ's batchgenerators, please make sure you run version >= 0.20.1.")
+ gt_class_ids = batch['class_target']
gt_boxes = batch['bb_target']
#axes = (0, 2, 3, 1) if self.cf.dim == 2 else (0, 2, 3, 4, 1)
#gt_masks = [np.transpose(batch['roi_masks'][ii], axes=axes) for ii in range(len(batch['roi_masks']))]
# --> now GT masks has c==channels in last dimension.
gt_masks = batch['roi_masks']
img = torch.from_numpy(img).float().cuda()
batch_rpn_class_loss = torch.FloatTensor([0]).cuda()
batch_rpn_bbox_loss = torch.FloatTensor([0]).cuda()
# list of output boxes for monitoring/plotting. each element is a list of boxes per batch element.
box_results_list = [[] for _ in range(img.shape[0])]
#forward passes. 1. general forward pass, where no activations are saved in second stage (for performance
# monitoring and loss sampling). 2. second stage forward pass of sampled rois with stored activations for backprop.
rpn_class_logits, rpn_pred_deltas, proposal_boxes, detections, detection_masks = self.forward(img)
mrcnn_class_logits, mrcnn_pred_deltas, mrcnn_pred_mask, target_class_ids, mrcnn_target_deltas, target_mask, \
sample_proposals = self.loss_samples_forward(gt_class_ids, gt_boxes, gt_masks)
# loop over batch
for b in range(img.shape[0]):
if len(gt_boxes[b]) > 0:
# add gt boxes to output list for monitoring.
for ix in range(len(gt_boxes[b])):
box_results_list[b].append({'box_coords': batch['bb_target'][b][ix],
- 'box_label': batch['roi_labels'][b][ix], 'box_type': 'gt'})
+ 'box_label': batch['class_target'][b][ix], 'box_type': 'gt'})
# match gt boxes with anchors to generate targets for RPN losses.
rpn_match, rpn_target_deltas = mutils.gt_anchor_matching(self.cf, self.np_anchors, gt_boxes[b])
# add positive anchors used for loss to output list for monitoring.
pos_anchors = mutils.clip_boxes_numpy(self.np_anchors[np.argwhere(rpn_match == 1)][:, 0], img.shape[2:])
for p in pos_anchors:
box_results_list[b].append({'box_coords': p, 'box_type': 'pos_anchor'})
else:
rpn_match = np.array([-1]*self.np_anchors.shape[0])
rpn_target_deltas = np.array([0])
rpn_match_gpu = torch.from_numpy(rpn_match).cuda()
rpn_target_deltas = torch.from_numpy(rpn_target_deltas).float().cuda()
# compute RPN losses.
rpn_class_loss, neg_anchor_ix = compute_rpn_class_loss(rpn_match_gpu, rpn_class_logits[b], self.cf.shem_poolsize)
rpn_bbox_loss = compute_rpn_bbox_loss(rpn_target_deltas, rpn_pred_deltas[b], rpn_match_gpu)
batch_rpn_class_loss += rpn_class_loss / img.shape[0]
batch_rpn_bbox_loss += rpn_bbox_loss / img.shape[0]
# add negative anchors used for loss to output list for monitoring.
neg_anchors = mutils.clip_boxes_numpy(self.np_anchors[rpn_match == -1][neg_anchor_ix], img.shape[2:])
for n in neg_anchors:
box_results_list[b].append({'box_coords': n, 'box_type': 'neg_anchor'})
# add highest scoring proposals to output list for monitoring.
rpn_proposals = proposal_boxes[b][proposal_boxes[b, :, -1].argsort()][::-1]
for r in rpn_proposals[:self.cf.n_plot_rpn_props, :-1]:
box_results_list[b].append({'box_coords': r, 'box_type': 'prop'})
# add positive and negative roi samples used for mrcnn losses to output list for monitoring.
if 0 not in sample_proposals.shape:
rois = mutils.clip_to_window(self.cf.window, sample_proposals).cpu().data.numpy()
for ix, r in enumerate(rois):
box_results_list[int(r[-1])].append({'box_coords': r[:-1] * self.cf.scale,
'box_type': 'pos_class' if target_class_ids[ix] > 0 else 'neg_class'})
batch_rpn_class_loss = batch_rpn_class_loss
batch_rpn_bbox_loss = batch_rpn_bbox_loss
# compute mrcnn losses.
mrcnn_class_loss = compute_mrcnn_class_loss(target_class_ids, mrcnn_class_logits)
mrcnn_bbox_loss = compute_mrcnn_bbox_loss(mrcnn_target_deltas, mrcnn_pred_deltas, target_class_ids)
# mrcnn can be run without pixelwise annotations available (Faster R-CNN mode).
# In this case, the mask_loss is taken out of training.
if not self.cf.frcnn_mode:
mrcnn_mask_loss = compute_mrcnn_mask_loss(target_mask, mrcnn_pred_mask, target_class_ids)
else:
mrcnn_mask_loss = torch.FloatTensor([0]).cuda()
loss = batch_rpn_class_loss + batch_rpn_bbox_loss + mrcnn_class_loss + mrcnn_bbox_loss + mrcnn_mask_loss
# monitor RPN performance: detection count = the number of correctly matched proposals per fg-class.
dcount = [list(target_class_ids.cpu().data.numpy()).count(c) for c in np.arange(self.cf.head_classes)[1:]]
# run unmolding of predictions for monitoring and merge all results to one dictionary.
return_masks = True#self.cf.return_masks_in_val if is_validation else False
results_dict = get_results(self.cf, img.shape, detections, detection_masks,
box_results_list, return_masks=return_masks)
results_dict['torch_loss'] = loss
results_dict['monitor_values'] = {'loss': loss.item(), 'class_loss': mrcnn_class_loss.item()}
results_dict['logger_string'] = \
"loss: {0:.2f}, rpn_class: {1:.2f}, rpn_bbox: {2:.2f}, mrcnn_class: {3:.2f}, mrcnn_bbox: {4:.2f}, " \
"mrcnn_mask: {5:.2f}, dcount {6}".format(loss.item(), batch_rpn_class_loss.item(),
batch_rpn_bbox_loss.item(), mrcnn_class_loss.item(),
mrcnn_bbox_loss.item(), mrcnn_mask_loss.item(), dcount)
return results_dict
def test_forward(self, batch, return_masks=True):
"""
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 return_masks: boolean. If True, full resolution masks are returned for all proposals (speed trade-off).
: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 = batch['data']
img = torch.from_numpy(img).float().cuda()
_, _, _, detections, detection_masks = self.forward(img)
results_dict = get_results(self.cf, img.shape, detections, detection_masks, return_masks=return_masks)
return results_dict
def forward(self, img, is_training=True):
"""
:param img: input images (b, c, y, x, (z)).
:return: rpn_pred_logits: (b, n_anchors, 2)
:return: rpn_pred_deltas: (b, n_anchors, (y, x, (z), log(h), log(w), (log(d))))
:return: batch_proposal_boxes: (b, n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ix)) only for monitoring/plotting.
:return: detections: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score)
:return: detection_masks: (n_final_detections, n_classes, y, x, (z)) raw molded masks as returned by mask-head.
"""
# extract features.
fpn_outs = self.fpn(img)
rpn_feature_maps = [fpn_outs[i] for i in self.cf.pyramid_levels]
self.mrcnn_feature_maps = rpn_feature_maps
# loop through pyramid layers and apply RPN.
layer_outputs = [] # list of lists
for p in rpn_feature_maps:
layer_outputs.append(self.rpn(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]]
outputs = list(zip(*layer_outputs))
outputs = [torch.cat(list(o), dim=1) for o in outputs]
rpn_pred_logits, rpn_pred_probs, rpn_pred_deltas = outputs
# generate proposals: apply predicted deltas to anchors and filter by foreground scores from RPN classifier.
proposal_count = self.cf.post_nms_rois_training if is_training else self.cf.post_nms_rois_inference
batch_rpn_rois, batch_proposal_boxes = refine_proposals(rpn_pred_probs, rpn_pred_deltas, proposal_count, self.anchors, self.cf)
# merge batch dimension of proposals while storing allocation info in coordinate dimension.
batch_ixs = torch.from_numpy(np.repeat(np.arange(batch_rpn_rois.shape[0]), batch_rpn_rois.shape[1])).float().cuda()
rpn_rois = batch_rpn_rois.view(-1, batch_rpn_rois.shape[2])
self.rpn_rois_batch_info = torch.cat((rpn_rois, batch_ixs.unsqueeze(1)), dim=1)
# this is the first of two forward passes in the second stage, where no activations are stored for backprop.
# here, all proposals are forwarded (with virtual_batch_size = batch_size * post_nms_rois.)
# for inference/monitoring as well as sampling of rois for the loss functions.
# processed in chunks of roi_chunk_size to re-adjust to gpu-memory.
chunked_rpn_rois = self.rpn_rois_batch_info.split(self.cf.roi_chunk_size)
class_logits_list, bboxes_list = [], []
with torch.no_grad():
for chunk in chunked_rpn_rois:
chunk_class_logits, chunk_bboxes = self.classifier(self.mrcnn_feature_maps, chunk)
class_logits_list.append(chunk_class_logits)
bboxes_list.append(chunk_bboxes)
batch_mrcnn_class_logits = torch.cat(class_logits_list, 0)
batch_mrcnn_bbox = torch.cat(bboxes_list, 0)
self.batch_mrcnn_class_scores = F.softmax(batch_mrcnn_class_logits, dim=1)
# refine classified proposals, filter and return final detections.
detections = refine_detections(self.cf, batch_ixs, rpn_rois, batch_mrcnn_bbox, self.batch_mrcnn_class_scores)
# forward remaining detections through mask-head to generate corresponding masks.
scale = [img.shape[2]] * 4 + [img.shape[-1]] * 2
scale = torch.from_numpy(np.array(scale[:self.cf.dim * 2] + [1])[None]).float().cuda()
detection_boxes = detections[:, :self.cf.dim * 2 + 1] / scale
with torch.no_grad():
detection_masks = self.mask(self.mrcnn_feature_maps, detection_boxes)
return [rpn_pred_logits, rpn_pred_deltas, batch_proposal_boxes, detections, detection_masks]
def loss_samples_forward(self, batch_gt_class_ids, batch_gt_boxes, batch_gt_masks):
"""
this is the second forward pass through the second stage (features from stage one are re-used).
samples few rois in detection_target_layer and forwards only those for loss computation.
:param batch_gt_class_ids: list over batch elements. Each element is a list over the corresponding roi target labels.
: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)
:return: sample_logits: (n_sampled_rois, n_classes) predicted class scores.
:return: sample_boxes: (n_sampled_rois, n_classes, 2 * dim) predicted corrections to be applied to proposals for refinement.
:return: sample_mask: (n_sampled_rois, n_classes, y, x, (z)) predicted masks per class and proposal.
:return: sample_target_class_ids: (n_sampled_rois) target class labels of sampled proposals.
:return: sample_target_deltas: (n_sampled_rois, 2 * dim) target deltas of sampled proposals for box refinement.
:return: sample_target_masks: (n_sampled_rois, y, x, (z)) target masks of sampled proposals.
:return: sample_proposals: (n_sampled_rois, 2 * dim) RPN output for sampled proposals. only for monitoring/plotting.
"""
# sample rois for loss and get corresponding targets for all Mask R-CNN head network losses.
sample_ix, sample_target_class_ids, sample_target_deltas, sample_target_mask = \
detection_target_layer(self.rpn_rois_batch_info, self.batch_mrcnn_class_scores,
batch_gt_class_ids, batch_gt_boxes, batch_gt_masks, self.cf)
# re-use feature maps and RPN output from first forward pass.
sample_proposals = self.rpn_rois_batch_info[sample_ix]
if 0 not in sample_proposals.size():
sample_logits, sample_boxes = self.classifier(self.mrcnn_feature_maps, sample_proposals)
sample_mask = self.mask(self.mrcnn_feature_maps, sample_proposals)
else:
sample_logits = torch.FloatTensor().cuda()
sample_boxes = torch.FloatTensor().cuda()
sample_mask = torch.FloatTensor().cuda()
return [sample_logits, sample_boxes, sample_mask, sample_target_class_ids, sample_target_deltas,
sample_target_mask, sample_proposals]
\ No newline at end of file
diff --git a/models/retina_net.py b/models/retina_net.py
index a514a91..31140b4 100644
--- a/models/retina_net.py
+++ b/models/retina_net.py
@@ -1,504 +1,507 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""
Retina Net. According to https://arxiv.org/abs/1708.02002
Retina U-Net. According to https://arxiv.org/abs/1811.08661
"""
import utils.model_utils as mutils
import utils.exp_utils as utils
import sys
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils
sys.path.append('..')
from custom_extensions.nms import nms
############################################################
# Network Heads
############################################################
class Classifier(nn.Module):
def __init__(self, cf, conv):
"""
Builds the classifier sub-network.
"""
super(Classifier, self).__init__()
self.dim = conv.dim
self.n_classes = cf.head_classes
n_input_channels = cf.end_filts
n_features = cf.n_rpn_features
n_output_channels = cf.n_anchors_per_pos * cf.head_classes
anchor_stride = cf.rpn_anchor_stride
self.conv_1 = conv(n_input_channels, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_2 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_3 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_4 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_final = conv(n_features, n_output_channels, ks=3, stride=anchor_stride, pad=1, relu=None)
def forward(self, x):
"""
:param x: input feature map (b, in_c, y, x, (z))
:return: class_logits (b, n_anchors, n_classes)
"""
x = self.conv_1(x)
x = self.conv_2(x)
x = self.conv_3(x)
x = self.conv_4(x)
class_logits = self.conv_final(x)
axes = (0, 2, 3, 1) if self.dim == 2 else (0, 2, 3, 4, 1)
class_logits = class_logits.permute(*axes)
class_logits = class_logits.contiguous()
class_logits = class_logits.view(x.size()[0], -1, self.n_classes)
return [class_logits]
class BBRegressor(nn.Module):
def __init__(self, cf, conv):
"""
Builds the bb-regression sub-network.
"""
super(BBRegressor, self).__init__()
self.dim = conv.dim
n_input_channels = cf.end_filts
n_features = cf.n_rpn_features
n_output_channels = cf.n_anchors_per_pos * self.dim * 2
anchor_stride = cf.rpn_anchor_stride
self.conv_1 = conv(n_input_channels, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_2 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_3 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_4 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_final = conv(n_features, n_output_channels, ks=3, stride=anchor_stride,
pad=1, relu=None)
def forward(self, x):
"""
:param x: input feature map (b, in_c, y, x, (z))
:return: bb_logits (b, n_anchors, dim * 2)
"""
x = self.conv_1(x)
x = self.conv_2(x)
x = self.conv_3(x)
x = self.conv_4(x)
bb_logits = self.conv_final(x)
axes = (0, 2, 3, 1) if self.dim == 2 else (0, 2, 3, 4, 1)
bb_logits = bb_logits.permute(*axes)
bb_logits = bb_logits.contiguous()
bb_logits = bb_logits.view(x.size()[0], -1, self.dim * 2)
return [bb_logits]
############################################################
# Loss Functions
############################################################
def compute_class_loss(anchor_matches, class_pred_logits, shem_poolsize=20):
"""
:param anchor_matches: (n_anchors). [-1, 0, class_id] for negative, neutral, and positive matched anchors.
:param class_pred_logits: (n_anchors, n_classes). logits from classifier sub-network.
:param shem_poolsize: int. factor of top-k candidates to draw from per negative sample (online-hard-example-mining).
:return: loss: torch tensor.
:return: np_neg_ix: 1D array containing indices of the neg_roi_logits, which have been sampled for training.
"""
# Positive and Negative anchors contribute to the loss,
# but neutral anchors (match value = 0) don't.
pos_indices = torch.nonzero(anchor_matches > 0)
neg_indices = torch.nonzero(anchor_matches == -1)
# get positive samples and calculate loss.
if 0 not in pos_indices.shape:
pos_indices = pos_indices.squeeze(1)
roi_logits_pos = class_pred_logits[pos_indices]
targets_pos = anchor_matches[pos_indices]
pos_loss = F.cross_entropy(roi_logits_pos, targets_pos.long())
else:
pos_loss = torch.FloatTensor([0]).cuda()
# get negative samples, such that the amount matches the number of positive samples, but at least 1.
# get high scoring negatives by applying online-hard-example-mining.
if 0 not in neg_indices.shape:
neg_indices = neg_indices.squeeze(1)
roi_logits_neg = class_pred_logits[neg_indices]
negative_count = np.max((1, pos_indices.shape[0]))
roi_probs_neg = F.softmax(roi_logits_neg, dim=1)
neg_ix = mutils.shem(roi_probs_neg, negative_count, shem_poolsize)
neg_loss = F.cross_entropy(roi_logits_neg[neg_ix], torch.LongTensor([0] * neg_ix.shape[0]).cuda())
# return the indices of negative samples, which contributed to the loss (for monitoring plots).
np_neg_ix = neg_ix.cpu().data.numpy()
else:
neg_loss = torch.FloatTensor([0]).cuda()
np_neg_ix = np.array([]).astype('int32')
loss = (pos_loss + neg_loss) / 2
return loss, np_neg_ix
def compute_bbox_loss(target_deltas, pred_deltas, anchor_matches):
"""
:param target_deltas: (b, n_positive_anchors, (dy, dx, (dz), log(dh), log(dw), (log(dd)))).
Uses 0 padding to fill in unsed bbox deltas.
:param pred_deltas: predicted deltas from bbox regression head. (b, n_anchors, (dy, dx, (dz), log(dh), log(dw), (log(dd))))
:param anchor_matches: (n_anchors). [-1, 0, class_id] for negative, neutral, and positive matched anchors.
:return: loss: torch 1D tensor.
"""
if 0 not in torch.nonzero(anchor_matches > 0).size():
indices = torch.nonzero(anchor_matches > 0).squeeze(1)
# Pick bbox deltas that contribute to the loss
pred_deltas = pred_deltas[indices]
# Trim target bounding box deltas to the same length as pred_deltas.
target_deltas = target_deltas[:pred_deltas.size()[0], :]
# Smooth L1 loss
loss = F.smooth_l1_loss(pred_deltas, target_deltas)
else:
loss = torch.FloatTensor([0]).cuda()
return loss
############################################################
# Output Handler
############################################################
def refine_detections(anchors, probs, deltas, batch_ixs, cf):
"""
Refine classified proposals, filter overlaps and return final
detections. n_proposals here is typically a very large number: batch_size * n_anchors.
This function is hence optimized on trimming down n_proposals.
:param anchors: (n_anchors, 2 * dim)
:param probs: (n_proposals, n_classes) softmax probabilities for all rois as predicted by classifier head.
:param deltas: (n_proposals, n_classes, 2 * dim) box refinement deltas as predicted by bbox regressor head.
:param batch_ixs: (n_proposals) batch element assignemnt info for re-allocation.
:return: result: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score))
"""
anchors = anchors.repeat(batch_ixs.unique().shape[0], 1)
# flatten foreground probabilities, sort and trim down to highest confidences by pre_nms limit.
fg_probs = probs[:, 1:].contiguous()
flat_probs, flat_probs_order = fg_probs.view(-1).sort(descending=True)
keep_ix = flat_probs_order[:cf.pre_nms_limit]
# reshape indices to 2D index array with shape like fg_probs.
keep_arr = torch.cat(((keep_ix / fg_probs.shape[1]).unsqueeze(1), (keep_ix % fg_probs.shape[1]).unsqueeze(1)), 1)
pre_nms_scores = flat_probs[:cf.pre_nms_limit]
pre_nms_class_ids = keep_arr[:, 1] + 1 # add background again.
pre_nms_batch_ixs = batch_ixs[keep_arr[:, 0]]
pre_nms_anchors = anchors[keep_arr[:, 0]]
pre_nms_deltas = deltas[keep_arr[:, 0]]
keep = torch.arange(pre_nms_scores.shape[0]).long().cuda()
# apply bounding box deltas. re-scale to image coordinates.
std_dev = torch.from_numpy(np.reshape(cf.rpn_bbox_std_dev, [1, cf.dim * 2])).float().cuda()
scale = torch.from_numpy(cf.scale).float().cuda()
refined_rois = mutils.apply_box_deltas_2D(pre_nms_anchors / scale, pre_nms_deltas * std_dev) * scale \
if cf.dim == 2 else mutils.apply_box_deltas_3D(pre_nms_anchors / scale, pre_nms_deltas * std_dev) * scale
# round and cast to int since we're deadling with pixels now
refined_rois = mutils.clip_to_window(cf.window, refined_rois)
pre_nms_rois = torch.round(refined_rois)
for j, b in enumerate(mutils.unique1d(pre_nms_batch_ixs)):
bixs = torch.nonzero(pre_nms_batch_ixs == b)[:, 0]
bix_class_ids = pre_nms_class_ids[bixs]
bix_rois = pre_nms_rois[bixs]
bix_scores = pre_nms_scores[bixs]
for i, class_id in enumerate(mutils.unique1d(bix_class_ids)):
ixs = torch.nonzero(bix_class_ids == class_id)[:, 0]
# nms expects boxes sorted by score.
ix_rois = bix_rois[ixs]
ix_scores = bix_scores[ixs]
ix_scores, order = ix_scores.sort(descending=True)
ix_rois = ix_rois[order, :]
class_keep = nms.nms(ix_rois, ix_scores, cf.detection_nms_threshold)
# map indices back.
class_keep = keep[bixs[ixs[order[class_keep]]]]
# merge indices over classes for current batch element
b_keep = class_keep if i == 0 else mutils.unique1d(torch.cat((b_keep, class_keep)))
# only keep top-k boxes of current batch-element.
top_ids = pre_nms_scores[b_keep].sort(descending=True)[1][:cf.model_max_instances_per_batch_element]
b_keep = b_keep[top_ids]
# merge indices over batch elements.
batch_keep = b_keep if j == 0 else mutils.unique1d(torch.cat((batch_keep, b_keep)))
keep = batch_keep
# arrange output.
result = torch.cat((pre_nms_rois[keep],
pre_nms_batch_ixs[keep].unsqueeze(1).float(),
pre_nms_class_ids[keep].unsqueeze(1).float(),
pre_nms_scores[keep].unsqueeze(1)), dim=1)
return result
def get_results(cf, img_shape, detections, seg_logits, box_results_list=None):
"""
Restores batch dimension of merged detections, unmolds detections, creates and fills results dict.
:param img_shape:
:param detections: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score)
:param box_results_list: None or list of output boxes for monitoring/plotting.
each element is a list of boxes per batch element.
:return: results_dict: dictionary with keys:
'boxes': list over batch elements. each batch element is a list of boxes. each box is a dictionary:
[[{box_0}, ... {box_n}], [{box_0}, ... {box_n}], ...]
'seg_preds': pixel-wise class predictions (b, 1, y, x, (z)) with values [0, ..., n_classes] for
retina_unet and dummy array for retina_net.
"""
detections = detections.cpu().data.numpy()
batch_ixs = detections[:, cf.dim*2]
detections = [detections[batch_ixs == ix] for ix in range(img_shape[0])]
# for test_forward, where no previous list exists.
if box_results_list is None:
box_results_list = [[] for _ in range(img_shape[0])]
for ix in range(img_shape[0]):
if 0 not in detections[ix].shape:
boxes = detections[ix][:, :2 * cf.dim].astype(np.int32)
class_ids = detections[ix][:, 2 * cf.dim + 1].astype(np.int32)
scores = detections[ix][:, 2 * cf.dim + 2]
# Filter out detections with zero area. Often only happens in early
# stages of training when the network weights are still a bit random.
if cf.dim == 2:
exclude_ix = np.where((boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) <= 0)[0]
else:
exclude_ix = np.where(
(boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 5] - boxes[:, 4]) <= 0)[0]
if exclude_ix.shape[0] > 0:
boxes = np.delete(boxes, exclude_ix, axis=0)
class_ids = np.delete(class_ids, exclude_ix, axis=0)
scores = np.delete(scores, exclude_ix, axis=0)
if 0 not in boxes.shape:
for ix2, score in enumerate(scores):
if score >= cf.model_min_confidence:
box_results_list[ix].append({'box_coords': boxes[ix2],
'box_score': score,
'box_type': 'det',
'box_pred_class_id': class_ids[ix2]})
results_dict = {'boxes': box_results_list}
if seg_logits is None:
# output dummy segmentation for retina_net.
results_dict['seg_preds'] = np.zeros(img_shape)[:, 0][:, np.newaxis]
else:
# output label maps for retina_unet.
results_dict['seg_preds'] = F.softmax(seg_logits, 1).argmax(1).cpu().data.numpy()[:, np.newaxis].astype('uint8')
return results_dict
############################################################
# Retina (U-)Net Class
############################################################
class net(nn.Module):
def __init__(self, cf, logger):
super(net, self).__init__()
self.cf = cf
self.logger = logger
self.build()
if self.cf.weight_init is not None:
logger.info("using pytorch weight init of type {}".format(self.cf.weight_init))
mutils.initialize_weights(self)
else:
logger.info("using default pytorch weight init")
def build(self):
"""
Build Retina Net architecture.
"""
# Image size must be dividable by 2 multiple times.
h, w = self.cf.patch_size[:2]
if h / 2 ** 5 != int(h / 2 ** 5) or w / 2 ** 5 != int(w / 2 ** 5):
raise Exception("Image size must be dividable by 2 at least 5 times "
"to avoid fractions when downscaling and upscaling."
"For example, use 256, 320, 384, 448, 512, ... etc. ")
# instanciate abstract multi dimensional conv class and backbone model.
conv = mutils.NDConvGenerator(self.cf.dim)
backbone = utils.import_module('bbone', self.cf.backbone_path)
# build Anchors, FPN, Classifier / Bbox-Regressor -head
self.np_anchors = mutils.generate_pyramid_anchors(self.logger, self.cf)
self.anchors = torch.from_numpy(self.np_anchors).float().cuda()
self.Fpn = backbone.FPN(self.cf, conv, operate_stride1=self.cf.operate_stride1)
self.Classifier = Classifier(self.cf, conv)
self.BBRegressor = BBRegressor(self.cf, conv)
def train_forward(self, batch, **kwargs):
"""
train method (also used for validation monitoring). wrapper around forward pass of network. prepares input data
for processing, computes losses, and stores outputs in a dictionary.
:param batch: dictionary containing 'data', 'seg', etc.
:return: results_dict: dictionary with keys:
'boxes': list over batch elements. each batch element is a list of boxes. each box is a dictionary:
[[{box_0}, ... {box_n}], [{box_0}, ... {box_n}], ...]
'seg_preds': pixelwise segmentation output (b, c, y, x, (z)) with values [0, .., n_classes].
'monitor_values': dict of values to be monitored.
"""
img = batch['data']
- gt_class_ids = batch['roi_labels']
+ if "roi_labels" in batch.keys():
+ raise Exception("Key for roi-wise class targets changed in v0.1.0 from 'roi_labels' to 'class_target'.\n"
+ "If you use DKFZ's batchgenerators, please make sure you run version >= 0.20.1.")
+ gt_class_ids = batch['class_target']
gt_boxes = batch['bb_target']
img = torch.from_numpy(img).float().cuda()
batch_class_loss = torch.FloatTensor([0]).cuda()
batch_bbox_loss = torch.FloatTensor([0]).cuda()
# list of output boxes for monitoring/plotting. each element is a list of boxes per batch element.
box_results_list = [[] for _ in range(img.shape[0])]
detections, class_logits, pred_deltas, seg_logits = self.forward(img)
# loop over batch
for b in range(img.shape[0]):
# add gt boxes to results dict for monitoring.
if len(gt_boxes[b]) > 0:
for ix in range(len(gt_boxes[b])):
box_results_list[b].append({'box_coords': batch['bb_target'][b][ix],
- 'box_label': batch['roi_labels'][b][ix], 'box_type': 'gt'})
+ 'box_label': batch['class_target'][b][ix], 'box_type': 'gt'})
# match gt boxes with anchors to generate targets.
anchor_class_match, anchor_target_deltas = mutils.gt_anchor_matching(
self.cf, self.np_anchors, gt_boxes[b], gt_class_ids[b])
# add positive anchors used for loss to results_dict for monitoring.
pos_anchors = mutils.clip_boxes_numpy(
self.np_anchors[np.argwhere(anchor_class_match > 0)][:, 0], img.shape[2:])
for p in pos_anchors:
box_results_list[b].append({'box_coords': p, 'box_type': 'pos_anchor'})
else:
anchor_class_match = np.array([-1]*self.np_anchors.shape[0])
anchor_target_deltas = np.array([0])
anchor_class_match = torch.from_numpy(anchor_class_match).cuda()
anchor_target_deltas = torch.from_numpy(anchor_target_deltas).float().cuda()
# compute losses.
class_loss, neg_anchor_ix = compute_class_loss(anchor_class_match, class_logits[b])
bbox_loss = compute_bbox_loss(anchor_target_deltas, pred_deltas[b], anchor_class_match)
# add negative anchors used for loss to results_dict for monitoring.
neg_anchors = mutils.clip_boxes_numpy(
self.np_anchors[np.argwhere(anchor_class_match.cpu().numpy() == -1)][neg_anchor_ix, 0],
img.shape[2:])
for n in neg_anchors:
box_results_list[b].append({'box_coords': n, 'box_type': 'neg_anchor'})
batch_class_loss += class_loss / img.shape[0]
batch_bbox_loss += bbox_loss / img.shape[0]
results_dict = get_results(self.cf, img.shape, detections, seg_logits, box_results_list)
loss = batch_class_loss + batch_bbox_loss
results_dict['torch_loss'] = loss
results_dict['class_loss'] = batch_class_loss.item()
results_dict['logger_string'] = "loss: {0:.2f}, class: {1:.2f}, bbox: {2:.2f}"\
.format(loss.item(), batch_class_loss.item(), batch_bbox_loss.item())
return results_dict
def test_forward(self, batch, **kwargs):
"""
test method. wrapper around forward pass of network without usage of any ground truth information.
prepares input data for processing and stores outputs in a dictionary.
:param batch: dictionary containing 'data'
:return: results_dict: dictionary with keys:
'boxes': list over batch elements. each batch element is a list of boxes. each box is a dictionary:
[[{box_0}, ... {box_n}], [{box_0}, ... {box_n}], ...]
'seg_preds': pixel-wise class predictions (b, 1, y, x, (z)) with values [0, ..., n_classes] for
retina_unet and dummy array for retina_net.
"""
img = batch['data']
img = torch.from_numpy(img).float().cuda()
detections, _, _, seg_logits = self.forward(img)
results_dict = get_results(self.cf, img.shape, detections, seg_logits)
return results_dict
def forward(self, img):
"""
forward pass of the model.
:param img: input img (b, c, y, x, (z)).
:return: rpn_pred_logits: (b, n_anchors, 2)
:return: rpn_pred_deltas: (b, n_anchors, (y, x, (z), log(h), log(w), (log(d))))
:return: batch_proposal_boxes: (b, n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ix)) only for monitoring/plotting.
:return: detections: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score)
:return: detection_masks: (n_final_detections, n_classes, y, x, (z)) raw molded masks as returned by mask-head.
"""
# Feature extraction
fpn_outs = self.Fpn(img)
seg_logits = None
selected_fmaps = [fpn_outs[i] for i in self.cf.pyramid_levels]
# Loop through pyramid layers
class_layer_outputs, bb_reg_layer_outputs = [], [] # list of lists
for p in selected_fmaps:
class_layer_outputs.append(self.Classifier(p))
bb_reg_layer_outputs.append(self.BBRegressor(p))
# Concatenate layer outputs
# Convert from list of lists of level outputs to list of lists
# of outputs across levels.
# e.g. [[a1, b1, c1], [a2, b2, c2]] => [[a1, a2], [b1, b2], [c1, c2]]
class_logits = list(zip(*class_layer_outputs))
class_logits = [torch.cat(list(o), dim=1) for o in class_logits][0]
bb_outputs = list(zip(*bb_reg_layer_outputs))
bb_outputs = [torch.cat(list(o), dim=1) for o in bb_outputs][0]
# merge batch_dimension and store info in batch_ixs for re-allocation.
batch_ixs = torch.arange(class_logits.shape[0]).unsqueeze(1).repeat(1, class_logits.shape[1]).view(-1).cuda()
flat_class_softmax = F.softmax(class_logits.view(-1, class_logits.shape[-1]), 1)
flat_bb_outputs = bb_outputs.view(-1, bb_outputs.shape[-1])
detections = refine_detections(self.anchors, flat_class_softmax, flat_bb_outputs, batch_ixs, self.cf)
return detections, class_logits, bb_outputs, seg_logits
diff --git a/models/retina_unet.py b/models/retina_unet.py
index a0d6c17..4582f57 100644
--- a/models/retina_unet.py
+++ b/models/retina_unet.py
@@ -1,510 +1,513 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""
Retina Net. According to https://arxiv.org/abs/1708.02002
Retina U-Net. According to https://arxiv.org/abs/1811.08661
"""
import utils.model_utils as mutils
import utils.exp_utils as utils
import sys
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils
sys.path.append('../')
from custom_extensions.nms import nms
############################################################
# Network Heads
############################################################
class Classifier(nn.Module):
def __init__(self, cf, conv):
"""
Builds the classifier sub-network.
"""
super(Classifier, self).__init__()
self.dim = conv.dim
self.n_classes = cf.head_classes
n_input_channels = cf.end_filts
n_features = cf.n_rpn_features
n_output_channels = cf.n_anchors_per_pos * cf.head_classes
anchor_stride = cf.rpn_anchor_stride
self.conv_1 = conv(n_input_channels, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_2 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_3 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_4 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_final = conv(n_features, n_output_channels, ks=3, stride=anchor_stride, pad=1, relu=None)
def forward(self, x):
"""
:param x: input feature map (b, in_c, y, x, (z))
:return: class_logits (b, n_anchors, n_classes)
"""
x = self.conv_1(x)
x = self.conv_2(x)
x = self.conv_3(x)
x = self.conv_4(x)
class_logits = self.conv_final(x)
axes = (0, 2, 3, 1) if self.dim == 2 else (0, 2, 3, 4, 1)
class_logits = class_logits.permute(*axes)
class_logits = class_logits.contiguous()
class_logits = class_logits.view(x.size()[0], -1, self.n_classes)
return [class_logits]
class BBRegressor(nn.Module):
def __init__(self, cf, conv):
"""
Builds the bb-regression sub-network.
"""
super(BBRegressor, self).__init__()
self.dim = conv.dim
n_input_channels = cf.end_filts
n_features = cf.n_rpn_features
n_output_channels = cf.n_anchors_per_pos * self.dim * 2
anchor_stride = cf.rpn_anchor_stride
self.conv_1 = conv(n_input_channels, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_2 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_3 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_4 = conv(n_features, n_features, ks=3, stride=anchor_stride, pad=1, relu=cf.relu)
self.conv_final = conv(n_features, n_output_channels, ks=3, stride=anchor_stride,
pad=1, relu=None)
def forward(self, x):
"""
:param x: input feature map (b, in_c, y, x, (z))
:return: bb_logits (b, n_anchors, dim * 2)
"""
x = self.conv_1(x)
x = self.conv_2(x)
x = self.conv_3(x)
x = self.conv_4(x)
bb_logits = self.conv_final(x)
axes = (0, 2, 3, 1) if self.dim == 2 else (0, 2, 3, 4, 1)
bb_logits = bb_logits.permute(*axes)
bb_logits = bb_logits.contiguous()
bb_logits = bb_logits.view(x.size()[0], -1, self.dim * 2)
return [bb_logits]
############################################################
# Loss Functions
############################################################
def compute_class_loss(anchor_matches, class_pred_logits, shem_poolsize=20):
"""
:param anchor_matches: (n_anchors). [-1, 0, class_id] for negative, neutral, and positive matched anchors.
:param class_pred_logits: (n_anchors, n_classes). logits from classifier sub-network.
:param shem_poolsize: int. factor of top-k candidates to draw from per negative sample (online-hard-example-mining).
:return: loss: torch tensor.
:return: np_neg_ix: 1D array containing indices of the neg_roi_logits, which have been sampled for training.
"""
# Positive and Negative anchors contribute to the loss,
# but neutral anchors (match value = 0) don't.
pos_indices = torch.nonzero(anchor_matches > 0)
neg_indices = torch.nonzero(anchor_matches == -1)
# get positive samples and calucalte loss.
if 0 not in pos_indices.size():
pos_indices = pos_indices.squeeze(1)
roi_logits_pos = class_pred_logits[pos_indices]
targets_pos = anchor_matches[pos_indices]
pos_loss = F.cross_entropy(roi_logits_pos, targets_pos.long())
else:
pos_loss = torch.FloatTensor([0]).cuda()
# get negative samples, such that the amount matches the number of positive samples, but at least 1.
# get high scoring negatives by applying online-hard-example-mining.
if 0 not in neg_indices.size():
neg_indices = neg_indices.squeeze(1)
roi_logits_neg = class_pred_logits[neg_indices]
negative_count = np.max((1, pos_indices.size()[0]))
roi_probs_neg = F.softmax(roi_logits_neg, dim=1)
neg_ix = mutils.shem(roi_probs_neg, negative_count, shem_poolsize)
neg_loss = F.cross_entropy(roi_logits_neg[neg_ix], torch.LongTensor([0] * neg_ix.shape[0]).cuda())
# return the indices of negative samples, which contributed to the loss (for monitoring plots).
np_neg_ix = neg_ix.cpu().data.numpy()
else:
neg_loss = torch.FloatTensor([0]).cuda()
np_neg_ix = np.array([]).astype('int32')
loss = (pos_loss + neg_loss) / 2
return loss, np_neg_ix
def compute_bbox_loss(target_deltas, pred_deltas, anchor_matches):
"""
:param target_deltas: (b, n_positive_anchors, (dy, dx, (dz), log(dh), log(dw), (log(dd)))).
Uses 0 padding to fill in unsed bbox deltas.
:param pred_deltas: predicted deltas from bbox regression head. (b, n_anchors, (dy, dx, (dz), log(dh), log(dw), (log(dd))))
:param anchor_matches: (n_anchors). [-1, 0, class_id] for negative, neutral, and positive matched anchors.
:return: loss: torch 1D tensor.
"""
if 0 not in torch.nonzero(anchor_matches > 0).size():
indices = torch.nonzero(anchor_matches > 0).squeeze(1)
# Pick bbox deltas that contribute to the loss
pred_deltas = pred_deltas[indices]
# Trim target bounding box deltas to the same length as pred_deltas.
target_deltas = target_deltas[:pred_deltas.size()[0], :]
# Smooth L1 loss
loss = F.smooth_l1_loss(pred_deltas, target_deltas)
else:
loss = torch.FloatTensor([0]).cuda()
return loss
############################################################
# Output Handler
############################################################
def refine_detections(anchors, probs, deltas, batch_ixs, cf):
"""
Refine classified proposals, filter overlaps and return final
detections. n_proposals here is typically a very large number: batch_size * n_anchors.
This function is hence optimized on trimming down n_proposals.
:param anchors: (n_anchors, 2 * dim)
:param probs: (n_proposals, n_classes) softmax probabilities for all rois as predicted by classifier head.
:param deltas: (n_proposals, n_classes, 2 * dim) box refinement deltas as predicted by bbox regressor head.
:param batch_ixs: (n_proposals) batch element assignemnt info for re-allocation.
:return: result: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score))
"""
anchors = anchors.repeat(batch_ixs.unique().shape[0], 1)
# flatten foreground probabilities, sort and trim down to highest confidences by pre_nms limit.
fg_probs = probs[:, 1:].contiguous()
flat_probs, flat_probs_order = fg_probs.view(-1).sort(descending=True)
keep_ix = flat_probs_order[:cf.pre_nms_limit]
# reshape indices to 2D index array with shape like fg_probs.
keep_arr = torch.cat(((keep_ix / fg_probs.shape[1]).unsqueeze(1), (keep_ix % fg_probs.shape[1]).unsqueeze(1)), 1)
pre_nms_scores = flat_probs[:cf.pre_nms_limit]
pre_nms_class_ids = keep_arr[:, 1] + 1 # add background again.
pre_nms_batch_ixs = batch_ixs[keep_arr[:, 0]]
pre_nms_anchors = anchors[keep_arr[:, 0]]
pre_nms_deltas = deltas[keep_arr[:, 0]]
keep = torch.arange(pre_nms_scores.size()[0]).long().cuda()
# apply bounding box deltas. re-scale to image coordinates.
std_dev = torch.from_numpy(np.reshape(cf.rpn_bbox_std_dev, [1, cf.dim * 2])).float().cuda()
scale = torch.from_numpy(cf.scale).float().cuda()
refined_rois = mutils.apply_box_deltas_2D(pre_nms_anchors / scale, pre_nms_deltas * std_dev) * scale \
if cf.dim == 2 else mutils.apply_box_deltas_3D(pre_nms_anchors / scale, pre_nms_deltas * std_dev) * scale
# round and cast to int since we're deadling with pixels now
refined_rois = mutils.clip_to_window(cf.window, refined_rois)
pre_nms_rois = torch.round(refined_rois)
for j, b in enumerate(mutils.unique1d(pre_nms_batch_ixs)):
bixs = torch.nonzero(pre_nms_batch_ixs == b)[:, 0]
bix_class_ids = pre_nms_class_ids[bixs]
bix_rois = pre_nms_rois[bixs]
bix_scores = pre_nms_scores[bixs]
for i, class_id in enumerate(mutils.unique1d(bix_class_ids)):
ixs = torch.nonzero(bix_class_ids == class_id)[:, 0]
# nms expects boxes sorted by score.
ix_rois = bix_rois[ixs]
ix_scores = bix_scores[ixs]
ix_scores, order = ix_scores.sort(descending=True)
ix_rois = ix_rois[order, :]
ix_scores = ix_scores
class_keep = nms.nms(ix_rois, ix_scores, cf.detection_nms_threshold)
# map indices back.
class_keep = keep[bixs[ixs[order[class_keep]]]]
# merge indices over classes for current batch element
b_keep = class_keep if i == 0 else mutils.unique1d(torch.cat((b_keep, class_keep)))
# only keep top-k boxes of current batch-element.
top_ids = pre_nms_scores[b_keep].sort(descending=True)[1][:cf.model_max_instances_per_batch_element]
b_keep = b_keep[top_ids]
# merge indices over batch elements.
batch_keep = b_keep if j == 0 else mutils.unique1d(torch.cat((batch_keep, b_keep)))
keep = batch_keep
# arrange output.
result = torch.cat((pre_nms_rois[keep],
pre_nms_batch_ixs[keep].unsqueeze(1).float(),
pre_nms_class_ids[keep].unsqueeze(1).float(),
pre_nms_scores[keep].unsqueeze(1)), dim=1)
return result
def get_results(cf, img_shape, detections, seg_logits, box_results_list=None):
"""
Restores batch dimension of merged detections, unmolds detections, creates and fills results dict.
:param img_shape:
:param detections: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score)
:param box_results_list: None or list of output boxes for monitoring/plotting.
each element is a list of boxes per batch element.
:return: results_dict: dictionary with keys:
'boxes': list over batch elements. each batch element is a list of boxes. each box is a dictionary:
[[{box_0}, ... {box_n}], [{box_0}, ... {box_n}], ...]
'seg_preds': pixel-wise class predictions (b, 1, y, x, (z)) with values [0, ..., n_classes] for
retina_unet and dummy array for retina_net.
"""
detections = detections.cpu().data.numpy()
batch_ixs = detections[:, cf.dim*2]
detections = [detections[batch_ixs == ix] for ix in range(img_shape[0])]
# for test_forward, where no previous list exists.
if box_results_list is None:
box_results_list = [[] for _ in range(img_shape[0])]
for ix in range(img_shape[0]):
if 0 not in detections[ix].shape:
boxes = detections[ix][:, :2 * cf.dim].astype(np.int32)
class_ids = detections[ix][:, 2 * cf.dim + 1].astype(np.int32)
scores = detections[ix][:, 2 * cf.dim + 2]
# Filter out detections with zero area. Often only happens in early
# stages of training when the network weights are still a bit random.
if cf.dim == 2:
exclude_ix = np.where((boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) <= 0)[0]
else:
exclude_ix = np.where(
(boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 5] - boxes[:, 4]) <= 0)[0]
if exclude_ix.shape[0] > 0:
boxes = np.delete(boxes, exclude_ix, axis=0)
class_ids = np.delete(class_ids, exclude_ix, axis=0)
scores = np.delete(scores, exclude_ix, axis=0)
if 0 not in boxes.shape:
for ix2, score in enumerate(scores):
if score >= cf.model_min_confidence:
box_results_list[ix].append({'box_coords': boxes[ix2],
'box_score': score,
'box_type': 'det',
'box_pred_class_id': class_ids[ix2]})
results_dict = {'boxes': box_results_list}
if seg_logits is None:
# output dummy segmentation for retina_net.
results_dict['seg_preds'] = np.zeros(img_shape)[:, 0][:, np.newaxis]
else:
# output label maps for retina_unet.
results_dict['seg_preds'] = F.softmax(seg_logits, 1).argmax(1).cpu().data.numpy()[:, np.newaxis].astype('uint8')
return results_dict
############################################################
# Retina (U-)Net Class
############################################################
class net(nn.Module):
def __init__(self, cf, logger):
super(net, self).__init__()
self.cf = cf
self.logger = logger
self.build()
if self.cf.weight_init is not None:
logger.info("using pytorch weight init of type {}".format(self.cf.weight_init))
mutils.initialize_weights(self)
else:
logger.info("using default pytorch weight init")
def build(self):
"""
Build Retina Net architecture.
"""
# Image size must be dividable by 2 multiple times.
h, w = self.cf.patch_size[:2]
if h / 2 ** 5 != int(h / 2 ** 5) or w / 2 ** 5 != int(w / 2 ** 5):
raise Exception("Image size must be dividable by 2 at least 5 times "
"to avoid fractions when downscaling and upscaling."
"For example, use 256, 320, 384, 448, 512, ... etc. ")
# instanciate abstract multi dimensional conv class and backbone model.
conv = mutils.NDConvGenerator(self.cf.dim)
backbone = utils.import_module('bbone', self.cf.backbone_path)
# build Anchors, FPN, Classifier / Bbox-Regressor -head
self.np_anchors = mutils.generate_pyramid_anchors(self.logger, self.cf)
self.anchors = torch.from_numpy(self.np_anchors).float().cuda()
self.Fpn = backbone.FPN(self.cf, conv, operate_stride1=self.cf.operate_stride1)
self.Classifier = Classifier(self.cf, conv)
self.BBRegressor = BBRegressor(self.cf, conv)
self.final_conv = conv(self.cf.end_filts, self.cf.num_seg_classes, ks=1, pad=0, norm=None, relu=None)
def train_forward(self, batch, **kwargs):
"""
train method (also used for validation monitoring). wrapper around forward pass of network. prepares input data
for processing, computes losses, and stores outputs in a dictionary.
:param batch: dictionary containing 'data', 'seg', etc.
:return: results_dict: dictionary with keys:
'boxes': list over batch elements. each batch element is a list of boxes. each box is a dictionary:
[[{box_0}, ... {box_n}], [{box_0}, ... {box_n}], ...]
'seg_preds': pixelwise segmentation output (b, c, y, x, (z)) with values [0, .., n_classes].
'monitor_values': dict of values to be monitored.
"""
img = batch['data']
- gt_class_ids = batch['roi_labels']
+ if "roi_labels" in batch.keys():
+ raise Exception("Key for roi-wise class targets changed in v0.1.0 from 'roi_labels' to 'class_target'.\n"
+ "If you use DKFZ's batchgenerators, please make sure you run version >= 0.20.1.")
+ gt_class_ids = batch['class_target']
gt_boxes = batch['bb_target']
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()
batch_class_loss = torch.FloatTensor([0]).cuda()
batch_bbox_loss = torch.FloatTensor([0]).cuda()
# list of output boxes for monitoring/plotting. each element is a list of boxes per batch element.
box_results_list = [[] for _ in range(img.shape[0])]
detections, class_logits, pred_deltas, seg_logits = self.forward(img)
# loop over batch
for b in range(img.shape[0]):
# add gt boxes to results dict for monitoring.
if len(gt_boxes[b]) > 0:
for ix in range(len(gt_boxes[b])):
box_results_list[b].append({'box_coords': batch['bb_target'][b][ix],
- 'box_label': batch['roi_labels'][b][ix], 'box_type': 'gt'})
+ 'box_label': batch['class_target'][b][ix], 'box_type': 'gt'})
# match gt boxes with anchors to generate targets.
anchor_class_match, anchor_target_deltas = mutils.gt_anchor_matching(
self.cf, self.np_anchors, gt_boxes[b], gt_class_ids[b])
# add positive anchors used for loss to results_dict for monitoring.
pos_anchors = mutils.clip_boxes_numpy(
self.np_anchors[np.argwhere(anchor_class_match > 0)][:, 0], img.shape[2:])
for p in pos_anchors:
box_results_list[b].append({'box_coords': p, 'box_type': 'pos_anchor'})
else:
anchor_class_match = np.array([-1]*self.np_anchors.shape[0])
anchor_target_deltas = np.array([0])
anchor_class_match = torch.from_numpy(anchor_class_match).cuda()
anchor_target_deltas = torch.from_numpy(anchor_target_deltas).float().cuda()
# compute losses.
class_loss, neg_anchor_ix = compute_class_loss(anchor_class_match, class_logits[b])
bbox_loss = compute_bbox_loss(anchor_target_deltas, pred_deltas[b], anchor_class_match)
# add negative anchors used for loss to results_dict for monitoring.
neg_anchors = mutils.clip_boxes_numpy(
self.np_anchors[np.argwhere(anchor_class_match.cpu().numpy() == -1)][neg_anchor_ix, 0],
img.shape[2:])
for n in neg_anchors:
box_results_list[b].append({'box_coords': n, 'box_type': 'neg_anchor'})
batch_class_loss += class_loss / img.shape[0]
batch_bbox_loss += bbox_loss / img.shape[0]
results_dict = get_results(self.cf, img.shape, detections, seg_logits, box_results_list)
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])
loss = batch_class_loss + batch_bbox_loss + (seg_loss_dice + seg_loss_ce) / 2
results_dict['torch_loss'] = loss
results_dict['monitor_values'] = {'loss': loss.item(), 'class_loss': batch_class_loss.item()}
results_dict['logger_string'] = \
"loss: {0:.2f}, class: {1:.2f}, bbox: {2:.2f}, seg dice: {3:.3f}, seg ce: {4:.3f}, mean pix. pr.: {5:.5f}"\
.format(loss.item(), batch_class_loss.item(), batch_bbox_loss.item(), seg_loss_dice.item(),
seg_loss_ce.item(), np.mean(results_dict['seg_preds']))
return results_dict
def test_forward(self, batch, **kwargs):
"""
test method. wrapper around forward pass of network without usage of any ground truth information.
prepares input data for processing and stores outputs in a dictionary.
:param batch: dictionary containing 'data'
:return: results_dict: dictionary with keys:
'boxes': list over batch elements. each batch element is a list of boxes. each box is a dictionary:
[[{box_0}, ... {box_n}], [{box_0}, ... {box_n}], ...]
'seg_preds': pixel-wise class predictions (b, 1, y, x, (z)) with values [0, ..., n_classes] for
retina_unet and dummy array for retina_net.
"""
img = batch['data']
img = torch.from_numpy(img).float().cuda()
detections, _, _, seg_logits = self.forward(img)
results_dict = get_results(self.cf, img.shape, detections, seg_logits)
return results_dict
def forward(self, img):
"""
forward pass of the model.
:param img: input img (b, c, y, x, (z)).
:return: rpn_pred_logits: (b, n_anchors, 2)
:return: rpn_pred_deltas: (b, n_anchors, (y, x, (z), log(h), log(w), (log(d))))
:return: batch_proposal_boxes: (b, n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ix)) only for monitoring/plotting.
:return: detections: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score)
:return: detection_masks: (n_final_detections, n_classes, y, x, (z)) raw molded masks as returned by mask-head.
"""
# Feature extraction
fpn_outs = self.Fpn(img)
seg_logits = self.final_conv(fpn_outs[0])
selected_fmaps = [fpn_outs[i + 1] for i in self.cf.pyramid_levels]
# Loop through pyramid layers
class_layer_outputs, bb_reg_layer_outputs = [], [] # list of lists
for p in selected_fmaps:
class_layer_outputs.append(self.Classifier(p))
bb_reg_layer_outputs.append(self.BBRegressor(p))
# Concatenate layer outputs
# Convert from list of lists of level outputs to list of lists
# of outputs across levels.
# e.g. [[a1, b1, c1], [a2, b2, c2]] => [[a1, a2], [b1, b2], [c1, c2]]
class_logits = list(zip(*class_layer_outputs))
class_logits = [torch.cat(list(o), dim=1) for o in class_logits][0]
bb_outputs = list(zip(*bb_reg_layer_outputs))
bb_outputs = [torch.cat(list(o), dim=1) for o in bb_outputs][0]
# merge batch_dimension and store info in batch_ixs for re-allocation.
batch_ixs = torch.arange(class_logits.shape[0]).unsqueeze(1).repeat(1, class_logits.shape[1]).view(-1).cuda()
flat_class_softmax = F.softmax(class_logits.view(-1, class_logits.shape[-1]), 1)
flat_bb_outputs = bb_outputs.view(-1, bb_outputs.shape[-1])
detections = refine_detections(self.anchors, flat_class_softmax, flat_bb_outputs, batch_ixs, self.cf)
return detections, class_logits, bb_outputs, seg_logits
diff --git a/models/ufrcnn.py b/models/ufrcnn.py
index 5f60e45..c04cba9 100644
--- a/models/ufrcnn.py
+++ b/models/ufrcnn.py
@@ -1,1273 +1,1276 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""
Parts are based on https://github.com/multimodallearning/pytorch-mask-rcnn
published under MIT license.
"""
import sys
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils
sys.path.append("..")
import utils.model_utils as mutils
import utils.exp_utils as utils
from custom_extensions.nms import nms
from custom_extensions.roi_align import roi_align
############################################################
# Networks on top of backbone
############################################################
class RPN(nn.Module):
"""
Region Proposal Network.
"""
def __init__(self, cf, conv):
super(RPN, self).__init__()
self.dim = conv.dim
self.conv_shared = conv(cf.end_filts, cf.n_rpn_features, ks=3, stride=cf.rpn_anchor_stride, pad=1, relu=cf.relu)
self.conv_class = conv(cf.n_rpn_features, 2 * len(cf.rpn_anchor_ratios), ks=1, stride=1, relu=None)
self.conv_bbox = conv(cf.n_rpn_features, 2 * self.dim * len(cf.rpn_anchor_ratios), ks=1, stride=1, relu=None)
def forward(self, x):
"""
:param x: input feature maps (b, in_channels, y, x, (z))
:return: rpn_class_logits (b, 2, n_anchors)
:return: rpn_probs_logits (b, 2, n_anchors)
:return: rpn_bbox (b, 2 * dim, n_anchors)
"""
# Shared convolutional base of the RPN.
x = self.conv_shared(x)
# Anchor Score. (batch, anchors per location * 2, y, x, (z)).
rpn_class_logits = self.conv_class(x)
# Reshape to (batch, 2, anchors)
axes = (0, 2, 3, 1) if self.dim == 2 else (0, 2, 3, 4, 1)
rpn_class_logits = rpn_class_logits.permute(*axes)
rpn_class_logits = rpn_class_logits.contiguous()
rpn_class_logits = rpn_class_logits.view(x.size()[0], -1, 2)
# Softmax on last dimension (fg vs. bg).
rpn_probs = F.softmax(rpn_class_logits, dim=2)
# Bounding box refinement. (batch, anchors_per_location * (y, x, (z), log(h), log(w), (log(d)), y, x, (z))
rpn_bbox = self.conv_bbox(x)
# Reshape to (batch, 2*dim, anchors)
rpn_bbox = rpn_bbox.permute(*axes)
rpn_bbox = rpn_bbox.contiguous()
rpn_bbox = rpn_bbox.view(x.size()[0], -1, self.dim * 2)
return [rpn_class_logits, rpn_probs, rpn_bbox]
class Classifier(nn.Module):
"""
Head network for classification and bounding box refinement. Performs RoiAlign, processes resulting features through a
shared convolutional base and finally branches off the classifier- and regression head.
"""
def __init__(self, cf, conv):
super(Classifier, self).__init__()
self.dim = conv.dim
self.in_channels = cf.end_filts
self.pool_size = cf.pool_size
self.pyramid_levels = cf.pyramid_levels
# instance_norm does not work with spatial dims (1, 1, (1))
norm = cf.norm if cf.norm != 'instance_norm' else None
self.conv1 = conv(cf.end_filts, cf.end_filts * 4, ks=self.pool_size, stride=1, norm=norm, relu=cf.relu)
self.conv2 = conv(cf.end_filts * 4, cf.end_filts * 4, ks=1, stride=1, norm=norm, relu=cf.relu)
self.linear_class = nn.Linear(cf.end_filts * 4, cf.head_classes)
self.linear_bbox = nn.Linear(cf.end_filts * 4, cf.head_classes * 2 * self.dim)
def forward(self, x, rois):
"""
:param x: input feature maps (b, in_channels, y, x, (z))
:param rois: normalized box coordinates as proposed by the RPN to be forwarded through
the second stage (n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ix). Proposals of all batch elements
have been merged to one vector, while the origin info has been stored for re-allocation.
:return: mrcnn_class_logits (n_proposals, n_head_classes)
:return: mrcnn_bbox (n_proposals, n_head_classes, 2 * dim) predicted corrections to be applied to proposals for refinement.
"""
x = pyramid_roi_align(x, rois, self.pool_size, self.pyramid_levels, self.dim)
x = self.conv1(x)
x = self.conv2(x)
x = x.view(-1, self.in_channels * 4)
mrcnn_class_logits = self.linear_class(x)
mrcnn_bbox = self.linear_bbox(x)
mrcnn_bbox = mrcnn_bbox.view(mrcnn_bbox.size()[0], -1, self.dim * 2)
return [mrcnn_class_logits, mrcnn_bbox]
class Mask(nn.Module):
"""
Head network for proposal-based mask segmentation. Performs RoiAlign, some convolutions and applies sigmoid on the
output logits to allow for overlapping classes.
"""
def __init__(self, cf, conv):
super(Mask, self).__init__()
self.pool_size = cf.mask_pool_size
self.pyramid_levels = cf.pyramid_levels
self.dim = conv.dim
self.conv1 = conv(cf.end_filts, cf.end_filts, ks=3, stride=1, pad=1, norm=cf.norm, relu=cf.relu)
self.conv2 = conv(cf.end_filts, cf.end_filts, ks=3, stride=1, pad=1, norm=cf.norm, relu=cf.relu)
self.conv3 = conv(cf.end_filts, cf.end_filts, ks=3, stride=1, pad=1, norm=cf.norm, relu=cf.relu)
self.conv4 = conv(cf.end_filts, cf.end_filts, ks=3, stride=1, pad=1, norm=cf.norm, relu=cf.relu)
if conv.dim == 2:
self.deconv = nn.ConvTranspose2d(cf.end_filts, cf.end_filts, kernel_size=2, stride=2)
else:
self.deconv = nn.ConvTranspose3d(cf.end_filts, cf.end_filts, kernel_size=2, stride=2)
self.relu = nn.ReLU(inplace=True) if cf.relu == 'relu' else nn.LeakyReLU(inplace=True)
self.conv5 = conv(cf.end_filts, cf.head_classes, ks=1, stride=1, relu=None)
self.sigmoid = nn.Sigmoid()
def forward(self, x, rois):
"""
:param x: input feature maps (b, in_channels, y, x, (z))
:param rois: normalized box coordinates as proposed by the RPN to be forwarded through
the second stage (n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ix). Proposals of all batch elements
have been merged to one vector, while the origin info has been stored for re-allocation.
:return: x: masks (n_sampled_proposals (n_detections in inference), n_classes, y, x, (z))
"""
x = pyramid_roi_align(x, rois, self.pool_size, self.pyramid_levels, self.dim)
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)
x = self.conv4(x)
x = self.relu(self.deconv(x))
x = self.conv5(x)
x = self.sigmoid(x)
return x
############################################################
# Loss Functions
############################################################
def compute_rpn_class_loss(rpn_match, rpn_class_logits, shem_poolsize):
"""
:param rpn_match: (n_anchors). [-1, 0, 1] for negative, neutral, and positive matched anchors.
:param rpn_class_logits: (n_anchors, 2). logits from RPN classifier.
:param shem_poolsize: int. factor of top-k candidates to draw from per negative sample
(stochastic-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.
"""
# filter out neutral anchors.
pos_indices = torch.nonzero(rpn_match == 1)
neg_indices = torch.nonzero(rpn_match == -1)
# loss for positive samples
if 0 not in pos_indices.size():
pos_indices = pos_indices.squeeze(1)
roi_logits_pos = rpn_class_logits[pos_indices]
pos_loss = F.cross_entropy(roi_logits_pos, torch.LongTensor([1] * pos_indices.shape[0]).cuda())
else:
pos_loss = torch.FloatTensor([0]).cuda()
# loss for negative samples: draw hard negative examples (SHEM)
# that match the number of positive samples, but at least 1.
if 0 not in neg_indices.size():
neg_indices = neg_indices.squeeze(1)
roi_logits_neg = rpn_class_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())
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_rpn_bbox_loss(rpn_target_deltas, rpn_pred_deltas, rpn_match):
"""
:param rpn_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 rpn_pred_deltas: predicted deltas from RPN. (b, n_anchors, (dy, dx, (dz), log(dh), log(dw), (log(dd))))
:param rpn_match: (n_anchors). [-1, 0, 1] for negative, neutral, and positive matched anchors.
:return: loss: torch 1D tensor.
"""
if 0 not in torch.nonzero(rpn_match == 1).size():
indices = torch.nonzero(rpn_match == 1).squeeze(1)
# Pick bbox deltas that contribute to the loss
rpn_pred_deltas = rpn_pred_deltas[indices]
# Trim target bounding box deltas to the same length as rpn_bbox.
target_deltas = rpn_target_deltas[:rpn_pred_deltas.size()[0], :]
# Smooth L1 loss
loss = F.smooth_l1_loss(rpn_pred_deltas, target_deltas)
else:
loss = torch.FloatTensor([0]).cuda()
return loss
def compute_mrcnn_class_loss(target_class_ids, pred_class_logits):
"""
:param target_class_ids: (n_sampled_rois) batch dimension was merged into roi dimension.
:param pred_class_logits: (n_sampled_rois, n_classes)
:return: loss: torch 1D tensor.
"""
if 0 not in target_class_ids.size():
loss = F.cross_entropy(pred_class_logits, target_class_ids.long())
else:
loss = torch.FloatTensor([0.]).cuda()
return loss
def compute_mrcnn_bbox_loss(mrcnn_target_deltas, mrcnn_pred_deltas, target_class_ids):
"""
:param mrcnn_target_deltas: (n_sampled_rois, (dy, dx, (dz), log(dh), log(dw), (log(dh)))
:param mrcnn_pred_deltas: (n_sampled_rois, n_classes, (dy, dx, (dz), log(dh), log(dw), (log(dh)))
:param target_class_ids: (n_sampled_rois)
:return: loss: torch 1D tensor.
"""
if 0 not in torch.nonzero(target_class_ids > 0).size():
positive_roi_ix = torch.nonzero(target_class_ids > 0)[:, 0]
positive_roi_class_ids = target_class_ids[positive_roi_ix].long()
target_bbox = mrcnn_target_deltas[positive_roi_ix, :].detach()
pred_bbox = mrcnn_pred_deltas[positive_roi_ix, positive_roi_class_ids, :]
loss = F.smooth_l1_loss(pred_bbox, target_bbox)
else:
loss = torch.FloatTensor([0]).cuda()
return loss
def compute_mrcnn_mask_loss(target_masks, pred_masks, target_class_ids):
"""
:param target_masks: (n_sampled_rois, y, x, (z)) A float32 tensor of values 0 or 1. Uses zero padding to fill array.
:param pred_masks: (n_sampled_rois, n_classes, y, x, (z)) float32 tensor with values between [0, 1].
:param target_class_ids: (n_sampled_rois)
:return: loss: torch 1D tensor.
"""
if 0 not in torch.nonzero(target_class_ids > 0).size():
# Only positive ROIs contribute to the loss. And only
# the class specific mask of each ROI.
positive_ix = torch.nonzero(target_class_ids > 0)[:, 0]
positive_class_ids = target_class_ids[positive_ix].long()
y_true = target_masks[positive_ix, :, :].detach()
y_pred = pred_masks[positive_ix, positive_class_ids, :, :]
loss = F.binary_cross_entropy(y_pred, y_true)
else:
loss = torch.FloatTensor([0]).cuda()
return loss
############################################################
# Helper Layers
############################################################
# def proposal_layer(rpn_pred_probs, rpn_pred_deltas, proposal_count, 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 detals 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_boxes: Proposals in normalized coordinates
# (b, proposal_count, (y1, x1, y2, x2, (z1), (z2)))
# :return: batch_out_proposals: Box coords + RPN foreground scores
# for monitoring/plotting (b, proposal_count, (y1, x1, y2, x2, (z1), (z2), score))
# """
# batch_scores = rpn_pred_probs[:, :, 1]
# batch_deltas = rpn_pred_deltas
# batch_anchors = anchors
# batch_normalized_boxes = []
# batch_out_proposals = []
#
# # loop over batch dimension.
# for ix in range(batch_scores.shape[0]):
#
# scores = batch_scores[ix]
# deltas = batch_deltas[ix]
# anchors = batch_anchors.clone()
# # norm deltas
# std_dev = torch.from_numpy(cf.rpn_bbox_std_dev[None]).float().cuda()
# deltas = deltas * std_dev
#
# # 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, :]
# anchors = anchors[order, :]
#
# # apply deltas to anchors to get refined anchors and filter with non-maximum surpression.
# if batch_deltas.shape[-1] == 4:
# boxes = mutils.apply_box_deltas_2D(anchors, deltas)
# boxes = mutils.clip_boxes_2D(boxes, cf.window)
# keep = nms_2D(torch.cat((boxes, scores.unsqueeze(1)), 1), cf.rpn_nms_threshold)
# norm = torch.from_numpy(cf.scale).float().cuda()
#
# else:
# boxes = mutils.apply_box_deltas_3D(anchors, deltas)
# boxes = mutils.clip_boxes_3D(boxes, cf.window)
# keep = nms_3D(torch.cat((boxes, scores.unsqueeze(1)), 1), cf.rpn_nms_threshold)
# norm = torch.from_numpy(cf.scale).float().cuda()
#
# 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
# # add back batch dimension
# batch_normalized_boxes.append(normalized_boxes.unsqueeze(0))
#
# batch_normalized_boxes = torch.cat(batch_normalized_boxes)
# batch_out_proposals = np.array(batch_out_proposals)
# return batch_normalized_boxes, batch_out_proposals
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 = mutils.apply_box_deltas_2D(anchors[order, :], deltas)
boxes = mutils.clip_boxes_2D(boxes, cf.window)
else:
boxes = mutils.apply_box_deltas_3D(anchors[order, :], deltas)
boxes = mutils.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(normalized_boxes.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 + mutils.log2(torch.sqrt(h*w))).round().int().clamp(pyramid_levels[0], pyramid_levels[-1])
# # if Pyramid contains additional level P6, adapt the roi_level assignemnt 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 = []
# 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()
#
# # Crop and Resize
# # From Mask R-CNN paper: "We sample four regular locations, so
# # that we can evaluate either max or average pooling. In fact,
# # interpolating only a single value at each bin center (without
# # pooling) is nearly as effective."
# #
# # Here we use the simplified approach of a single value per bin,
# # which is how is done in tf.crop_and_resize()
# #
# # Also fixed a bug from original implementation, reported in:
# # https://hackernoon.com/how-tensorflows-tf-image-resize-stole-60-days-of-my-life-aba5eb093f35
#
# if len(pool_size) == 2:
# pooled_features = ra2D(pool_size[0], pool_size[1], 0)(feature_maps[level_ix], level_boxes, ind)
# else:
# pooled_features = ra3D(pool_size[0], pool_size[1], pool_size[2], 0)(feature_maps[level_ix], level_boxes, ind)
#
# 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 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 + 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 detection_target_layer(batch_proposals, batch_mrcnn_class_scores, batch_gt_class_ids, batch_gt_boxes, cf):
"""
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 are 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".)
: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 batch_mrcnn_class_scores: (n_proposals, n_classes)
:param batch_gt_class_ids: list over batch elements. Each element is a list over the corresponding roi target labels.
: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)
: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
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_class_ids = []
std_dev = torch.from_numpy(cf.bbox_std_dev).float().cuda()
# loop over batch and get positive and negative sample rois.
for b in range(len(batch_gt_class_ids)):
gt_class_ids = torch.from_numpy(batch_gt_class_ids[b]).int().cuda()
if np.any(batch_gt_class_ids[b] > 0): # skip roi selection for no gt images.
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 0 not in gt_boxes.size():
if gt_boxes.shape[1] == 4:
overlaps = mutils.bbox_overlaps_2D(proposals, gt_boxes)
else:
overlaps = mutils.bbox_overlaps_3D(proposals, gt_boxes)
# Determine postive 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 0 not 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]
# Compute bbox refinement targets for positive ROIs
deltas = mutils.box_refinement(positive_rois, roi_gt_boxes)
deltas /= std_dev
sample_positive_indices.append(batch_element_indices[positive_indices])
sample_deltas.append(deltas)
sample_class_ids.append(roi_gt_class_ids)
positive_count += positive_samples
else:
positive_samples = 0
# Negative ROIs. Add enough to maintain positive:negative ratio, but at least 1. Sample via SHEM.
if 0 not 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_probs_neg = batch_mrcnn_class_scores[batch_element_indices[negative_indices]]
raw_sampled_indices = mutils.shem(roi_probs_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_class_ids = torch.cat(sample_class_ids)
# 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).int().cuda()
target_class_ids = torch.cat([target_class_ids, zeros], dim=0)
zeros = torch.zeros(negative_count, cf.dim * 2).cuda()
target_deltas = torch.cat([target_deltas, 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)
zeros = torch.zeros(negative_count).int().cuda()
target_class_ids = zeros
zeros = torch.zeros(negative_count, cf.dim * 2).cuda()
target_deltas = zeros
else:
sample_indices = torch.LongTensor().cuda()
target_class_ids = torch.IntTensor().cuda()
target_deltas = torch.FloatTensor().cuda()
return sample_indices, target_class_ids, target_deltas
############################################################
# Output Handler
############################################################
# def refine_detections(rois, probs, deltas, batch_ixs, cf):
# """
# Refine classified proposals, filter overlaps and return final detections.
#
# :param rois: (n_proposals, 2 * dim) normalized boxes as proposed by RPN. n_proposals = batch_size * POST_NMS_ROIS
# :param probs: (n_proposals, n_classes) softmax probabilities for all rois as predicted by mrcnn classifier.
# :param deltas: (n_proposals, n_classes, 2 * dim) box refinement deltas as predicted by mrcnn bbox regressor.
# :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))
# """
# # 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()
#
# rois = rois.repeat(fg_classes, 1)
# probs = probs.repeat(fg_classes, 1)
# deltas = deltas.repeat(fg_classes, 1, 1)
# batch_ixs = batch_ixs.repeat(fg_classes)
#
# # get class-specific scores and bounding box deltas
# idx = torch.arange(class_ids.size()[0]).long().cuda()
# class_scores = probs[idx, class_ids]
# deltas_specific = deltas[idx, class_ids]
# batch_ixs = batch_ixs[idx]
#
# # 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(rois, deltas_specific * std_dev) * scale if cf.dim == 2 else \
# mutils.apply_box_deltas_3D(rois, deltas_specific * 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)
# refined_rois = torch.round(refined_rois)
#
# # filter out low confidence boxes
# keep = idx
# keep_bool = (class_scores >= cf.model_min_confidence)
# if 0 not 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(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, :]
#
# 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)
#
# # 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 mutils.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 mutils.unique1d(torch.cat((batch_keep, b_keep)))
#
# keep = batch_keep
#
# else:
# keep = torch.tensor([0]).long().cuda()
#
# # arrange output
# result = torch.cat((refined_rois[keep],
# batch_ixs[keep].unsqueeze(1),
# class_ids[keep].unsqueeze(1).float(),
# class_scores[keep].unsqueeze(1)), dim=1)
#
# return result
def refine_detections(cf, batch_ixs, rois, deltas, scores):
"""
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.
: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)
# 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]
# 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(rois, deltas_specific * std_dev) * scale if cf.dim == 2 else \
mutils.apply_box_deltas_3D(rois, deltas_specific * std_dev) * scale
# round and cast to int since we're dealing with pixels now
refined_rois = mutils.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(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, :]
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 mutils.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 mutils.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)]
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 get_results(cf, img_shape, detections, seg_logits, box_results_list=None):
"""
Restores batch dimension of merged detections, unmolds detections, creates and fills results dict.
:param img_shape:
:param detections: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score)
:param detection_masks: (n_final_detections, n_classes, y, x, (z)) raw molded masks as returned by mask-head.
:param box_results_list: None or list of output boxes for monitoring/plotting.
each element is a list of boxes per batch element.
:param return_masks: boolean. If True, full resolution masks are returned for all proposals (speed trade-off).
: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()
# restore batch dimension of merged detections using the batch_ix info.
batch_ixs = detections[:, cf.dim*2]
detections = [detections[batch_ixs == ix] for ix in range(img_shape[0])]
# for test_forward, where no previous list exists.
if box_results_list is None:
box_results_list = [[] for _ in range(img_shape[0])]
seg_preds = []
# loop over batch and unmold detections.
for ix in range(img_shape[0]):
if 0 not in detections[ix].shape:
boxes = detections[ix][:, :2 * cf.dim].astype(np.int32)
class_ids = detections[ix][:, 2 * cf.dim + 1].astype(np.int32)
scores = detections[ix][:, 2 * cf.dim + 2]
# Filter out detections with zero area. Often only happens in early
# stages of training when the network weights are still a bit random.
if cf.dim == 2:
exclude_ix = np.where((boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) <= 0)[0]
else:
exclude_ix = np.where(
(boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 5] - boxes[:, 4]) <= 0)[0]
if exclude_ix.shape[0] > 0:
boxes = np.delete(boxes, exclude_ix, axis=0)
class_ids = np.delete(class_ids, exclude_ix, axis=0)
scores = np.delete(scores, exclude_ix, axis=0)
# add final perdictions to results.
if 0 not in boxes.shape:
for ix2, score in enumerate(scores):
if score >= cf.model_min_confidence:
box_results_list[ix].append({'box_coords': boxes[ix2], 'box_score': score,
'box_type': 'det', 'box_pred_class_id': class_ids[ix2]})
# create and fill results dictionary.
results_dict = {'boxes': box_results_list}
if seg_logits is None:
# output dummy segmentation for retina_net.
results_dict['seg_preds'] = np.zeros(img_shape)[:, 0][:, np.newaxis]
else:
# output label maps for retina_unet.
results_dict['seg_preds'] = F.softmax(seg_logits, 1).argmax(1).cpu().data.numpy()[:, np.newaxis].astype('uint8')
return results_dict
############################################################
# Mask R-CNN Class
############################################################
class net(nn.Module):
def __init__(self, cf, logger):
super(net, self).__init__()
self.cf = cf
self.logger = logger
self.build()
if self.cf.weight_init is not None:
logger.info("using pytorch weight init of type {}".format(self.cf.weight_init))
mutils.initialize_weights(self)
else:
logger.info("using default pytorch weight init")
def build(self):
"""Build Mask R-CNN architecture."""
# Image size must be dividable by 2 multiple times.
h, w = self.cf.patch_size[:2]
if h / 2**5 != int(h / 2**5) or w / 2**5 != int(w / 2**5):
raise Exception("Image size must be dividable by 2 at least 5 times "
"to avoid fractions when downscaling and upscaling."
"For example, use 256, 320, 384, 448, 512, ... etc. ")
# instanciate abstract multi dimensional conv class and backbone class.
conv = mutils.NDConvGenerator(self.cf.dim)
backbone = utils.import_module('bbone', self.cf.backbone_path)
# build Anchors, FPN, RPN, Classifier / Bbox-Regressor -head, Mask-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=True)
self.rpn = RPN(self.cf, conv)
self.classifier = Classifier(self.cf, conv)
self.mask = Mask(self.cf, conv)
self.final_conv = conv(self.cf.end_filts, self.cf.num_seg_classes, ks=1, pad=0, norm=self.cf.norm, relu=None)
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': 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.
"""
img = batch['data']
- gt_class_ids = batch['roi_labels']
+ if "roi_labels" in batch.keys():
+ raise Exception("Key for roi-wise class targets changed in v0.1.0 from 'roi_labels' to 'class_target'.\n"
+ "If you use DKFZ's batchgenerators, please make sure you run version >= 0.20.1.")
+ gt_class_ids = batch['class_target']
gt_boxes = batch['bb_target']
axes = (0, 2, 3, 1) if self.cf.dim == 2 else (0, 2, 3, 4, 1)
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()
batch_rpn_class_loss = torch.FloatTensor([0]).cuda()
batch_rpn_bbox_loss = torch.FloatTensor([0]).cuda()
# list of output boxes for monitoring/plotting. each element is a list of boxes per batch element.
box_results_list = [[] for _ in range(img.shape[0])]
#forward passes. 1. general forward pass, where no activations are saved in second stage (for performance
# monitoring and loss sampling). 2. second stage forward pass of sampled rois with stored activations for backprop.
rpn_class_logits, rpn_pred_deltas, proposal_boxes, detections, seg_logits = self.forward(img)
mrcnn_class_logits, mrcnn_pred_deltas, target_class_ids, mrcnn_target_deltas, \
sample_proposals = self.loss_samples_forward(gt_class_ids, gt_boxes)
# loop over batch
for b in range(img.shape[0]):
if len(gt_boxes[b]) > 0:
# add gt boxes to output list for monitoring.
for ix in range(len(gt_boxes[b])):
box_results_list[b].append({'box_coords': batch['bb_target'][b][ix],
- 'box_label': batch['roi_labels'][b][ix], 'box_type': 'gt'})
+ 'box_label': batch['class_target'][b][ix], 'box_type': 'gt'})
# match gt boxes with anchors to generate targets for RPN losses.
rpn_match, rpn_target_deltas = mutils.gt_anchor_matching(self.cf, self.np_anchors, gt_boxes[b])
# add positive anchors used for loss to output list for monitoring.
pos_anchors = mutils.clip_boxes_numpy(self.np_anchors[np.argwhere(rpn_match == 1)][:, 0], img.shape[2:])
for p in pos_anchors:
box_results_list[b].append({'box_coords': p, 'box_type': 'pos_anchor'})
else:
rpn_match = np.array([-1]*self.np_anchors.shape[0])
rpn_target_deltas = np.array([0])
rpn_match_gpu = torch.from_numpy(rpn_match).cuda()
rpn_target_deltas = torch.from_numpy(rpn_target_deltas).float().cuda()
# compute RPN losses.
rpn_class_loss, neg_anchor_ix = compute_rpn_class_loss(rpn_match_gpu, rpn_class_logits[b], self.cf.shem_poolsize)
rpn_bbox_loss = compute_rpn_bbox_loss(rpn_target_deltas, rpn_pred_deltas[b], rpn_match_gpu)
batch_rpn_class_loss += rpn_class_loss / img.shape[0]
batch_rpn_bbox_loss += rpn_bbox_loss / img.shape[0]
# add negative anchors used for loss to output list for monitoring.
neg_anchors = mutils.clip_boxes_numpy(self.np_anchors[rpn_match == -1][neg_anchor_ix], img.shape[2:])
for n in neg_anchors:
box_results_list[b].append({'box_coords': n, 'box_type': 'neg_anchor'})
# add highest scoring proposals to output list for monitoring.
rpn_proposals = proposal_boxes[b][proposal_boxes[b, :, -1].argsort()][::-1]
for r in rpn_proposals[:self.cf.n_plot_rpn_props, :-1]:
box_results_list[b].append({'box_coords': r, 'box_type': 'prop'})
# add positive and negative roi samples used for mrcnn losses to output list for monitoring.
if 0 not in sample_proposals.shape:
rois = mutils.clip_to_window(self.cf.window, sample_proposals).cpu().data.numpy()
for ix, r in enumerate(rois):
box_results_list[int(r[-1])].append({'box_coords': r[:-1] * self.cf.scale,
'box_type': 'pos_class' if target_class_ids[ix] > 0 else 'neg_class'})
batch_rpn_class_loss = batch_rpn_class_loss
batch_rpn_bbox_loss = batch_rpn_bbox_loss
# compute mrcnn losses.
mrcnn_class_loss = compute_mrcnn_class_loss(target_class_ids, mrcnn_class_logits)
mrcnn_bbox_loss = compute_mrcnn_bbox_loss(mrcnn_target_deltas, mrcnn_pred_deltas, target_class_ids)
# mrcnn can be run without pixelwise annotations available (Faster R-CNN mode).
# In this case, the mask_loss is taken out of training.
# if not self.cf.frcnn_mode:
# mrcnn_mask_loss = compute_mrcnn_mask_loss(target_mask, mrcnn_pred_mask, target_class_ids)
# else:
# mrcnn_mask_loss = torch.FloatTensor([0]).cuda()
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])
loss = batch_rpn_class_loss + batch_rpn_bbox_loss + mrcnn_class_loss + mrcnn_bbox_loss + (seg_loss_dice + seg_loss_ce) / 2
# monitor RPN performance: detection count = the number of correctly matched proposals per fg-class.
dcount = [list(target_class_ids.cpu().data.numpy()).count(c) for c in np.arange(self.cf.head_classes)[1:]]
# run unmolding of predictions for monitoring and merge all results to one dictionary.
results_dict = get_results(self.cf, img.shape, detections, seg_logits, box_results_list)
results_dict['torch_loss'] = loss
results_dict['monitor_values'] = {'loss': loss.item(), 'class_loss': mrcnn_class_loss.item()}
results_dict['logger_string'] = "loss: {0:.2f}, rpn_class: {1:.2f}, rpn_bbox: {2:.2f}, mrcnn_class: {3:.2f}, " \
"mrcnn_bbox: {4:.2f}, dice_loss: {5:.2f}, dcount {6}"\
.format(loss.item(), batch_rpn_class_loss.item(), batch_rpn_bbox_loss.item(), mrcnn_class_loss.item(),
mrcnn_bbox_loss.item(), seg_loss_dice.item(), dcount)
return results_dict
def test_forward(self, batch, return_masks=True):
"""
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 return_masks: boolean. If True, full resolution masks are returned for all proposals (speed trade-off).
: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 = batch['data']
img = torch.from_numpy(img).float().cuda()
_, _, _, detections, seg_logits = self.forward(img)
results_dict = get_results(self.cf, img.shape, detections, seg_logits)
return results_dict
def forward(self, img, is_training=True):
"""
:param img: input images (b, c, y, x, (z)).
:return: rpn_pred_logits: (b, n_anchors, 2)
:return: rpn_pred_deltas: (b, n_anchors, (y, x, (z), log(h), log(w), (log(d))))
:return: batch_proposal_boxes: (b, n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ix)) only for monitoring/plotting.
:return: detections: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score)
:return: detection_masks: (n_final_detections, n_classes, y, x, (z)) raw molded masks as returned by mask-head.
"""
# extract features.
fpn_outs = self.fpn(img)
seg_logits = self.final_conv(fpn_outs[0])
rpn_feature_maps = [fpn_outs[i + 1] for i in self.cf.pyramid_levels]
self.mrcnn_feature_maps = rpn_feature_maps
# loop through pyramid layers and apply RPN.
layer_outputs = [] # list of lists
for p in rpn_feature_maps:
layer_outputs.append(self.rpn(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]]
outputs = list(zip(*layer_outputs))
outputs = [torch.cat(list(o), dim=1) for o in outputs]
rpn_pred_logits, rpn_pred_probs, rpn_pred_deltas = outputs
# generate proposals: apply predicted deltas to anchors and filter by foreground scores from RPN classifier.
proposal_count = self.cf.post_nms_rois_training if is_training else self.cf.post_nms_rois_inference
batch_rpn_rois, batch_proposal_boxes = refine_proposals(rpn_pred_probs, rpn_pred_deltas, proposal_count, self.anchors, self.cf)
# merge batch dimension of proposals while storing allocation info in coordinate dimension.
batch_ixs = torch.from_numpy(np.repeat(np.arange(batch_rpn_rois.shape[0]), batch_rpn_rois.shape[1])).float().cuda()
rpn_rois = batch_rpn_rois.view(-1, batch_rpn_rois.shape[2])
self.rpn_rois_batch_info = torch.cat((rpn_rois, batch_ixs.unsqueeze(1)), dim=1)
# this is the first of two forward passes in the second stage, where no activations are stored for backprop.
# here, all proposals are forwarded (with virtual_batch_size = batch_size * post_nms_rois.)
# for inference/monitoring as well as sampling of rois for the loss functions.
# processed in chunks of roi_chunk_size to re-adjust to gpu-memory.
chunked_rpn_rois = self.rpn_rois_batch_info.split(self.cf.roi_chunk_size)
class_logits_list, bboxes_list = [], []
with torch.no_grad():
for chunk in chunked_rpn_rois:
chunk_class_logits, chunk_bboxes = self.classifier(self.mrcnn_feature_maps, chunk)
class_logits_list.append(chunk_class_logits)
bboxes_list.append(chunk_bboxes)
batch_mrcnn_class_logits = torch.cat(class_logits_list, 0)
batch_mrcnn_bbox = torch.cat(bboxes_list, 0)
self.batch_mrcnn_class_scores = F.softmax(batch_mrcnn_class_logits, dim=1)
# refine classified proposals, filter and return final detections.
detections = refine_detections(self.cf, batch_ixs, rpn_rois, batch_mrcnn_bbox, self.batch_mrcnn_class_scores)
return [rpn_pred_logits, rpn_pred_deltas, batch_proposal_boxes, detections, seg_logits]
def loss_samples_forward(self, batch_gt_class_ids, batch_gt_boxes):
"""
this is the second forward pass through the second stage (features from stage one are re-used).
samples few rois in detection_target_layer and forwards only those for loss computation.
:param batch_gt_class_ids: list over batch elements. Each element is a list over the corresponding roi target labels.
: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)
:return: sample_logits: (n_sampled_rois, n_classes) predicted class scores.
:return: sample_boxes: (n_sampled_rois, n_classes, 2 * dim) predicted corrections to be applied to proposals for refinement.
:return: sample_mask: (n_sampled_rois, n_classes, y, x, (z)) predicted masks per class and proposal.
:return: sample_target_class_ids: (n_sampled_rois) target class labels of sampled proposals.
:return: sample_target_deltas: (n_sampled_rois, 2 * dim) target deltas of sampled proposals for box refinement.
:return: sample_target_masks: (n_sampled_rois, y, x, (z)) target masks of sampled proposals.
:return: sample_proposals: (n_sampled_rois, 2 * dim) RPN output for sampled proposals. only for monitoring/plotting.
"""
# sample rois for loss and get corresponding targets for all Mask R-CNN head network losses.
sample_ix, sample_target_class_ids, sample_target_deltas = \
detection_target_layer(self.rpn_rois_batch_info, self.batch_mrcnn_class_scores,
batch_gt_class_ids, batch_gt_boxes, self.cf)
# re-use feature maps and RPN output from first forward pass.
sample_proposals = self.rpn_rois_batch_info[sample_ix]
if 0 not in sample_proposals.size():
sample_logits, sample_boxes = self.classifier(self.mrcnn_feature_maps, sample_proposals)
else:
sample_logits = torch.FloatTensor().cuda()
sample_boxes = torch.FloatTensor().cuda()
return [sample_logits, sample_boxes, sample_target_class_ids, sample_target_deltas, sample_proposals]
diff --git a/plotting.py b/plotting.py
index a5b3565..ccda616 100644
--- a/plotting.py
+++ b/plotting.py
@@ -1,295 +1,297 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
+from typing import Union
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import numpy as np
import os
from copy import deepcopy
def suppress_axes_lines(ax):
"""
:param ax: pyplot axes object
"""
ax.axes.get_xaxis().set_ticks([])
ax.axes.get_yaxis().set_ticks([])
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['bottom'].set_visible(False)
ax.spines['left'].set_visible(False)
return
-def plot_batch_prediction(batch, results_dict, cf, outfile=None, suptitle=None):
+def plot_batch_prediction(batch: dict, results_dict: dict, cf, outfile: Union[str, None]=None,
+ suptitle: Union[str, None]=None):
"""
plot the input images, ground truth annotations, and output predictions of a batch. If 3D batch, plots a 2D projection
of one randomly sampled element (patient) in the batch. Since plotting all slices of patient volume blows up costs of
time and space, only a section containing a randomly sampled ground truth annotation is plotted.
:param batch: dict with keys: 'data' (input image), 'seg' (pixelwise annotations), 'pid'
:param results_dict: list over batch element. Each element is a list of boxes (prediction and ground truth),
where every box is a dictionary containing box_coords, box_score and box_type.
"""
if outfile is None:
outfile = os.path.join(cf.plot_dir, 'pred_example_{}.png'.format(cf.fold))
data = batch['data']
segs = batch['seg']
pids = batch['pid']
# for 3D, repeat pid over batch elements.
if len(set(pids)) == 1:
pids = [pids] * data.shape[0]
seg_preds = results_dict['seg_preds']
roi_results = deepcopy(results_dict['boxes'])
# Randomly sampled one patient of batch and project data into 2D slices for plotting.
if cf.dim == 3:
patient_ix = np.random.choice(data.shape[0])
data = np.transpose(data[patient_ix], axes=(3, 0, 1, 2))
# select interesting foreground section to plot.
gt_boxes = [box['box_coords'] for box in roi_results[patient_ix] if box['box_type'] == 'gt']
if len(gt_boxes) > 0:
z_cuts = [np.max((int(gt_boxes[0][4]) - 5, 0)), np.min((int(gt_boxes[0][5]) + 5, data.shape[0]))]
else:
z_cuts = [data.shape[0]//2 - 5, int(data.shape[0]//2 + np.min([10, data.shape[0]//2]))]
p_roi_results = roi_results[patient_ix]
roi_results = [[] for _ in range(data.shape[0])]
# iterate over cubes and spread across slices.
for box in p_roi_results:
b = box['box_coords']
# dismiss negative anchor slices.
slices = np.round(np.unique(np.clip(np.arange(b[4], b[5] + 1), 0, data.shape[0]-1)))
for s in slices:
roi_results[int(s)].append(box)
roi_results[int(s)][-1]['box_coords'] = b[:4]
roi_results = roi_results[z_cuts[0]: z_cuts[1]]
data = data[z_cuts[0]: z_cuts[1]]
segs = np.transpose(segs[patient_ix], axes=(3, 0, 1, 2))[z_cuts[0]: z_cuts[1]]
seg_preds = np.transpose(seg_preds[patient_ix], axes=(3, 0, 1, 2))[z_cuts[0]: z_cuts[1]]
pids = [pids[patient_ix]] * data.shape[0]
try:
# all dimensions except for the 'channel-dimension' are required to match
for i in [0, 2, 3]:
assert data.shape[i] == segs.shape[i] == seg_preds.shape[i]
except:
raise Warning('Shapes of arrays to plot not in agreement!'
'Shapes {} vs. {} vs {}'.format(data.shape, segs.shape, seg_preds.shape))
show_arrays = np.concatenate([data, segs, seg_preds, data[:, 0][:, None]], axis=1).astype(float)
approx_figshape = (4 * show_arrays.shape[0], 4 * show_arrays.shape[1])
fig = plt.figure(figsize=approx_figshape)
gs = gridspec.GridSpec(show_arrays.shape[1] + 1, show_arrays.shape[0])
gs.update(wspace=0.1, hspace=0.1)
for b in range(show_arrays.shape[0]):
for m in range(show_arrays.shape[1]):
ax = plt.subplot(gs[m, b])
suppress_axes_lines(ax)
if m < show_arrays.shape[1]:
arr = show_arrays[b, m]
if m < data.shape[1] or m == show_arrays.shape[1] - 1:
if b == 0:
ax.set_ylabel("Input" + (" + GT & Pred Box" if m == show_arrays.shape[1] - 1 else ""))
cmap = 'gray'
vmin = None
vmax = None
else:
cmap = None
vmin = 0
vmax = cf.num_seg_classes - 1
if m == 0:
plt.title('{}'.format(pids[b][:10]), fontsize=20)
plt.imshow(arr, cmap=cmap, vmin=vmin, vmax=vmax)
if m >= (data.shape[1]):
if b == 0:
if m == data.shape[1]:
ax.set_ylabel("GT Box & Seg")
if m == data.shape[1]+1:
ax.set_ylabel("GT Box + Pred Seg & Box")
for box in roi_results[b]:
if box['box_type'] != 'patient_tn_box': # don't plot true negative dummy boxes.
coords = box['box_coords']
if box['box_type'] == 'det':
# dont plot background preds or low confidence boxes.
if box['box_pred_class_id'] > 0 and box['box_score'] > 0.1:
plot_text = True
score = np.max(box['box_score'])
score_text = '{}|{:.0f}'.format(box['box_pred_class_id'], score*100)
# if prob detection: plot only boxes from correct sampling instance.
if 'sample_id' in box.keys() and int(box['sample_id']) != m - data.shape[1] - 2:
continue
# if prob detection: plot reconstructed boxes only in corresponding line.
if not 'sample_id' in box.keys() and m != data.shape[1] + 1:
continue
score_font_size = 7
text_color = 'w'
text_x = coords[1] + 10*(box['box_pred_class_id'] -1) #avoid overlap of scores in plot.
text_y = coords[2] + 5
else:
continue
elif box['box_type'] == 'gt':
plot_text = True
score_text = int(box['box_label'])
score_font_size = 7
text_color = 'r'
text_x = coords[1]
text_y = coords[0] - 1
else:
plot_text = False
color_var = 'extra_usage' if 'extra_usage' in list(box.keys()) else 'box_type'
color = cf.box_color_palette[box[color_var]]
plt.plot([coords[1], coords[3]], [coords[0], coords[0]], color=color, linewidth=1, alpha=1) # up
plt.plot([coords[1], coords[3]], [coords[2], coords[2]], color=color, linewidth=1, alpha=1) # down
plt.plot([coords[1], coords[1]], [coords[0], coords[2]], color=color, linewidth=1, alpha=1) # left
plt.plot([coords[3], coords[3]], [coords[0], coords[2]], color=color, linewidth=1, alpha=1) # right
if plot_text:
plt.text(text_x, text_y, score_text, fontsize=score_font_size, color=text_color)
if suptitle is not None:
plt.suptitle(suptitle, fontsize=22)
try:
plt.savefig(outfile)
except:
raise Warning('failed to save plot.')
plt.close(fig)
class TrainingPlot_2Panel():
# todo remove since replaced by tensorboard?
def __init__(self, cf):
self.file_name = cf.plot_dir + '/monitor_{}'.format(cf.fold)
self.exp_name = cf.fold_dir
self.do_validation = cf.do_validation
self.separate_values_dict = cf.assign_values_to_extra_figure
self.figure_list = []
for n in range(cf.n_monitoring_figures):
self.figure_list.append(plt.figure(figsize=(10, 6)))
self.figure_list[-1].ax1 = plt.subplot(111)
self.figure_list[-1].ax1.set_xlabel('epochs')
self.figure_list[-1].ax1.set_ylabel('loss / metrics')
self.figure_list[-1].ax1.set_xlim(0, cf.num_epochs)
self.figure_list[-1].ax1.grid()
self.figure_list[0].ax1.set_ylim(0, 1.5)
self.color_palette = ['b', 'c', 'r', 'purple', 'm', 'y', 'k', 'tab:gray']
def update_and_save(self, metrics, epoch):
for figure_ix in range(len(self.figure_list)):
fig = self.figure_list[figure_ix]
detection_monitoring_plot(fig.ax1, metrics, self.exp_name, self.color_palette, epoch, figure_ix,
self.separate_values_dict,
self.do_validation)
fig.savefig(self.file_name + '_{}'.format(figure_ix))
def detection_monitoring_plot(ax1, metrics, exp_name, color_palette, epoch, figure_ix, separate_values_dict, do_validation):
# todo remove since replaced by tensorboard?
monitor_values_keys = metrics['train']['monitor_values'][1][0].keys()
separate_values = [v for fig_ix in separate_values_dict.values() for v in fig_ix]
if figure_ix == 0:
plot_keys = [ii for ii in monitor_values_keys if ii not in separate_values]
plot_keys += [k for k in metrics['train'].keys() if k != 'monitor_values']
else:
plot_keys = separate_values_dict[figure_ix]
x = np.arange(1, epoch + 1)
for kix, pk in enumerate(plot_keys):
if pk in metrics['train'].keys():
y_train = metrics['train'][pk][1:]
if do_validation:
y_val = metrics['val'][pk][1:]
else:
y_train = [np.mean([er[pk] for er in metrics['train']['monitor_values'][e]]) for e in x]
if do_validation:
y_val = [np.mean([er[pk] for er in metrics['val']['monitor_values'][e]]) for e in x]
ax1.plot(x, y_train, label='train_{}'.format(pk), linestyle='--', color=color_palette[kix])
if do_validation:
ax1.plot(x, y_val, label='val_{}'.format(pk), linestyle='-', color=color_palette[kix])
if epoch == 1:
box = ax1.get_position()
ax1.set_position([box.x0, box.y0, box.width * 0.8, box.height])
ax1.legend(loc='center left', bbox_to_anchor=(1, 0.5))
ax1.set_title(exp_name)
-def plot_prediction_hist(label_list, pred_list, type_list, outfile):
+def plot_prediction_hist(label_list: list, pred_list: list, type_list: list, outfile: str):
"""
plot histogram of predictions for a specific class.
:param label_list: list of 1s and 0s specifying whether prediction is a true positive match (1) or a false positive (0).
False negatives (missed ground truth objects) are artificially added predictions with score 0 and label 1.
:param pred_list: list of prediction-scores.
:param type_list: list of prediction-types for stastic-info in title.
"""
preds = np.array(pred_list)
labels = np.array(label_list)
title = outfile.split('/')[-1] + ' count:{}'.format(len(label_list))
plt.figure()
plt.yscale('log')
if 0 in labels:
plt.hist(preds[labels == 0], alpha=0.3, color='g', range=(0, 1), bins=50, label='false pos.')
if 1 in labels:
plt.hist(preds[labels == 1], alpha=0.3, color='b', range=(0, 1), bins=50, label='true pos. (false neg. @ score=0)')
if type_list is not None:
fp_count = type_list.count('det_fp')
fn_count = type_list.count('det_fn')
tp_count = type_list.count('det_tp')
pos_count = fn_count + tp_count
title += ' tp:{} fp:{} fn:{} pos:{}'. format(tp_count, fp_count, fn_count, pos_count)
plt.legend()
plt.title(title)
plt.xlabel('confidence score')
plt.ylabel('log n')
plt.savefig(outfile)
plt.close()
-def plot_stat_curves(stats, outfile):
+def plot_stat_curves(stats: list, outfile: str):
for c in ['roc', 'prc']:
plt.figure()
for s in stats:
- if s[c] is not None:
+ if not (isinstance(s[c], float) and np.isnan(s[c])):
plt.plot(s[c][0], s[c][1], label=s['name'] + '_' + c)
plt.title(outfile.split('/')[-1] + '_' + c)
plt.legend(loc=3 if c == 'prc' else 4)
plt.xlabel('precision' if c == 'prc' else '1-spec.')
plt.ylabel('recall')
plt.savefig(outfile + '_' + c)
plt.close()
diff --git a/predictor.py b/predictor.py
index a68b405..aa2131a 100644
--- a/predictor.py
+++ b/predictor.py
@@ -1,876 +1,884 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import os
import numpy as np
import torch
from scipy.stats import norm
from collections import OrderedDict
from multiprocessing import Pool
import pickle
+from copy import deepcopy
import pandas as pd
import utils.exp_utils as utils
from plotting import plot_batch_prediction
class Predictor:
"""
Prediction pipeline:
- receives a patched patient image (n_patches, c, y, x, (z)) from patient data loader.
- forwards patches through model in chunks of batch_size. (method: batch_tiling_forward)
- unmolds predictions (boxes and segmentations) to original patient coordinates. (method: spatial_tiling_forward)
Ensembling (mode == 'test'):
- for inference, forwards 4 mirrored versions of image to through model and unmolds predictions afterwards
accordingly (method: data_aug_forward)
- for inference, loads multiple parameter-sets of the trained model corresponding to different epochs. for each
parameter-set loops over entire test set, runs prediction pipeline for each patient. (method: predict_test_set)
Consolidation of predictions:
- consolidates a patient's predictions (boxes, segmentations) collected over patches, data_aug- and temporal ensembling,
performs clustering and weighted averaging (external function: apply_wbc_to_patient) to obtain consistent outptus.
- for 2D networks, consolidates box predictions to 3D cubes via clustering (adaption of non-maximum surpression).
(external function: merge_2D_to_3D_preds_per_patient)
Ground truth handling:
- dissmisses any ground truth boxes returned by the model (happens in validation mode, patch-based groundtruth)
- if provided by data loader, adds 3D ground truth to the final predictions to be passed to the evaluator.
"""
def __init__(self, cf, net, logger, mode):
self.cf = cf
self.logger = logger
# mode is 'val' for patient-based validation/monitoring and 'test' for inference.
self.mode = mode
# model instance. In validation mode, contains parameters of current epoch.
self.net = net
# rank of current epoch loaded (for temporal averaging). this info is added to each prediction,
# for correct weighting during consolidation.
self.rank_ix = '0'
# number of ensembled models. used to calculate the number of expected predictions per position
# during consolidation of predictions. Default is 1 (no ensembling, e.g. in validation).
self.n_ens = 1
if self.mode == 'test':
try:
self.epoch_ranking = np.load(os.path.join(self.cf.fold_dir, 'epoch_ranking.npy'))[:cf.test_n_epochs]
except:
raise RuntimeError('no epoch ranking file in fold directory. '
'seems like you are trying to run testing without prior training...')
self.n_ens = cf.test_n_epochs
if self.cf.test_aug:
self.n_ens *= 4
self.example_plot_dir = os.path.join(cf.test_dir, "example_plots")
os.makedirs(self.example_plot_dir, exist_ok=True)
def predict_patient(self, batch):
"""
predicts one patient.
called either directly via loop over validation set in exec.py (mode=='val')
or from self.predict_test_set (mode=='test).
in val mode: adds 3D ground truth info to predictions and runs consolidation and 2Dto3D merging of predictions.
in test mode: returns raw predictions (ground truth addition, consolidation, 2D to 3D merging are
done in self.predict_test_set, because patient predictions across several epochs might be needed
to be collected first, in case of temporal ensembling).
:return. results_dict: stores the results for one patient. dictionary with keys:
- 'boxes': list over batch elements. each element is a list over boxes, where each box is
one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions
(if not merged to 3D), and a dummy batch dimension of 1 for 3D predictions.
- 'seg_preds': pixel-wise predictions. (b, 1, y, x, (z))
- losses (only in validation mode)
"""
#self.logger.info('\revaluating patient {} for fold {} '.format(batch['pid'], self.cf.fold))
print('\revaluating patient {} for fold {} '.format(batch['pid'], self.cf.fold), end="", flush=True)
# True if patient is provided in patches and predictions need to be tiled.
- self.patched_patient = True if 'patch_crop_coords' in list(batch.keys()) else False
+ self.patched_patient = 'patch_crop_coords' in batch.keys()
# forward batch through prediction pipeline.
results_dict = self.data_aug_forward(batch)
if self.mode == 'val':
for b in range(batch['patient_bb_target'].shape[0]):
for t in range(len(batch['patient_bb_target'][b])):
results_dict['boxes'][b].append({'box_coords': batch['patient_bb_target'][b][t],
'box_label': batch['patient_roi_labels'][b][t],
'box_type': 'gt'})
if self.patched_patient:
wcs_input = [results_dict['boxes'], 'dummy_pid', self.cf.class_dict, self.cf.wcs_iou, self.n_ens]
results_dict['boxes'] = apply_wbc_to_patient(wcs_input)[0]
if self.cf.merge_2D_to_3D_preds:
merge_dims_inputs = [results_dict['boxes'], 'dummy_pid', self.cf.class_dict, self.cf.merge_3D_iou]
results_dict['boxes'] = merge_2D_to_3D_preds_per_patient(merge_dims_inputs)[0]
return results_dict
def predict_test_set(self, batch_gen, return_results=True):
"""
wrapper around test method, which loads multiple (or one) epoch parameters (temporal ensembling), loops through
the test set and collects predictions per patient. Also flattens the results per patient and epoch
and adds optional ground truth boxes for evaluation. Saves out the raw result list for later analysis and
optionally consolidates and returns predictions immediately.
:return: (optionally) list_of_results_per_patient: list over patient results. each entry is a dict with keys:
- 'boxes': list over batch elements. each element is a list over boxes, where each box is
one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions
(if not merged to 3D), and a dummy batch dimension of 1 for 3D predictions.
- 'seg_preds': not implemented yet. todo for evaluation of instance/semantic segmentation.
"""
dict_of_patient_results = OrderedDict()
# get paths of all parameter sets to be loaded for temporal ensembling. (or just one for no temp. ensembling).
weight_paths = [os.path.join(self.cf.fold_dir, '{}_best_checkpoint'.format(epoch), 'params.pth') for epoch in
self.epoch_ranking]
n_test_plots = min(batch_gen['n_test'], 1)
for rank_ix, weight_path in enumerate(weight_paths):
self.logger.info(('tmp ensembling over rank_ix:{} epoch:{}'.format(rank_ix, weight_path)))
self.net.load_state_dict(torch.load(weight_path))
self.net.eval()
self.rank_ix = str(rank_ix) # get string of current rank for unique patch ids.
plot_batches = np.random.choice(np.arange(batch_gen['n_test']), size=n_test_plots, replace=False)
with torch.no_grad():
for i in range(batch_gen['n_test']):
batch = next(batch_gen['test'])
# store batch info in patient entry of results dict.
if rank_ix == 0:
dict_of_patient_results[batch['pid']] = {}
dict_of_patient_results[batch['pid']]['results_dicts'] = []
dict_of_patient_results[batch['pid']]['patient_bb_target'] = batch['patient_bb_target']
dict_of_patient_results[batch['pid']]['patient_roi_labels'] = batch['patient_roi_labels']
# call prediction pipeline and store results in dict.
results_dict = self.predict_patient(batch)
dict_of_patient_results[batch['pid']]['results_dicts'].append({"boxes": results_dict['boxes']})
- if i in plot_batches and (not self.patched_patient or 'patient_data' in batch.keys()):
+ if i in plot_batches and not self.patched_patient:
+ # view qualitative results of random test case
+ # plotting for patched patients is too expensive, thus not done. Change at will.
try:
- # view qualitative results of random test case
out_file = os.path.join(self.example_plot_dir,
- 'batch_example_test_{}_rank_{}.png'.format(self.cf.fold, rank_ix))
- plot_results = results_dict.copy()
- # seg preds of test augs are included separately. for viewing only show aug 0 (merging
+ 'batch_example_test_{}_rank_{}.png'.format(self.cf.fold,
+ rank_ix))
+ results_for_plotting = deepcopy(results_dict)
+ # seg preds of test augs are included separately. for viewing, only show aug 0 (merging
# would need multiple changes, incl in every model).
- if plot_results["seg_preds"].shape[1] > 1:
- plot_results["seg_preds"] = results_dict['seg_preds'][:,[0]]
- utils.split_off_process(plot_batch_prediction, batch, results_dict, self.cf,
- outfile=out_file)
+ if results_for_plotting["seg_preds"].shape[1] > 1:
+ results_for_plotting["seg_preds"] = results_dict['seg_preds'][:, [0]]
+ for bix in range(batch["seg"].shape[0]): # batch dim should be 1
+ for tix in range(len(batch['bb_target'][bix])):
+ results_for_plotting['boxes'][bix].append({'box_coords': batch['bb_target'][bix][tix],
+ 'box_label': batch['class_target'][bix][tix],
+ 'box_type': 'gt'})
+ utils.split_off_process(plot_batch_prediction, batch, results_for_plotting, self.cf,
+ outfile=out_file, suptitle="Test plot:\nunmerged TTA overlayed.")
except Exception as e:
self.logger.info("WARNING: error in plotting example test batch: {}".format(e))
self.logger.info('finished predicting test set. starting post-processing of predictions.')
results_per_patient = []
# loop over patients again to flatten results across epoch predictions.
# if provided, add ground truth boxes for evaluation.
for pid, p_dict in dict_of_patient_results.items():
tmp_ens_list = p_dict['results_dicts']
results_dict = {}
# collect all boxes/seg_preds of same batch_instance over temporal instances.
b_size = len(tmp_ens_list[0]["boxes"])
results_dict['boxes'] = [[item for rank_dict in tmp_ens_list for item in rank_dict["boxes"][batch_instance]]
for batch_instance in range(b_size)]
# TODO return for instance segmentation:
# results_dict['seg_preds'] = np.mean(results_dict['seg_preds'], 1)[:, None]
# results_dict['seg_preds'] = np.array([[item for d in tmp_ens_list for item in d['seg_preds'][batch_instance]]
# for batch_instance in range(len(tmp_ens_list[0]['boxes']))])
# add 3D ground truth boxes for evaluation.
for b in range(p_dict['patient_bb_target'].shape[0]):
for t in range(len(p_dict['patient_bb_target'][b])):
results_dict['boxes'][b].append({'box_coords': p_dict['patient_bb_target'][b][t],
'box_label': p_dict['patient_roi_labels'][b][t],
'box_type': 'gt'})
results_per_patient.append([results_dict, pid])
# save out raw predictions.
out_string = 'raw_pred_boxes_hold_out_list' if self.cf.hold_out_test_set else 'raw_pred_boxes_list'
with open(os.path.join(self.cf.fold_dir, '{}.pickle'.format(out_string)), 'wb') as handle:
pickle.dump(results_per_patient, handle)
if return_results:
final_patient_box_results = [(res_dict["boxes"], pid) for res_dict, pid in results_per_patient]
# consolidate predictions.
self.logger.info('applying wcs to test set predictions with iou = {} and n_ens = {}.'.format(
self.cf.wcs_iou, self.n_ens))
pool = Pool(processes=6)
mp_inputs = [[ii[0], ii[1], self.cf.class_dict, self.cf.wcs_iou, self.n_ens] for ii in final_patient_box_results]
final_patient_box_results = pool.map(apply_wbc_to_patient, mp_inputs, chunksize=1)
pool.close()
pool.join()
# merge 2D boxes to 3D cubes. (if model predicts 2D but evaluation is run in 3D)
if self.cf.merge_2D_to_3D_preds:
self.logger.info('applying 2Dto3D merging to test set predictions with iou = {}.'.format(self.cf.merge_3D_iou))
pool = Pool(processes=6)
mp_inputs = [[ii[0], ii[1], self.cf.class_dict, self.cf.merge_3D_iou] for ii in final_patient_box_results]
final_patient_box_results = pool.map(merge_2D_to_3D_preds_per_patient, mp_inputs, chunksize=1)
pool.close()
pool.join()
# final_patient_box_results holds [avg_boxes, pid] if wbc
for ix in range(len(results_per_patient)):
assert results_per_patient[ix][1] == final_patient_box_results[ix][1], "should be same pid"
results_per_patient[ix][0]["boxes"] = final_patient_box_results[ix][0]
return results_per_patient
def load_saved_predictions(self, apply_wbc=False):
"""
loads raw predictions saved by self.predict_test_set. consolidates and merges 2D boxes to 3D cubes for evaluation.
(if model predicts 2D but evaluation is run in 3D)
:return: (optionally) results_list: list over patient results. each entry is a dict with keys:
- 'boxes': list over batch elements. each element is a list over boxes, where each box is
one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions
(if not merged to 3D), and a dummy batch dimension of 1 for 3D predictions.
- 'seg_preds': not implemented yet. todo for evaluation of instance/semantic segmentation.
"""
# load predictions for a single test-set fold.
results_file = 'raw_pred_boxes_hold_out_list.pickle' if self.cf.hold_out_test_set else 'raw_pred_boxes_list.pickle'
if not self.cf.hold_out_test_set or not self.cf.ensemble_folds:
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 = 4 if self.cf.test_aug else 1
n_ens = self.cf.test_n_epochs * da_factor
self.logger.info('loaded raw test set predictions with n_patients = {} and n_ens = {}'.format(
len(results_list), n_ens))
# if hold out test set was perdicted, aggregate predictions of all trained models
# corresponding to all CV-folds and flatten them.
else:
self.logger.info("loading saved predictions of hold-out test set and ensembling over folds.")
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))
da_factor = 4 if self.cf.test_aug else 1
n_ens = self.cf.test_n_epochs * da_factor * folds_loaded
# consolidate predictions.
if apply_wbc:
self.logger.info('applying wcs to test set predictions with iou = {} and n_ens = {}.'.format(
self.cf.wcs_iou, n_ens))
pool = Pool(processes=6)
mp_inputs = [[ii[0], ii[1], self.cf.class_dict, self.cf.wcs_iou, n_ens] for ii in box_results_list]
box_results_list = pool.map(apply_wbc_to_patient, mp_inputs, chunksize=1)
pool.close()
pool.join()
# merge 2D box predictions to 3D cubes (if model predicts 2D but evaluation is run in 3D)
if self.cf.merge_2D_to_3D_preds:
self.logger.info(
'applying 2Dto3D merging to test set predictions with iou = {}.'.format(self.cf.merge_3D_iou))
pool = Pool(processes=6)
mp_inputs = [[ii[0], ii[1], self.cf.class_dict, self.cf.merge_3D_iou] for ii in box_results_list]
box_results_list = pool.map(merge_2D_to_3D_preds_per_patient, mp_inputs, chunksize=1)
pool.close()
pool.join()
for ix in range(len(results_list)):
assert np.all(
results_list[ix][1] == box_results_list[ix][1]), "pid mismatch between loaded and aggregated results"
results_list[ix][0]["boxes"] = box_results_list[ix][0]
return results_list # holds (results_dict, pid)
def data_aug_forward(self, batch):
"""
in val_mode: passes batch through to spatial_tiling method without data_aug.
in test_mode: if cf.test_aug is set in configs, createst 4 mirrored versions of the input image,
passes all of them to the next processing step (spatial_tiling method) and re-transforms returned predictions
to original image version.
:return. results_dict: stores the results for one patient. dictionary with keys:
- 'boxes': list over batch elements. each element is a list over boxes, where each box is
one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions,
and a dummy batch dimension of 1 for 3D predictions.
- 'seg_preds': pixel-wise predictions. (b, 1, y, x, (z))
- losses (only in validation mode)
"""
patch_crops = batch['patch_crop_coords'] if self.patched_patient else None
results_list = [self.spatial_tiling_forward(batch, patch_crops)]
org_img_shape = batch['original_img_shape']
if self.mode == 'test' and self.cf.test_aug:
if self.patched_patient:
# apply mirror transformations to patch-crop coordinates, for correct tiling in spatial_tiling method.
mirrored_patch_crops = get_mirrored_patch_crops(patch_crops, batch['original_img_shape'])
else:
mirrored_patch_crops = [None] * 3
img = np.copy(batch['data'])
# first mirroring: y-axis.
batch['data'] = np.flip(img, axis=2).copy()
chunk_dict = self.spatial_tiling_forward(batch, mirrored_patch_crops[0], n_aug='1')
# re-transform coordinates.
for ix in range(len(chunk_dict['boxes'])):
for boxix in range(len(chunk_dict['boxes'][ix])):
coords = chunk_dict['boxes'][ix][boxix]['box_coords'].copy()
coords[0] = org_img_shape[2] - chunk_dict['boxes'][ix][boxix]['box_coords'][2]
coords[2] = org_img_shape[2] - chunk_dict['boxes'][ix][boxix]['box_coords'][0]
assert coords[2] >= coords[0], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]
assert coords[3] >= coords[1], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]
chunk_dict['boxes'][ix][boxix]['box_coords'] = coords
# re-transform segmentation predictions.
chunk_dict['seg_preds'] = np.flip(chunk_dict['seg_preds'], axis=2)
results_list.append(chunk_dict)
# second mirroring: x-axis.
batch['data'] = np.flip(img, axis=3).copy()
chunk_dict = self.spatial_tiling_forward(batch, mirrored_patch_crops[1], n_aug='2')
# re-transform coordinates.
for ix in range(len(chunk_dict['boxes'])):
for boxix in range(len(chunk_dict['boxes'][ix])):
coords = chunk_dict['boxes'][ix][boxix]['box_coords'].copy()
coords[1] = org_img_shape[3] - chunk_dict['boxes'][ix][boxix]['box_coords'][3]
coords[3] = org_img_shape[3] - chunk_dict['boxes'][ix][boxix]['box_coords'][1]
assert coords[2] >= coords[0], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]
assert coords[3] >= coords[1], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]
chunk_dict['boxes'][ix][boxix]['box_coords'] = coords
# re-transform segmentation predictions.
chunk_dict['seg_preds'] = np.flip(chunk_dict['seg_preds'], axis=3)
results_list.append(chunk_dict)
# third mirroring: y-axis and x-axis.
batch['data'] = np.flip(np.flip(img, axis=2), axis=3).copy()
chunk_dict = self.spatial_tiling_forward(batch, mirrored_patch_crops[2], n_aug='3')
# re-transform coordinates.
for ix in range(len(chunk_dict['boxes'])):
for boxix in range(len(chunk_dict['boxes'][ix])):
coords = chunk_dict['boxes'][ix][boxix]['box_coords'].copy()
coords[0] = org_img_shape[2] - chunk_dict['boxes'][ix][boxix]['box_coords'][2]
coords[2] = org_img_shape[2] - chunk_dict['boxes'][ix][boxix]['box_coords'][0]
coords[1] = org_img_shape[3] - chunk_dict['boxes'][ix][boxix]['box_coords'][3]
coords[3] = org_img_shape[3] - chunk_dict['boxes'][ix][boxix]['box_coords'][1]
assert coords[2] >= coords[0], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]
assert coords[3] >= coords[1], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]
chunk_dict['boxes'][ix][boxix]['box_coords'] = coords
# re-transform segmentation predictions.
chunk_dict['seg_preds'] = np.flip(np.flip(chunk_dict['seg_preds'], axis=2), axis=3).copy()
results_list.append(chunk_dict)
batch['data'] = img
# aggregate all boxes/seg_preds per batch element from data_aug predictions.
results_dict = {}
results_dict['boxes'] = [[item for d in results_list for item in d['boxes'][batch_instance]]
for batch_instance in range(org_img_shape[0])]
results_dict['seg_preds'] = np.array([[item for d in results_list for item in d['seg_preds'][batch_instance]]
for batch_instance in range(org_img_shape[0])])
if self.mode == 'val':
try:
results_dict['torch_loss'] = results_list[0]['torch_loss']
results_dict['class_loss'] = results_list[0]['class_loss']
except KeyError:
pass
return results_dict
def spatial_tiling_forward(self, batch, patch_crops=None, n_aug='0'):
"""
forwards batch to batch_tiling_forward method and receives and returns a dictionary with results.
if patch-based prediction, the results received from batch_tiling_forward will be on a per-patch-basis.
this method uses the provided patch_crops to re-transform all predictions to whole-image coordinates.
Patch-origin information of all box-predictions will be needed for consolidation, hence it is stored as
'patch_id', which is a unique string for each patch (also takes current data aug and temporal epoch instances
into account). all box predictions get additional information about the amount overlapping patches at the
respective position (used for consolidation).
:return. results_dict: stores the results for one patient. dictionary with keys:
- 'boxes': list over batch elements. each element is a list over boxes, where each box is
one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions,
and a dummy batch dimension of 1 for 3D predictions.
- 'seg_preds': pixel-wise predictions. (b, 1, y, x, (z))
- losses (only in validation mode)
"""
if patch_crops is not None:
patches_dict = self.batch_tiling_forward(batch)
results_dict = {'boxes': [[] for _ in range(batch['original_img_shape'][0])]}
# instanciate segemntation output array. Will contain averages over patch predictions.
out_seg_preds = np.zeros(batch['original_img_shape'], dtype=np.float16)[:, 0][:, None]
# counts patch instances per pixel-position.
patch_overlap_map = np.zeros_like(out_seg_preds, dtype='uint8')
#unmold segmentation outputs. loop over patches.
for pix, pc in enumerate(patch_crops):
if self.cf.dim == 3:
out_seg_preds[:, :, pc[0]:pc[1], pc[2]:pc[3], pc[4]:pc[5]] += patches_dict['seg_preds'][pix][None]
patch_overlap_map[:, :, pc[0]:pc[1], pc[2]:pc[3], pc[4]:pc[5]] += 1
else:
out_seg_preds[pc[4]:pc[5], :, pc[0]:pc[1], pc[2]:pc[3], ] += patches_dict['seg_preds'][pix]
patch_overlap_map[pc[4]:pc[5], :, pc[0]:pc[1], pc[2]:pc[3], ] += 1
# take mean in overlapping areas.
out_seg_preds[patch_overlap_map > 0] /= patch_overlap_map[patch_overlap_map > 0]
results_dict['seg_preds'] = out_seg_preds
# unmold box outputs. loop over patches.
for pix, pc in enumerate(patch_crops):
patch_boxes = patches_dict['boxes'][pix]
for box in patch_boxes:
# add unique patch id for consolidation of predictions.
box['patch_id'] = self.rank_ix + '_' + n_aug + '_' + str(pix)
# boxes from the edges of a patch have a lower prediction quality, than the ones at patch-centers.
# hence they will be downweighted for consolidation, using the 'box_patch_center_factor', which is
# obtained by a normal distribution over positions in the patch and average over spatial dimensions.
# Also the info 'box_n_overlaps' is stored for consolidation, which depicts the amount over
# overlapping patches at the box's position.
c = box['box_coords']
box_centers = [(c[ii] + c[ii + 2]) / 2 for ii in range(2)]
if self.cf.dim == 3:
box_centers.append((c[4] + c[5]) / 2)
box['box_patch_center_factor'] = np.mean(
[norm.pdf(bc, loc=pc, scale=pc * 0.8) * np.sqrt(2 * np.pi) * pc * 0.8 for bc, pc in
zip(box_centers, np.array(self.cf.patch_size) / 2)])
if self.cf.dim == 3:
c += np.array([pc[0], pc[2], pc[0], pc[2], pc[4], pc[4]])
int_c = [int(np.floor(ii)) if ix%2 == 0 else int(np.ceil(ii)) for ix, ii in enumerate(c)]
box['box_n_overlaps'] = np.mean(patch_overlap_map[:, :, int_c[1]:int_c[3], int_c[0]:int_c[2], int_c[4]:int_c[5]])
results_dict['boxes'][0].append(box)
else:
c += np.array([pc[0], pc[2], pc[0], pc[2]])
int_c = [int(np.floor(ii)) if ix % 2 == 0 else int(np.ceil(ii)) for ix, ii in enumerate(c)]
box['box_n_overlaps'] = np.mean(patch_overlap_map[pc[4], :, int_c[1]:int_c[3], int_c[0]:int_c[2]])
results_dict['boxes'][pc[4]].append(box)
if self.mode == 'val':
try:
results_dict['torch_loss'] = patches_dict['torch_loss']
results_dict['class_loss'] = patches_dict['class_loss']
except KeyError:
pass
# if predictions are not patch-based:
# add patch-origin info to boxes (entire image is the same patch with overlap=1) and return results.
else:
results_dict = self.batch_tiling_forward(batch)
for b in results_dict['boxes']:
for box in b:
box['box_patch_center_factor'] = 1
box['box_n_overlaps'] = 1
box['patch_id'] = self.rank_ix + '_' + n_aug
return results_dict
def batch_tiling_forward(self, batch):
"""
calls the actual network forward method. in patch-based prediction, the batch dimension might be overladed
with n_patches >> batch_size, which would exceed gpu memory. In this case, batches are processed in chunks of
batch_size. validation mode calls the train method to monitor losses (returned ground truth objects are discarded).
test mode calls the test forward method, no ground truth required / involved.
:return. results_dict: stores the results for one patient. dictionary with keys:
- 'boxes': list over batch elements. each element is a list over boxes, where each box is
one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions,
and a dummy batch dimension of 1 for 3D predictions.
- 'seg_preds': pixel-wise predictions. (b, 1, y, x, (z))
- losses (only in validation mode)
"""
#self.logger.info('forwarding (patched) patient with shape: {}'.format(batch['data'].shape))
img = batch['data']
if img.shape[0] <= self.cf.batch_size:
if self.mode == 'val':
# call training method to monitor losses
results_dict = self.net.train_forward(batch, is_validation=True)
# discard returned ground-truth boxes (also training info boxes).
results_dict['boxes'] = [[box for box in b if box['box_type'] == 'det'] for b in results_dict['boxes']]
else:
results_dict = self.net.test_forward(batch, return_masks=self.cf.return_masks_in_test)
else:
split_ixs = np.split(np.arange(img.shape[0]), np.arange(img.shape[0])[::self.cf.batch_size])
chunk_dicts = []
for chunk_ixs in split_ixs[1:]: # first split is elements before 0, so empty
b = {k: batch[k][chunk_ixs] for k in batch.keys()
if (isinstance(batch[k], np.ndarray) and batch[k].shape[0] == img.shape[0])}
if self.mode == 'val':
chunk_dicts += [self.net.train_forward(b, is_validation=True)]
else:
chunk_dicts += [self.net.test_forward(b, return_masks=self.cf.return_masks_in_test)]
results_dict = {}
# flatten out batch elements from chunks ([chunk, chunk] -> [b, b, b, b, ...])
results_dict['boxes'] = [item for d in chunk_dicts for item in d['boxes']]
results_dict['seg_preds'] = np.array([item for d in chunk_dicts for item in d['seg_preds']])
if self.mode == 'val':
try:
# estimate metrics by mean over batch_chunks. Most similar to training metrics.
results_dict['torch_loss'] = torch.mean(torch.cat([d['torch_loss'] for d in chunk_dicts]))
results_dict['class_loss'] = np.mean([d['class_loss'] for d in chunk_dicts])
except KeyError:
# losses are not necessarily monitored
pass
# discard returned ground-truth boxes (also training info boxes).
results_dict['boxes'] = [[box for box in b if box['box_type'] == 'det'] for b in results_dict['boxes']]
return results_dict
def apply_wbc_to_patient(inputs):
"""
wrapper around prediction box consolidation: weighted cluster scoring (wcs). processes a single patient.
loops over batch elements in patient results (1 in 3D, slices in 2D) and foreground classes,
aggregates and stores results in new list.
:return. patient_results_list: list over batch elements. each element is a list over boxes, where each box is
one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D
predictions, and a dummy batch dimension of 1 for 3D predictions.
:return. pid: string. patient id.
"""
in_patient_results_list, pid, class_dict, wcs_iou, n_ens = inputs
out_patient_results_list = [[] for _ in range(len(in_patient_results_list))]
for bix, b in enumerate(in_patient_results_list):
for cl in list(class_dict.keys()):
boxes = [(ix, box) for ix, box in enumerate(b) if (box['box_type'] == 'det' and box['box_pred_class_id'] == cl)]
box_coords = np.array([b[1]['box_coords'] for b in boxes])
box_scores = np.array([b[1]['box_score'] for b in boxes])
box_center_factor = np.array([b[1]['box_patch_center_factor'] for b in boxes])
box_n_overlaps = np.array([b[1]['box_n_overlaps'] for b in boxes])
box_patch_id = np.array([b[1]['patch_id'] for b in boxes])
if 0 not in box_scores.shape:
keep_scores, keep_coords = weighted_box_clustering(
np.concatenate((box_coords, box_scores[:, None], box_center_factor[:, None],
box_n_overlaps[:, None]), axis=1), box_patch_id, wcs_iou, n_ens)
for boxix in range(len(keep_scores)):
out_patient_results_list[bix].append({'box_type': 'det', 'box_coords': keep_coords[boxix],
'box_score': keep_scores[boxix], 'box_pred_class_id': cl})
# add gt boxes back to new output list.
out_patient_results_list[bix].extend([box for box in b if box['box_type'] == 'gt'])
return [out_patient_results_list, pid]
def merge_2D_to_3D_preds_per_patient(inputs):
"""
wrapper around 2Dto3D merging operation. Processes a single patient. Takes 2D patient results (slices in batch dimension)
and returns 3D patient results (dummy batch dimension of 1). Applies an adaption of Non-Maximum Surpression
(Detailed methodology is described in nms_2to3D).
:return. results_dict_boxes: list over batch elements (1 in 3D). each element is a list over boxes, where each box is
one dictionary: [[box_0, ...], [box_n,...]].
:return. pid: string. patient id.
"""
in_patient_results_list, pid, class_dict, merge_3D_iou = inputs
out_patient_results_list = []
for cl in list(class_dict.keys()):
boxes, slice_ids = [], []
# collect box predictions over batch dimension (slices) and store slice info as slice_ids.
for bix, b in enumerate(in_patient_results_list):
det_boxes = [(ix, box) for ix, box in enumerate(b) if
(box['box_type'] == 'det' and box['box_pred_class_id'] == cl)]
boxes += det_boxes
slice_ids += [bix] * len(det_boxes)
box_coords = np.array([b[1]['box_coords'] for b in boxes])
box_scores = np.array([b[1]['box_score'] for b in boxes])
slice_ids = np.array(slice_ids)
if 0 not in box_scores.shape:
keep_ix, keep_z = nms_2to3D(
np.concatenate((box_coords, box_scores[:, None], slice_ids[:, None]), axis=1), merge_3D_iou)
else:
keep_ix, keep_z = [], []
# store kept predictions in new results list and add corresponding z-dimension info to coordinates.
for kix, kz in zip(keep_ix, keep_z):
out_patient_results_list.append({'box_type': 'det', 'box_coords': list(box_coords[kix]) + kz,
'box_score': box_scores[kix], 'box_pred_class_id': cl})
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
# add dummy batch dimension 1 for 3D.
return [[out_patient_results_list], pid]
def weighted_box_clustering(dets, box_patch_id, thresh, n_ens):
"""
consolidates overlapping predictions resulting from patch overlaps, test data augmentations and temporal ensembling.
clusters predictions together with iou > thresh (like in NMS). Output score and coordinate for one cluster are the
average weighted by individual patch center factors (how trustworthy is this candidate measured by how centered
its position the patch is) and the size of the corresponding box.
The number of expected predictions at a position is n_data_aug * n_temp_ens * n_overlaps_at_position
(1 prediction per unique patch). Missing predictions at a cluster position are defined as the number of unique
patches in the cluster, which did not contribute any predict any boxes.
:param dets: (n_dets, (y1, x1, y2, x2, (z1), (z2), scores, box_pc_facts, box_n_ovs)
:param thresh: threshold for iou_matching.
:param n_ens: number of models, that are ensembled. (-> number of expected predicitions per position)
:return: keep_scores: (n_keep) new scores of boxes to be kept.
:return: keep_coords: (n_keep, (y1, x1, y2, x2, (z1), (z2)) new coordinates of boxes to be kept.
"""
dim = 2 if dets.shape[1] == 7 else 3
y1 = dets[:, 0]
x1 = dets[:, 1]
y2 = dets[:, 2]
x2 = dets[:, 3]
scores = dets[:, -3]
box_pc_facts = dets[:, -2]
box_n_ovs = dets[:, -1]
areas = (y2 - y1 + 1) * (x2 - x1 + 1)
if dim == 3:
z1 = dets[:, 4]
z2 = dets[:, 5]
areas *= (z2 - z1 + 1)
# order is the sorted index. maps order to index o[1] = 24 (rank1, ix 24)
order = scores.argsort()[::-1]
keep = []
keep_scores = []
keep_coords = []
while order.size > 0:
i = order[0] # higehst scoring element
xx1 = np.maximum(x1[i], x1[order])
yy1 = np.maximum(y1[i], y1[order])
xx2 = np.minimum(x2[i], x2[order])
yy2 = np.minimum(y2[i], y2[order])
w = np.maximum(0.0, xx2 - xx1 + 1)
h = np.maximum(0.0, yy2 - yy1 + 1)
inter = w * h
if dim == 3:
zz1 = np.maximum(z1[i], z1[order])
zz2 = np.minimum(z2[i], z2[order])
d = np.maximum(0.0, zz2 - zz1 + 1)
inter *= d
# overall between currently highest scoring box and all boxes.
ovr = inter / (areas[i] + areas[order] - inter)
# get all the predictions that match the current box to build one cluster.
matches = np.argwhere(ovr > thresh)
match_n_ovs = box_n_ovs[order[matches]]
match_pc_facts = box_pc_facts[order[matches]]
match_patch_id = box_patch_id[order[matches]]
match_ov_facts = ovr[matches]
match_areas = areas[order[matches]]
match_scores = scores[order[matches]]
# weight all socres in cluster by patch factors, and size.
match_score_weights = match_ov_facts * match_areas * match_pc_facts
match_scores *= match_score_weights
# for the weigted average, scores have to be divided by the number of total expected preds at the position
# of the current cluster. 1 Prediction per patch is expected. therefore, the number of ensembled models is
# multiplied by the mean overlaps of patches at this position (boxes of the cluster might partly be
# in areas of different overlaps).
n_expected_preds = n_ens * np.mean(match_n_ovs)
# the number of missing predictions is obtained as the number of patches,
# which did not contribute any prediction to the current cluster.
n_missing_preds = np.max((0, n_expected_preds - np.unique(match_patch_id).shape[0]))
# missing preds are given the mean weighting
# (expected prediction is the mean over all predictions in cluster).
denom = np.sum(match_score_weights) + n_missing_preds * np.mean(match_score_weights)
# compute weighted average score for the cluster
avg_score = np.sum(match_scores) / denom
# compute weighted average of coordinates for the cluster. now only take existing
# predictions into account.
avg_coords = [np.sum(y1[order[matches]] * match_scores) / np.sum(match_scores),
np.sum(x1[order[matches]] * match_scores) / np.sum(match_scores),
np.sum(y2[order[matches]] * match_scores) / np.sum(match_scores),
np.sum(x2[order[matches]] * match_scores) / np.sum(match_scores)]
if dim == 3:
avg_coords.append(np.sum(z1[order[matches]] * match_scores) / np.sum(match_scores))
avg_coords.append(np.sum(z2[order[matches]] * match_scores) / np.sum(match_scores))
# some clusters might have very low scores due to high amounts of missing predictions.
# filter out the with a conservative threshold, to speed up evaluation.
if avg_score > 0.01:
keep_scores.append(avg_score)
keep_coords.append(avg_coords)
# get index of all elements that were not matched and discard all others.
inds = np.where(ovr <= thresh)[0]
order = order[inds]
return keep_scores, keep_coords
def nms_2to3D(dets, thresh):
"""
Merges 2D boxes to 3D cubes. Therefore, boxes of all slices are projected into one slices. An adaptation of Non-maximum surpression
is applied, where clusters are found (like in NMS) with an extra constrained, that surpressed boxes have to have 'connected'
z-coordinates w.r.t the core slice (cluster center, highest scoring box). 'connected' z-coordinates are determined
as the z-coordinates with predictions until the first coordinate, where no prediction was found.
example: a cluster of predictions was found overlap > iou thresh in xy (like NMS). The z-coordinate of the highest
scoring box is 50. Other predictions have 23, 46, 48, 49, 51, 52, 53, 56, 57.
Only the coordinates connected with 50 are clustered to one cube: 48, 49, 51, 52, 53. (46 not because nothing was
found in 47, so 47 is a 'hole', which interrupts the connection). Only the boxes corresponding to these coordinates
are surpressed. All others are kept for building of further clusters.
This algorithm works better with a certain min_confidence of predictions, because low confidence (e.g. noisy/cluttery)
predictions can break the relatively strong assumption of defining cubes' z-boundaries at the first 'hole' in the cluster.
:param dets: (n_detections, (y1, x1, y2, x2, scores, slice_id)
:param thresh: iou matchin threshold (like in NMS).
:return: keep: (n_keep) 1D tensor of indices to be kept.
:return: keep_z: (n_keep, [z1, z2]) z-coordinates to be added to boxes, which are kept in order to form cubes.
"""
y1 = dets[:, 0]
x1 = dets[:, 1]
y2 = dets[:, 2]
x2 = dets[:, 3]
scores = dets[:, -2]
slice_id = dets[:, -1]
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
order = scores.argsort()[::-1]
keep = []
keep_z = []
while order.size > 0: # order is the sorted index. maps order to index o[1] = 24 (rank1, ix 24)
i = order[0] # pop higehst scoring element
xx1 = np.maximum(x1[i], x1[order])
yy1 = np.maximum(y1[i], y1[order])
xx2 = np.minimum(x2[i], x2[order])
yy2 = np.minimum(y2[i], y2[order])
w = np.maximum(0.0, xx2 - xx1 + 1)
h = np.maximum(0.0, yy2 - yy1 + 1)
inter = w * h
ovr = inter / (areas[i] + areas[order] - inter)
matches = np.argwhere(ovr > thresh) # get all the elements that match the current box and have a lower score
slice_ids = slice_id[order[matches]]
core_slice = slice_id[int(i)]
upper_wholes = [ii for ii in np.arange(core_slice, np.max(slice_ids)) if ii not in slice_ids]
lower_wholes = [ii for ii in np.arange(np.min(slice_ids), core_slice) if ii not in slice_ids]
max_valid_slice_id = np.min(upper_wholes) if len(upper_wholes) > 0 else np.max(slice_ids)
min_valid_slice_id = np.max(lower_wholes) if len(lower_wholes) > 0 else np.min(slice_ids)
z_matches = matches[(slice_ids <= max_valid_slice_id) & (slice_ids >= min_valid_slice_id)]
z1 = np.min(slice_id[order[z_matches]]) - 1
z2 = np.max(slice_id[order[z_matches]]) + 1
keep.append(i)
keep_z.append([z1, z2])
order = np.delete(order, z_matches, axis=0)
return keep, keep_z
def get_mirrored_patch_crops(patch_crops, org_img_shape):
"""
apply 3 mirrror transformations (x-axis, y-axis, x&y-axis)
to given patch crop coordinates and return the transformed coordinates.
Handles 2D and 3D coordinates.
:param patch_crops: list of crops: each element is a list of coordinates for given crop [[y1, x1, ...], [y1, x1, ..]]
:param org_img_shape: shape of patient volume used as world coordinates.
:return: list of mirrored patch crops: lenght=3. each element is a list of transformed patch crops.
"""
mirrored_patch_crops = []
# y-axis transform.
mirrored_patch_crops.append([[org_img_shape[2] - ii[1],
org_img_shape[2] - ii[0],
ii[2], ii[3]] if len(ii) == 4 else
[org_img_shape[2] - ii[1],
org_img_shape[2] - ii[0],
ii[2], ii[3], ii[4], ii[5]] for ii in patch_crops])
# x-axis transform.
mirrored_patch_crops.append([[ii[0], ii[1],
org_img_shape[3] - ii[3],
org_img_shape[3] - ii[2]] if len(ii) == 4 else
[ii[0], ii[1],
org_img_shape[3] - ii[3],
org_img_shape[3] - ii[2],
ii[4], ii[5]] for ii in patch_crops])
# y-axis and x-axis transform.
mirrored_patch_crops.append([[org_img_shape[2] - ii[1],
org_img_shape[2] - ii[0],
org_img_shape[3] - ii[3],
org_img_shape[3] - ii[2]] if len(ii) == 4 else
[org_img_shape[2] - ii[1],
org_img_shape[2] - ii[0],
org_img_shape[3] - ii[3],
org_img_shape[3] - ii[2],
ii[4], ii[5]] for ii in patch_crops])
return mirrored_patch_crops
diff --git a/requirements.txt b/requirements.txt
index 0eda747..034ef00 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,67 +1,12 @@
-absl-py==0.9.0
-backcall==0.1.0
-batchgenerators==0.19.7
-cachetools==4.0.0
-certifi==2019.11.28
-cffi==1.11.5
-chardet==3.0.4
-cycler==0.10.0
-Cython==0.29.14
-decorator==4.4.1
-future==0.18.2
-google-auth==1.10.0
-google-auth-oauthlib==0.4.1
-grpcio==1.26.0
-idna==2.8
-imageio==2.6.1
-ipython-genutils==0.2.0
-jedi==0.15.1
-joblib==0.14.1
-kiwisolver==1.1.0
-linecache2==1.0.0
-Markdown==3.1.1
-matplotlib==3.1.2
-medicaldetectiontoolkit==0.0.1
-networkx==2.4
+batchgenerators==0.20.1
nms-extension==0.0.0
-numpy==1.17.4
-oauthlib==3.1.0
pandas==0.25.3
-parso==0.5.2
-pexpect==4.7.0
-pickleshare==0.7.5
-Pillow==7.1.0
-prompt-toolkit==3.0.2
-protobuf==3.11.2
-ptyprocess==0.6.0
-pyasn1==0.4.8
-pyasn1-modules==0.2.7
-pycparser==2.19
-Pygments==2.5.2
-pyparsing==2.4.5
-python-dateutil==2.8.1
-pytz==2019.3
-PyWavelets==1.1.1
-requests==2.22.0
-requests-oauthlib==1.3.0
+Pillow<7.1
RoIAlign-extension-2D==0.0.0
RoIAlign-extension-3D==0.0.0
-rsa==4.0
-scikit-image==0.16.2
-scikit-learn==0.22.1
-scipy==1.3.3
SimpleITK==1.2.4
-six==1.13.0
-sklearn==0.0
tensorboard==2.2.0
-tensorboard-plugin-wit==1.6.0.post2
-threadpoolctl==1.1.0
torch==1.4.0
torchvision==0.5.0
-tqdm==4.40.2
-traceback2==1.4.0
-traitlets==4.3.3
-unittest2==1.1.0
-urllib3==1.25.7
-wcwidth==0.1.7
-Werkzeug==1.0.1
+tqdm
+
diff --git a/setup.py b/setup.py
index 750a5d6..9001753 100644
--- a/setup.py
+++ b/setup.py
@@ -1,73 +1,74 @@
#!/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.
# ==============================================================================
from setuptools import find_packages, setup
import os, sys, subprocess
def parse_requirements(filename, exclude=[]):
lineiter = (line.strip() for line in open(filename))
return [line for line in lineiter if line and not line.startswith("#") and not line.split("==")[0] in exclude]
def pip_install(item):
subprocess.check_call([sys.executable, "-m", "pip", "install", item])
def install_custom_ext(setup_path):
try:
pip_install(setup_path)
except Exception as e:
print("Could not install custom extension {} from source due to Error:\n{}\n".format(path, e) +
"Trying to install from pre-compiled wheel.")
dist_path = setup_path+"/dist"
wheel_file = [fn for fn in os.listdir(dist_path) if fn.endswith(".whl")][0]
pip_install(os.path.join(dist_path, wheel_file))
def clean():
"""Custom clean command to tidy up the project root."""
os.system('rm -vrf ./build ./dist ./*.pyc ./*.tgz ./*.egg-info')
if __name__ == "__main__":
req_file = "requirements.txt"
custom_exts = ["nms-extension", "RoIAlign-extension-2D", "RoIAlign-extension-3D"]
install_reqs = parse_requirements(req_file, exclude=custom_exts)
setup(name='medicaldetectiontoolkit',
- version='0.0.1',
+ version='0.1.0',
url="https://github.com/MIC-DKFZ/medicaldetectiontoolkit",
author='P. Jaeger, G. Ramien, MIC at DKFZ Heidelberg',
license="Apache 2.0",
description="Medical Object-Detection Toolkit.",
classifiers=[
"Development Status :: 4 - Beta",
"Intended Audience :: Developers",
"Programming Language :: Python :: 3.7"
],
packages=find_packages(exclude=['test', 'test.*']),
install_requires=install_reqs,
+ python_requires=">=3.7"
)
custom_exts = ["custom_extensions/nms", "custom_extensions/roi_align/2D", "custom_extensions/roi_align/3D"]
for path in custom_exts:
try:
install_custom_ext(path)
except Exception as e:
print("FAILED to install custom extension {} due to Error:\n{}".format(path, e))
clean()
\ No newline at end of file
diff --git a/shell_scripts/bpeek_wrapper.sh b/shell_scripts/bpeek_wrapper.sh
new file mode 100644
index 0000000..7f1383e
--- /dev/null
+++ b/shell_scripts/bpeek_wrapper.sh
@@ -0,0 +1,19 @@
+#!/bin/bash
+
+# wrapper around lsf-scheduler bpeek command writing bpeek to file and opening it.
+# amend default output dir and editor as desired.
+# positonal
+# -arg #1 lsf job id
+
+job_id="${1}"
+out_dir=$HOME/lsf_peek_output
+out_file="${out_dir}/${job_id}.out"
+
+
+mkdir -p ${out_dir}
+
+eval "bpeek ${job_id} > ${out_file}"
+
+nano ${out_file}
+#tail -F ${out_file}
+
diff --git a/shell_scripts/cluster_install_notes.txt b/shell_scripts/cluster_install_notes.txt
deleted file mode 100644
index e69cf39..0000000
--- a/shell_scripts/cluster_install_notes.txt
+++ /dev/null
@@ -1,24 +0,0 @@
-
- - copy all source files of mdt-public to cluster destination, e.g., by using update_scripts_on_cluster.sh.
-
- - log in to a COMPUTE NODE, e.g., e132-comp01, not one of the worker/submission since we need CUDA installed.
-
- - run:
-
- module load python/3.7.0
- module load gcc/7.2.0
-
- virtualenv -p python3.7
- source /bin/activate
-
- export CUDA_HOME=/usr/local/cuda-${CUDA}
- export TORCH_CUDA_ARCH_LIST="6.1;7.0;7.5"
-
- cd mdt-public
- python setup.py install #--> check that custom extension are installed successfully.
-
-
-
-after install:
- - until we have a better solution: submit jobs not from the recommended worker nodes but from a compute node (since we need /datasets to be mounted for job submission).
-
diff --git a/shell_scripts/cluster_install_usage_notes.txt b/shell_scripts/cluster_install_usage_notes.txt
new file mode 100644
index 0000000..5a404df
--- /dev/null
+++ b/shell_scripts/cluster_install_usage_notes.txt
@@ -0,0 +1,37 @@
+install:
+ - copy all source files of mdt-public to cluster destination, e.g., by using update_scripts_on_cluster.sh.
+
+ - log in to a COMPUTE NODE, e.g., e132-comp01, not one of the worker/submission nodes since we need CUDA installed. and
+ stay in your home directory.
+
+ - run:
+
+ module load python/3.7.0
+ module load gcc/7.2.0
+
+ virtualenv -p python3.7 .virtualenvs/mdt
+ source .virtualenvs/mdt/bin/activate
+
+ export CUDA_HOME=/usr/local/cuda-${CUDA}
+ export TORCH_CUDA_ARCH_LIST="6.1;7.0;7.5"
+
+ cd mdt-public
+ python setup.py install #--> check that custom extension are installed successfully.
+
+
+
+after install/ usage:
+ - until we have a better solution: submit jobs not from the recommended worker nodes but from a compute node (since we need /datasets to be mounted for job submission).
+ - adjust the paths in job_starter.sh (root_dir and exp_parent_dir) and in cluster_runner_meddec.sh (job_dir=/ssd//...).
+ - job submission routine:
+ - log in to node
+ - cd mdt-public
+ - sh job_starter.sh *OPTIONS, where
+ - is the directory name of the dataset-specific source code (lidc_exp or toy_exp)
+ - is the name of the experiment directory (not a full or relative path, only the name). The experiment will be located under the parent dir /experiments.
+ - see job_starter.sh for further optional arguments, e.g. -p change the default parent dir.
+ - pass flag -c to indicate you want to create a new experiment.
+ - if a job crashed and you want to continue it from the last checkpoint, simply add --resume to its submission command.
+
+
+
diff --git a/shell_scripts/cluster_runner_meddec.sh b/shell_scripts/cluster_runner_meddec.sh
index 703e1a4..76122f7 100644
--- a/shell_scripts/cluster_runner_meddec.sh
+++ b/shell_scripts/cluster_runner_meddec.sh
@@ -1,65 +1,59 @@
#!/bin/bash
#Usage:
# -->not true?: this script has to be started from the same directory the python files called below lie in (e.g. exec.py lies in meddetectiontkit).
# part of the slurm-job name you pass to sbatch will be the experiment folder's name.
# you need to pass 3 positional arguments to this script (cluster_runner_..sh #1 #2 #3):
# -#1 source directory in which main source code (framework) is located (e.g. medicaldetectiontoolkit/)
# -#2 the exp_dir where job-specific code was copied before by create_exp and exp results are safed by exec.py
# -#3 absolute path to dataset-specific code in source dir
# -#4 mode to run
# -#5 folds to run on
source_dir=${1}
exp_dir=${2}
dataset_abs_path=${3}
mode=${4}
folds=${5}
resume=$6
-#known problem: trap somehow does not execute the rm -r tmp_dir command when using scancel on job
-#trap clean_up EXIT KILL TERM ABRT QUIT
-
job_dir=/ssd/ramien/${LSB_JOBID}
tmp_dir_data=${job_dir}/data
mkdir $tmp_dir_data
tmp_dir_cache=${job_dir}/cache
mkdir $tmp_dir_cache
CUDA_CACHE_PATH=$tmp_dir_cache
export CUDA_CACHE_PATH
#data must not lie permantly on nodes' ssd, only during training time
#needs to be named with the SLURM_JOB_ID to not be automatically removed
#can permanently lie on /datasets drive --> copy from there before every experiment
#files on datasets are saved as npz (compressed) --> use data_manager.py to copy and unpack into .npy; is done implicitly in exec.py
#(tensorboard --logdir ${exp_dir}/.. --port 1337 || echo "tboard startup failed")& # || tensorboard --logdir ${exp_dir}/.. --port 1338)&
#tboard_pid=$!
-#clean_up() {
-# rm -rf ${job_dir};
-#}
export OMP_NUM_THREADS=1 # this is a work-around fix for batchgenerators to deal with numpy-inherent multi-threading.
launch_opts="${source_dir}/exec.py --use_stored_settings --server_env --exp_source ${dataset_abs_path} --data_dest ${tmp_dir_data} --exp_dir ${exp_dir} --mode ${mode}"
if [ ! -z "${resume}" ]; then
launch_opts="${launch_opts} --resume"
echo "Resuming from checkpoint(s)."
fi
if [ ! -z "${folds}" ]; then
launch_opts="${launch_opts} --folds ${folds}"
fi
echo "submitting with ${launch_opts}"
python ${launch_opts}
diff --git a/shell_scripts/job_starter.sh b/shell_scripts/job_starter.sh
index bb0b80e..3af2952 100644
--- a/shell_scripts/job_starter.sh
+++ b/shell_scripts/job_starter.sh
@@ -1,193 +1,185 @@
#!/bin/bash
#wrapper for cluster_runner_....sh which copies job-specific, frequently changing files (e.g. configs.py) before the actual sbatch job
#is submitted since the job might pend in queue before execution --> hazard of job-specific files being unintentionally changed during queue wait time.
#positonal
# -arg #1 identifies the folder name of the dataset-related code (e.g. >toy_exp< or >lidc_exp<) within the code source directory
# -arg #2 is the experiment and first part of the job name,
# optional args and flags:
# -c / --create: (flag) whether to create the exp, i.e., if this is a new start of the exp with configs etc from source dir.
# -f / --folds FOLDS: (option) fold(s) to run on (FOLDS needs to be only one int or string of multiple ints separated by space), default None (-->set to all in config)
# -m / --mode MODE: (option) string, one of "train", "train_test", "test", defaults to "train_test"
# -p / --exp_parent_dir: (option) name of parent_dir rel to dataset folder on cluster. exp_dir is exp_parent_dir/exp_name, if not given defaults to "experiments"
# -q / --queue: (option) which queue (-q parameter for bsub) to send job to. default: gputest. others: gputest-short (max 5h jobs).
# -w / --which: (option) same as argument -m to bsub; host or host list (string separated by space) to send the job to.
# use nodenameXX where XX==nr of node or nodenameXX,nodenameYY,... or nodename[XX-YY]. nodename is e.g. e132-comp.
# --gmem: (option) how much gpu memory to request for job (in gigabytes), defaults to 11.9. Currently, the smaller nodes have 11.9G, the larger ones 31.7G.
# --resume: (flag) only with explicit fold argument, if set, resumes from checkpoint in exp_dir/fold_x/last_state.pth.
# --no_parallel: (flag) if set, folds won't start as parallel jobs on cluster, but run sequentially in one job.
dataset_name="${1}"
exp_name="${2}"
#arguments not passed, e.g. $7 if no seventh argument, are null.
if [ ! -z "${18}" ]; then #-z checks if is null string
echo "Error: Received too many arguments."
exit
fi
#make args optional: move up if some args are missing inbetween
while [ ${#} -gt 2 ]; do
case "${3}" in
-c|--create)
create_exp="c"
shift
;;
-f|--folds)
folds="${4}"
shift; shift
;;
-m|--mode)
mode="${4}"
shift; shift
;;
-p|--exp_parent_dir)
exp_parent_dir="${4}"
shift; shift
;;
-q|--queue)
queue="${4}"
shift; shift
;;
-w|--which)
which="${4}"
shift; shift
;;
-R|--resource)
resource="${4}"
shift; shift
;;
--gmem)
gmem="${4}"
shift; shift
;;
--resume)
resume=true
shift
;;
--no_parallel)
no_parallel=true
shift
;;
*)
echo "Invalid argument/option passed: ${3}"
exit 1
;;
esac
done
# default values
if [ -z ${exp_parent_dir} ]; then
exp_parent_dir="experiments"
fi
if [ -z ${mode} ]; then
mode="train_test"
fi
if [ -z ${queue} ]; then
queue="gputest"
fi
if [ -z ${gmem} ]; then
gmem="11"
fi
root_dir=/home/ramien #assumes /home/ramien exists
-prep_node=ramien@e132-comp07 #node used for prep tasks like create_exp
#medicaldetectiontoolkit
source_dir=${root_dir}/mdt-public
dataset_abs_path=${source_dir}/experiments/${dataset_name} #set as second argument passed to this script
exp_parent_dir=/datasets/datasets_ramien/${dataset_name}/${exp_parent_dir}
-#exp_parent_dir=/home/gregor/Documents/medicaldetectiontoolkit/datasets/${dataset_name}/experiments #for testing this script
-# /dataset is not mounted on log-in/job submission nodes (would maybe make sense, I feel), only on queue gputest's nodes e132-compXX.
-#ssh ${prep_node} "mkdir -p ${exp_parent_dir}"
exp_dir=${exp_parent_dir}/${exp_name}
#activate virtualenv that has all the packages:
source_dl="module load python/3.7.0; module load gcc/7.2.0; source ${root_dir}/.virtualenvs/mdt/bin/activate;"
-# TODO as long as no fix available: this script needs to be started directly from the prep node. :/ would be nice if (most importantly
-# 'module ...') would also work over ssh, but somehow some commands are not availabe over the ssh-induced shell (even when using it as interactive).
eval ${source_dl}
-# ssh: (not working)
-#create_cmd="ssh ${prep_node} '${source_dl} python ${source_dir}/exec.py --server_env --mode create_exp --exp_dir ${exp_dir} --exp_source ${dataset_abs_path};'"
# directly from prep node:
create_cmd="python ${source_dir}/exec.py --server_env --mode create_exp --exp_dir ${exp_dir} --exp_source ${dataset_abs_path};"
#if create_exp, check if would overwrite existing exp_dir
if [ ! -z ${create_exp} ] && [ ${create_exp} = "c" ]; then #-n doesnt work as replacement for !-z
if [ -d ${exp_dir} ]; then
echo "Please confirm to overwrite exp ${exp_name} settings, (Y/n): "; read confirmation
if ([ "${confirmation}" = "y" ] || [ "${confirmation}" = "yes" ] || [ "${confirmation}" = "Y" ] || [ -z "${confirmation}" ]); then
echo "Overwriting ${exp_name}"
else
echo "Exiting due to overwrite denial. Adjust options."
exit
fi
fi
#echo "opts: name ${exp_name}, ${source_dir}/exec.py --server_env --mode create_exp --exp_dir ${exp_dir} --exp_source ${dataset_abs_path}"
echo "Creating ${exp_name}"
eval ${create_cmd}
else
if [ ! -d ${exp_dir} ]; then
echo "Experiment directory ${exp_dir} does not exist."
echo "Run create_exp? (Y/n): "; read confirmation
if ([ "${confirmation}" = "y" ] || [ "${confirmation}" = "yes" ] || [ "${confirmation}" = "Y" ] || [ -z "${confirmation}" ]); then
echo "Creating ${exp_name}"
eval ${create_cmd}
fi
fi
fi
#if not create_exp, check if would overwrite existing folds (possibly valuable trained params!)
if [ -z ${create_exp} ] && ([ ${mode} = "train" ] || [ ${mode} = "train_test" ]) && [ -z "${resume}" ]; then
for f in ${folds}; do #if folds is null this check won't apply and folds will be quietly overwritten.
if [ -d ${exp_dir}/fold_${f} ]; then #-d checks if is dir
echo "please confirm to overwrite fold_${f}, (Y/n):"; read confirmation
if ([ "${confirmation}" = "y" ] || [ "${confirmation}" = "yes" ] || [ "${confirmation}" = "Y" ] || [ -z "${confirmation}" ]); then
echo "Overwriting "${exp_name}/fold_${f}
else
echo "Exiting due to overwrite denial. Adjust options."
exit
fi
fi
done
fi
bsub_opts="bsub -N -q ${queue} -gpu num=1:j_exclusive=yes:mode=exclusive_process:gmem=${gmem}G"
if [ ! -z "$resource" ]; then
bsub_opts=$bsub_opts $resource
fi
if [ ! -z ${which} ]; then
bsub_opts="${bsub_opts} -m ${which}"
fi
#----- parallel/separate fold jobs (each fold in a single job) -----------
if [ ! -z "${folds}" ] && [ -z ${no_parallel} ]; then #WHY do i need to convert to string again?
for f in ${folds}; do
out_file=${exp_dir}/logs/fold_${f}_lsf_output.out
bsub_opts="$bsub_opts -J '${dataset_name} ${exp_name} fold ${f} ${mode}' -oo '${out_file}'"
eval "${bsub_opts} sh cluster_runner_meddec.sh ${source_dir} ${exp_dir} ${dataset_abs_path} ${mode} ${f} ${resume}"
done
#----- consecutive folds job (all folds in one single job) -----------
else
if [ ! -z ${resume} ]; then
echo "You need to explicitly specify folds if you would like to resume from a checkpoint. Exiting."
exit
fi
out_file=${exp_dir}/logs/lsf_output.out
bsub_opts="$bsub_opts -J '${dataset_name} ${exp_name} folds ${folds} ${mode}' -oo '${out_file}'"
eval "${bsub_opts} sh cluster_runner_meddec.sh ${source_dir} ${exp_dir} ${dataset_abs_path} ${mode} ${folds} ${resume}"
echo "Started in no parallel, folds:" ${folds}
fi
diff --git a/unittests.py b/unittests.py
index dadf275..de4c52a 100644
--- a/unittests.py
+++ b/unittests.py
@@ -1,345 +1,404 @@
#!/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 unittest
import os
import pickle
import time
from multiprocessing import Pool
import subprocess
import numpy as np
import pandas as pd
import torch
import torchvision as tv
import tqdm
import utils.exp_utils as utils
import utils.model_utils as mutils
""" Note on unittests: run this file either in the way intended for unittests by starting the script with
python -m unittest unittests.py or start it as a normal python file as python unittests.py.
You can selective run single tests by calling python -m unittest unittests.TestClassOfYourChoice, where
TestClassOfYourChoice is the name of the test defined below, e.g., CompareFoldSplits.
"""
def inspect_info_df(pp_dir):
""" use your debugger to look into the info df of a pp dir.
:param pp_dir: preprocessed-data directory
"""
info_df = pd.read_pickle(os.path.join(pp_dir, "info_df.pickle"))
return
def generate_boxes(count, dim=2, h=100, w=100, d=20, normalize=False, on_grid=False, seed=0):
""" generate boxes of format [y1, x1, y2, x2, (z1, z2)].
:param count: nr of boxes
:param dim: dimension of boxes (2 or 3)
:return: boxes in format (n_boxes, 4 or 6), scores
"""
np.random.seed(seed)
if on_grid:
lower_y = np.random.randint(0, h // 2, (count,))
lower_x = np.random.randint(0, w // 2, (count,))
upper_y = np.random.randint(h // 2, h, (count,))
upper_x = np.random.randint(w // 2, w, (count,))
if dim == 3:
lower_z = np.random.randint(0, d // 2, (count,))
upper_z = np.random.randint(d // 2, d, (count,))
else:
lower_y = np.random.rand(count) * h / 2.
lower_x = np.random.rand(count) * w / 2.
upper_y = (np.random.rand(count) + 1.) * h / 2.
upper_x = (np.random.rand(count) + 1.) * w / 2.
if dim == 3:
lower_z = np.random.rand(count) * d / 2.
upper_z = (np.random.rand(count) + 1.) * d / 2.
if dim == 3:
boxes = np.array(list(zip(lower_y, lower_x, upper_y, upper_x, lower_z, upper_z)))
# add an extreme box that tests the boundaries
boxes = np.concatenate((boxes, np.array([[0., 0., h, w, 0, d]])))
else:
boxes = np.array(list(zip(lower_y, lower_x, upper_y, upper_x)))
boxes = np.concatenate((boxes, np.array([[0., 0., h, w]])))
scores = np.random.rand(count + 1)
if normalize:
divisor = np.array([h, w, h, w, d, d]) if dim == 3 else np.array([h, w, h, w])
boxes = boxes / divisor
return boxes, scores
# -------- check own nms CUDA implement against own numpy implement ------
class CheckNMSImplementation(unittest.TestCase):
@staticmethod
def assert_res_equality(keep_ics1, keep_ics2, boxes, scores, tolerance=0, names=("res1", "res2")):
"""
:param keep_ics1: keep indices (results), torch.Tensor of shape (n_ics,)
:param keep_ics2:
:return:
"""
keep_ics1, keep_ics2 = keep_ics1.cpu().numpy(), keep_ics2.cpu().numpy()
discrepancies = np.setdiff1d(keep_ics1, keep_ics2)
try:
checks = np.array([
len(discrepancies) <= tolerance
])
except:
checks = np.zeros((1,)).astype("bool")
msgs = np.array([
"""{}: {} \n{}: {} \nboxes: {}\n {}\n""".format(names[0], keep_ics1, names[1], keep_ics2, boxes,
scores)
])
assert np.all(checks), "NMS: results mismatch: " + "\n".join(msgs[~checks])
def single_case(self, count=20, dim=3, threshold=0.2, seed=0):
boxes, scores = generate_boxes(count, dim, seed=seed, h=320, w=280, d=30)
keep_numpy = torch.tensor(mutils.nms_numpy(boxes, scores, threshold))
# for some reason torchvision nms requires box coords as floats.
boxes = torch.from_numpy(boxes).type(torch.float32)
scores = torch.from_numpy(scores).type(torch.float32)
if dim == 2:
"""need to wait until next pytorch release where they fixed nms on cpu (currently they have >= where it
needs to be >.)
"""
# keep_ops = tv.ops.nms(boxes, scores, threshold)
# self.assert_res_equality(keep_numpy, keep_ops, boxes, scores, tolerance=0, names=["np", "ops"])
pass
boxes = boxes.cuda()
scores = scores.cuda()
keep = self.nms_ext.nms(boxes, scores, threshold)
self.assert_res_equality(keep_numpy, keep, boxes, scores, tolerance=0, names=["np", "cuda"])
def manual_example(self):
"""
100 x 221 (y, x) image. 5 overlapping boxes, 4 of the same class, 3 of them overlapping above threshold.
"""
threshold = 0.3
boxes = torch.tensor([
[20, 30, 80, 130], #0 reference (needs to have highest score)
[30, 40, 70, 120], #1 IoU 0.35
[10, 50, 90, 80], #2 IoU 0.11
[40, 20, 75, 135], #3 IoU 0.34
[30, 40, 70, 120], #4 IoU 0.35 again but with lower score
]).cuda().float()
scores = torch.tensor([0.71, 0.94, 1.0, 0.82, 0.11]).cuda()
# expected: keep == [1, 2]
keep = self.nms_ext.nms(boxes, scores, threshold)
diff = np.setdiff1d(keep.cpu().numpy(), [1,2])
assert len(diff) == 0, "expected: {}, received: {}.".format([1,2], keep)
def test(self, n_cases=200, box_count=30, threshold=0.5):
# dynamically import module so that it doesn't affect other tests if import fails
self.nms_ext = utils.import_module("nms_ext", 'custom_extensions/nms/nms.py')
self.manual_example()
# change seed to something fix if you want exactly reproducible test
seed0 = np.random.randint(50)
print("NMS test progress (done/total box configurations) 2D:", end="\n")
for i in tqdm.tqdm(range(n_cases)):
self.single_case(count=box_count, dim=2, threshold=threshold, seed=seed0+i)
print("NMS test progress (done/total box configurations) 3D:", end="\n")
for i in tqdm.tqdm(range(n_cases)):
self.single_case(count=box_count, dim=3, threshold=threshold, seed=seed0+i)
return
class CheckRoIAlignImplementation(unittest.TestCase):
def prepare(self, dim=2):
b, c, h, w = 1, 3, 50, 50
# feature map, (b, c, h, w(, z))
if dim == 2:
fmap = torch.rand(b, c, h, w).cuda()
# rois = torch.tensor([[
# [0.1, 0.1, 0.3, 0.3],
# [0.2, 0.2, 0.4, 0.7],
# [0.5, 0.7, 0.7, 0.9],
# ]]).cuda()
pool_size = (7, 7)
rois = generate_boxes(5, dim=dim, h=h, w=w, on_grid=True, seed=np.random.randint(50))[0]
elif dim == 3:
d = 20
fmap = torch.rand(b, c, h, w, d).cuda()
# rois = torch.tensor([[
# [0.1, 0.1, 0.3, 0.3, 0.1, 0.1],
# [0.2, 0.2, 0.4, 0.7, 0.2, 0.4],
# [0.5, 0.0, 0.7, 1.0, 0.4, 0.5],
# [0.0, 0.0, 0.9, 1.0, 0.0, 1.0],
# ]]).cuda()
pool_size = (7, 7, 3)
rois = generate_boxes(5, dim=dim, h=h, w=w, d=d, on_grid=True, seed=np.random.randint(50),
normalize=False)[0]
else:
raise ValueError("dim needs to be 2 or 3")
rois = [torch.from_numpy(rois).type(dtype=torch.float32).cuda(), ]
fmap.requires_grad_(True)
return fmap, rois, pool_size
def check_2d(self):
""" check vs torchvision ops not possible as on purpose different approach.
:return:
"""
raise NotImplementedError
# fmap, rois, pool_size = self.prepare(dim=2)
# ra_object = self.ra_ext.RoIAlign(output_size=pool_size, spatial_scale=1., sampling_ratio=-1)
# align_ext = ra_object(fmap, rois)
# loss_ext = align_ext.sum()
# loss_ext.backward()
#
# rois_swapped = [rois[0][:, [1,3,0,2]]]
# align_ops = tv.ops.roi_align(fmap, rois_swapped, pool_size)
# loss_ops = align_ops.sum()
# loss_ops.backward()
#
# assert (loss_ops == loss_ext), "sum of roialign ops and extension 2D diverges"
# assert (align_ops == align_ext).all(), "ROIAlign failed 2D test"
def check_3d(self):
fmap, rois, pool_size = self.prepare(dim=3)
ra_object = self.ra_ext.RoIAlign(output_size=pool_size, spatial_scale=1., sampling_ratio=-1)
align_ext = ra_object(fmap, rois)
loss_ext = align_ext.sum()
loss_ext.backward()
align_np = mutils.roi_align_3d_numpy(fmap.cpu().detach().numpy(), [roi.cpu().numpy() for roi in rois],
pool_size)
align_np = np.squeeze(align_np) # remove singleton batch dim
align_ext = align_ext.cpu().detach().numpy()
assert np.allclose(align_np, align_ext, rtol=1e-5,
atol=1e-8), "RoIAlign differences in numpy and CUDA implement"
def specific_example_check(self):
# dummy input
self.ra_ext = utils.import_module("ra_ext", 'custom_extensions/roi_align/roi_align.py')
exp = 6
pool_size = (2,2)
fmap = torch.arange(exp**2).view(exp,exp).unsqueeze(0).unsqueeze(0).cuda().type(dtype=torch.float32)
boxes = torch.tensor([[1., 1., 5., 5.]]).cuda()/exp
ind = torch.tensor([0.]*len(boxes)).cuda().type(torch.float32)
y_exp, x_exp = fmap.shape[2:] # exp = expansion
boxes.mul_(torch.tensor([y_exp, x_exp, y_exp, x_exp], dtype=torch.float32).cuda())
boxes = torch.cat((ind.unsqueeze(1), boxes), dim=1)
aligned_tv = tv.ops.roi_align(fmap, boxes, output_size=pool_size, sampling_ratio=-1)
aligned = self.ra_ext.roi_align_2d(fmap, boxes, output_size=pool_size, sampling_ratio=-1)
boxes_3d = torch.cat((boxes, torch.tensor([[-1.,1.]]*len(boxes)).cuda()), dim=1)
fmap_3d = fmap.unsqueeze(dim=-1)
pool_size = (*pool_size,1)
ra_object = self.ra_ext.RoIAlign(output_size=pool_size, spatial_scale=1.,)
aligned_3d = ra_object(fmap_3d, boxes_3d)
# expected_res = torch.tensor([[[[10.5000, 12.5000], # this would be with an alternative grid-point setting
# [22.5000, 24.5000]]]]).cuda()
expected_res = torch.tensor([[[[14., 16.],
[26., 28.]]]]).cuda()
expected_res_3d = torch.tensor([[[[[14.],[16.]],
[[26.],[28.]]]]]).cuda()
assert torch.all(aligned==expected_res), "2D RoIAlign check vs. specific example failed. res: {}\n expected: {}\n".format(aligned, expected_res)
assert torch.all(aligned_3d==expected_res_3d), "3D RoIAlign check vs. specific example failed. res: {}\n expected: {}\n".format(aligned_3d, expected_res_3d)
def test(self):
# dynamically import module so that it doesn't affect other tests if import fails
self.ra_ext = utils.import_module("ra_ext", 'custom_extensions/roi_align/roi_align.py')
self.specific_example_check()
# 2d test
#self.check_2d()
# 3d test
self.check_3d()
return
class VerifyFoldSplits(unittest.TestCase):
""" Check, for a single fold_ids file, i.e., for a single experiment, if the assigned folds (assignment of data
identifiers) is actually incongruent. No overlaps between folds are allowed for a correct cross validation.
"""
@staticmethod
def verify_fold_ids(splits):
"""
Splits: list (n_splits). Each element: list (4) with: 0 == array of train ids, 1 == arr of val ids,
2 == arr of test ids, 3 == int of fold ix.
"""
for f_ix, split_settings in enumerate(splits):
split_ids, fold_ix = split_settings[:3], split_settings[3]
assert f_ix == fold_ix
# check fold ids within their folds
for i, ids1 in enumerate(split_ids):
for j, ids2 in enumerate(split_ids):
if j > i:
inter = np.intersect1d(ids1, ids2)
if len(inter) > 0:
raise Exception("Fold {}: Split {} and {} intersect by pids {}".format(fold_ix, i, j, inter))
# check val and test ids across folds
val_ids = split_ids[1]
test_ids = split_ids[2]
for other_f_ix in range(f_ix + 1, len(splits)):
other_val_ids = splits[other_f_ix][1]
other_test_ids = splits[other_f_ix][2]
inter_val = np.intersect1d(val_ids, other_val_ids)
inter_test = np.intersect1d(test_ids, other_test_ids)
if len(inter_test) > 0:
raise Exception("Folds {} and {}: Test splits intersect by pids {}".format(f_ix, other_f_ix, inter_test))
if len(inter_val) > 0:
raise Exception(
"Folds {} and {}: Val splits intersect by pids {}".format(f_ix, other_f_ix, inter_val))
def test(self):
exp_dir = "/home/gregor/networkdrives/E132-Cluster-Projects/lidc_exp/experiments/042/retinau2d"
check_file = os.path.join(exp_dir, 'fold_ids.pickle')
with open(check_file, 'rb') as handle:
splits = pickle.load(handle)
self.verify_fold_ids(splits)
+class CompareFoldSplits(unittest.TestCase):
+ """ Find evtl. differences in cross-val file splits across different experiments.
+ """
+ @staticmethod
+ def group_id_paths(ref_exp_dir, comp_exp_dirs):
+
+ f_name = 'fold_ids.pickle'
+
+ ref_paths = os.path.join(ref_exp_dir, f_name)
+ assert os.path.isfile(ref_paths), "ref file {} does not exist.".format(ref_paths)
+
+
+ ref_paths = [ref_paths for comp_ed in comp_exp_dirs]
+ comp_paths = [os.path.join(comp_ed, f_name) for comp_ed in comp_exp_dirs]
+
+ return zip(ref_paths, comp_paths)
+
+ @staticmethod
+ def comp_fold_ids(mp_input):
+ fold_ids1, fold_ids2 = mp_input
+ with open(fold_ids1, 'rb') as f:
+ fold_ids1 = pickle.load(f)
+ try:
+ with open(fold_ids2, 'rb') as f:
+ fold_ids2 = pickle.load(f)
+ except FileNotFoundError:
+ print("comp file {} does not exist.".format(fold_ids2))
+ return
+
+ n_splits = len(fold_ids1)
+ assert n_splits == len(fold_ids2), "mismatch n splits: ref has {}, comp {}".format(n_splits, len(fold_ids2))
+ # train, val test
+ split_diffs = np.concatenate([np.setdiff1d(fold_ids1[s][assignment], fold_ids2[s][assignment]) for s in range(n_splits) for assignment in range(3)])
+ all_equal = np.any(split_diffs)
+ return (split_diffs, all_equal)
+
+ def iterate_exp_dirs(self, ref_exp, comp_exps, processes=os.cpu_count()):
+
+ grouped_paths = list(self.group_id_paths(ref_exp, comp_exps))
+ print("performing {} comparisons of cross-val file splits".format(len(grouped_paths)))
+ p = Pool(processes)
+ split_diffs = p.map(self.comp_fold_ids, grouped_paths)
+ p.close(); p.join()
+
+ df = pd.DataFrame(index=range(0,len(grouped_paths)), columns=["ref", "comp", "all_equal"])#, "diffs"])
+ for ix, (ref, comp) in enumerate(grouped_paths):
+ df.iloc[ix] = [ref, comp, split_diffs[ix][1]]#, split_diffs[ix][0]]
+
+ print("Any splits not equal?", df.all_equal.any())
+ assert not df.all_equal.any(), "a split set is different from reference split set, {}".format(df[~df.all_equal])
+
+ def test(self):
+ exp_parent_dir = '/home/gregor/networkdrives/E132-Cluster-Projects/lidc_exp/experiments/1x/adamw_nonorm_nosched'
+ ref_exp = '/media/gregor/HDD1/experiments/mdt/lidc_exp/original_paper_settings'
+ comp_exps = [os.path.join(exp_parent_dir, p) for p in os.listdir(exp_parent_dir)]
+ comp_exps = [p for p in comp_exps if os.path.isdir(p) and p != ref_exp]
+ self.iterate_exp_dirs(ref_exp, comp_exps)
+
+
if __name__=="__main__":
stime = time.time()
unittest.main()
mins, secs = divmod((time.time() - stime), 60)
h, mins = divmod(mins, 60)
t = "{:d}h:{:02d}m:{:02d}s".format(int(h), int(mins), int(secs))
print("{} total runtime: {}".format(os.path.split(__file__)[1], t))
\ No newline at end of file
diff --git a/utils/exp_utils.py b/utils/exp_utils.py
index d0f5264..1e6a07f 100644
--- a/utils/exp_utils.py
+++ b/utils/exp_utils.py
@@ -1,488 +1,506 @@
#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-from typing import Iterable, Tuple, Any
-import sys
+from typing import Iterable, Tuple, Any, Union
+import os, sys
import subprocess
from multiprocessing import Process
-import os
-import plotting
import importlib.util
import pickle
import logging
from torch.utils.tensorboard import SummaryWriter
from collections import OrderedDict
import numpy as np
import torch
import pandas as pd
-def split_off_process(target, *args, daemon=False, **kwargs):
+def split_off_process(target, *args, daemon: bool=False, **kwargs):
"""Start a process that won't block parent script.
No join(), no return value. If daemon=False: before parent exits, it waits for this to finish.
+ :param target: the target function of the process.
+ :params *args: args to pass to target.
+ :param daemon: if False: before parent exits, it waits for this process to finish.
+ :params **kwargs: kwargs to pass to target.
"""
p = Process(target=target, args=tuple(args), kwargs=kwargs, daemon=daemon)
p.start()
return p
+def get_formatted_duration(seconds: float, format: str="hms") -> str:
+ """Format a time in seconds.
+ :param format: "hms" for hours mins secs or "ms" for min secs.
+ """
+ mins, secs = divmod(seconds, 60)
+ if format == "ms":
+ t = "{:d}m:{:02d}s".format(int(mins), int(secs))
+ elif format == "hms":
+ h, mins = divmod(mins, 60)
+ t = "{:d}h:{:02d}m:{:02d}s".format(int(h), int(mins), int(secs))
+ else:
+ raise Exception("Format {} not available, only 'hms' or 'ms'".format(format))
+ return t
+
class CombinedLogger(object):
"""Combine console and tensorboard logger and record system metrics.
"""
- def __init__(self, name, log_dir, server_env=True, fold="all"):
+ def __init__(self, name: str, log_dir: str, server_env: bool=True, fold: Union[int, str]="all"):
self.pylogger = logging.getLogger(name)
self.tboard = SummaryWriter(log_dir=os.path.join(log_dir, "tboard"))
self.log_dir = log_dir
self.fold = str(fold)
self.server_env = server_env
self.pylogger.setLevel(logging.DEBUG)
self.log_file = os.path.join(log_dir, "fold_"+self.fold, 'exec.log')
os.makedirs(os.path.dirname(self.log_file), exist_ok=True)
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
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)
print("logger attr not found")
- def set_logfile(self, fold=None, log_file=None):
+ def set_logfile(self, fold: Union[int, str, None]=None, log_file: Union[str, None]=None):
if fold is not None:
self.fold = str(fold)
if log_file is None:
self.log_file = os.path.join(self.log_dir, "fold_"+self.fold, 'exec.log')
else:
self.log_file = log_file
os.makedirs(os.path.dirname(self.log_file), exist_ok=True)
for hdlr in self.pylogger.handlers:
hdlr.close()
self.pylogger.handlers = []
self.pylogger.addHandler(logging.FileHandler(self.log_file))
if not self.server_env:
self.pylogger.addHandler(ColorHandler())
else:
self.pylogger.addHandler(logging.StreamHandler())
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 = {}
mon_met_series = {}
for tag, val in metrics[key].items():
val = val[-1] # maybe remove list wrapping, recording in evaluator?
if '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 + "/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 __del__(self): # otherwise might produce multiple prints e.g. in ipython console
for hdlr in self.pylogger.handlers:
hdlr.close()
self.pylogger.handlers = []
del self.pylogger
self.tboard.flush()
# close somehow prevents main script from exiting
# maybe revise this issue in a later pytorch version
#self.tboard.close()
-def get_logger(exp_dir, server_env=False):
+def get_logger(exp_dir: str, server_env: bool=False) -> CombinedLogger:
"""
creates logger instance. writing out info to file, to terminal and to tensorboard.
:param exp_dir: experiment directory, where exec.log file is stored.
:param server_env: True if operating in server environment (e.g., gpu cluster)
:return: custom CombinedLogger instance.
"""
log_dir = os.path.join(exp_dir, "logs")
logger = CombinedLogger('medicaldetectiontoolkit', log_dir, server_env=server_env)
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:
"""
if is_training:
if use_stored_settings:
cf_file = import_module('cf_file', 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
os.makedirs(exp_path, exist_ok=True)
# 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_file', 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.created_fold_id_pickle = False
return cf
-def import_module(name, path):
+def import_module(name: str, path: str):
"""
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 set_params_flag(module: torch.nn.Module, flag: Tuple[str, Any], check_overwrite: bool = True):
- """Set an attribute for all module parameters.
+def set_params_flag(module: torch.nn.Module, flag: Tuple[str, Any], check_overwrite: bool = True) -> torch.nn.Module:
+ """Set an attribute for all passed module parameters.
:param flag: tuple (str attribute name : attr value)
:param check_overwrite: if True, assert that attribute not already exists.
"""
for param in module.parameters():
if check_overwrite:
assert not hasattr(param, flag[0]), \
"param {} already has attr {} (w/ val {})".format(param, flag[0], getattr(param, flag[0]))
setattr(param, flag[0], flag[1])
return module
-def parse_params_for_optim(net: torch.nn.Module, weight_decay: float = 0., exclude_from_wd: Iterable = ("norm",)):
- """Format network parameters for the optimizer.
- Convenience function to include options for group-specific settings like weight decay.
- :param net:
- :param weight_decay:
+def parse_params_for_optim(net: torch.nn.Module, weight_decay: float = 0., exclude_from_wd: Iterable = ("norm",)) -> list:
+ """Split network parameters into weight-decay dependent groups for the optimizer.
+ :param net: network.
+ :param weight_decay: weight decay value for the parameters that it is applied to. excluded parameters will have
+ weight decay 0.
:param exclude_from_wd: List of strings of parameter-group names to exclude from weight decay. Options: "norm", "bias".
:return:
"""
+ if weight_decay is None:
+ weight_decay = 0.
# pytorch implements parameter groups as dicts {'params': ...} and
# weight decay as p.data.mul_(1 - group['lr'] * group['weight_decay'])
norm_types = [torch.nn.BatchNorm1d, torch.nn.BatchNorm2d, torch.nn.BatchNorm3d,
torch.nn.InstanceNorm1d, torch.nn.InstanceNorm2d, torch.nn.InstanceNorm3d,
- torch.nn.LayerNorm, torch.nn.GroupNorm, torch.nn.SyncBatchNorm, torch.nn.LocalResponseNorm
- ]
+ torch.nn.LayerNorm, torch.nn.GroupNorm, torch.nn.SyncBatchNorm, torch.nn.LocalResponseNorm]
level_map = {"bias": "weight",
"norm": "module"}
type_map = {"norm": norm_types}
exclude_from_wd = [str(name).lower() for name in exclude_from_wd]
exclude_weight_names = [k for k, v in level_map.items() if k in exclude_from_wd and v == "weight"]
exclude_module_types = tuple([type_ for k, v in level_map.items() if (k in exclude_from_wd and v == "module")
for type_ in type_map[k]])
if exclude_from_wd:
print("excluding {} from weight decay.".format(exclude_from_wd))
for module in net.modules():
if isinstance(module, exclude_module_types):
set_params_flag(module, ("no_wd", True))
for param_name, param in net.named_parameters():
if np.any([ename in param_name for ename in exclude_weight_names]):
setattr(param, "no_wd", True)
with_dec, no_dec = [], []
for param in net.parameters():
if hasattr(param, "no_wd") and param.no_wd == True:
no_dec.append(param)
else:
with_dec.append(param)
orig_ps = sum(p.numel() for p in net.parameters())
with_ps = sum(p.numel() for p in with_dec)
wo_ps = sum(p.numel() for p in no_dec)
assert orig_ps == with_ps + wo_ps, "orig n parameters {} unequals sum of with wd {} and w/o wd {}."\
.format(orig_ps, with_ps, wo_ps)
groups = [{'params': gr, 'weight_decay': wd} for (gr, wd) in [(no_dec, 0.), (with_dec, weight_decay)] if len(gr)>0]
return groups
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):
+ def run_model_selection(self, net: torch.nn.Module, optimizer: torch.optim.Optimizer,
+ monitor_metrics: dict, epoch: int):
# take the mean over all selection criteria in each epoch
non_nan_scores = np.mean(np.array([[0 if (ii is None or np.isnan(ii)) else ii for ii in monitor_metrics['val'][sc]] for sc in self.cf.model_selection_criteria]), 0)
epochs_scores = [ii for ii in non_nan_scores[1:]]
# ranking of epochs according to model_selection_criterion
epoch_ranking = np.argsort(epochs_scores, kind="stable")[::-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 epochs.
if epoch in epoch_ranking[:self.cf.save_n_models]:
save_dir = os.path.join(self.cf.fold_dir, '{}_best_checkpoint'.format(epoch))
if not os.path.exists(save_dir):
os.mkdir(save_dir)
torch.save(net.state_dict(), os.path.join(save_dir, 'params.pth'))
with open(os.path.join(save_dir, 'monitor_metrics.pickle'), 'wb') as handle:
pickle.dump(monitor_metrics, handle)
# save epoch_ranking to keep info for inference.
np.save(os.path.join(self.cf.fold_dir, 'epoch_ranking'), epoch_ranking[:self.cf.save_n_models])
np.save(os.path.join(save_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_checkpoint' in ii]:
if se in epoch_ranking[self.cf.save_n_models:]:
subprocess.call('rm -rf {}'.format(os.path.join(self.cf.fold_dir, '{}_best_checkpoint'.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(),
}
# save checkpoint of current epoch.
save_dir = os.path.join(self.cf.fold_dir, 'last_checkpoint'.format(epoch))
if not os.path.exists(save_dir):
os.mkdir(save_dir)
torch.save(state, os.path.join(save_dir, 'params.pth'))
np.save(os.path.join(save_dir, 'epoch_ranking'), epoch_ranking[:self.cf.save_n_models])
with open(os.path.join(save_dir, 'monitor_metrics.pickle'), 'wb') as handle:
pickle.dump(monitor_metrics, handle)
-def load_checkpoint(checkpoint_path, net, optimizer):
+def load_checkpoint(checkpoint_path: str, net: torch.nn.Module, optimizer: torch.optim.Optimizer) -> Tuple:
checkpoint = torch.load(os.path.join(checkpoint_path, 'params.pth'))
net.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
with open(os.path.join(checkpoint_path, 'monitor_metrics.pickle'), 'rb') as handle:
monitor_metrics = pickle.load(handle)
starting_epoch = checkpoint['epoch'] + 1
return starting_epoch, net, optimizer, monitor_metrics
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()
metrics['val'] = OrderedDict()
metric_classes = []
if 'rois' in cf.report_score_level:
metric_classes.extend([v for k, v in cf.class_dict.items()])
if 'patient' in cf.report_score_level:
metric_classes.extend(['patient'])
for cl in metric_classes:
metrics['train'][cl + '_ap'] = [np.nan]
metrics['val'][cl + '_ap'] = [np.nan]
if cl == 'patient':
metrics['train'][cl + '_auc'] = [np.nan]
metrics['val'][cl + '_auc'] = [np.nan]
return metrics
def create_csv_output(results_list, cf, logger):
"""
Write out test set predictions to .csv file. output format is one line per prediction:
PatientID | PredictionID | [y1 x1 y2 x2 (z1) (z2)] | score | pred_classID
Note, that prediction coordinates correspond to images as loaded for training/testing and need to be adapted when
plotted over raw data (before preprocessing/resampling).
:param results_list: [[patient_results, patient_id], [patient_results, patient_id], ...]
"""
logger.info('creating csv output file at {}'.format(os.path.join(cf.test_dir, 'results.csv')))
predictions_df = pd.DataFrame(columns = ['patientID', 'predictionID', 'coords', 'score', 'pred_classID'])
for r in results_list:
pid = r[1]
#optionally load resampling info from preprocessing to match output predictions with raw data.
#with open(os.path.join(cf.exp_dir, 'test_resampling_info', pid), 'rb') as handle:
# resampling_info = pickle.load(handle)
for bix, box in enumerate(r[0][0]):
if box["box_type"] == "gt":
continue
assert box['box_type'] == 'det', box['box_type']
coords = box['box_coords']
score = box['box_score']
pred_class_id = box['box_pred_class_id']
out_coords = []
if score >= cf.min_det_thresh:
out_coords.append(coords[0]) #* resampling_info['scale'][0])
out_coords.append(coords[1]) #* resampling_info['scale'][1])
out_coords.append(coords[2]) #* resampling_info['scale'][0])
out_coords.append(coords[3]) #* resampling_info['scale'][1])
if len(coords) > 4:
out_coords.append(coords[4]) #* resampling_info['scale'][2] + resampling_info['z_crop'])
out_coords.append(coords[5]) #* resampling_info['scale'][2] + resampling_info['z_crop'])
predictions_df.loc[len(predictions_df)] = [pid, bix, out_coords, score, pred_class_id]
try:
fold = cf.fold
except:
fold = 'hold_out'
predictions_df.to_csv(os.path.join(cf.exp_dir, 'results_{}.csv'.format(fold)), index=False)
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)