MIC-DKFZ / anatomy_informed_DA

Apache License 2.0
18 stars 1 forks source link

Anatomy-informed Augmentation

The data augmentation is presented on MICCAI 2023 in the proceedings of 'Anatomy-informed Data Augmentation for enhanced Prostate Cancer Detection'.

If you use this augmentation please cite the following work:

Kovacs, Balint, et al. "Anatomy-Informed Data Augmentation for Enhanced Prostate Cancer Detection."
International Conference on Medical Image Computing and Computer-Assisted Intervention.
Cham: Springer Nature Switzerland, 2023.

Usage

To use the anatomy-informed augmentation during network training, your segmentation files need additional organ segmentations. In this current example, we add the rectum and bladder segmentations additionally to the prostate cancer segmentations (not all organ segmentations are necessary, it can also be used individually). In that case, your dataset.json and segmentation files should contain the following labels:

...
"labels": {
    "background": "0",
    "malignant prostate lesion": "1",
    "rectum": "2",
    "bladder": "3" 
},...

The nnU-Net trainer nnUNetTrainer_AnatInf.py is the version for the training utilizing the anatomy-informed augmentation. During training the rectum and the bladder are realistically deformed to provide diverse organ & lesion morphologies resulted from physiological soft tissue deformations. There are only 2 modifications in the trainer: 1) tr_transforms is appended with the AnatomyInformedTransform transformation with the following parameters:

Important suggestions for its successful utilization:

Furthermore, the additional organ labels of the class variable _all_labels of the PlansManager class in utilities/plans_handling/plans_handler.py are also removed to reduce the number of output channels of the network.

label_manager_instance._all_labels = label_manager_instance._all_labels[:2]

Installation

The repository is forked from nnU-Net v2, but contains no pretrained network so it can be used as an integrative framework for model developement 1) Create a new conda new conda environment with the recent version of Python (nnU-Net v2 supports 3.9 or newer version), as an example: conda create --name nnUNet_AnatInf python=3.9 2) Install pytorch with the most recent CUDA version by following the instructions on the PyTorch Website. 3) Clone this repository and install its dependencies:

https://github.com/MIC-DKFZ/anatomy_informed_DA.git
cd anatomy-informed-da
pip install -e .

You can find more information in the nnU-Net installation instructions.

Copyright

Copyright German Cancer Research Center (DKFZ) and contributors. Please make sure that your usage of this code is in compliance with its license.