Adding new data augmentations

[shatz01]

Looking for general advice to add (some fairly complex) data augmentations to nnUNet.
If there is already documentation for doing this please point me to it, but I only found one line in documentation/extending_nnunet.md which said to look in nnunet.training.network_training.data_augmentation.*. That points me to the network trainers, as there is no data_augmentation/ folder or file (at least not anymore). My assumption was that our data augmentations to the training data must happen before nnUNet is even aware of our network trainer.

Clearly, im very confused.

For some clarification, I am working on the kits21 data, and we already have a baseline nnUNet working with nnUNetTrainerV2. So I think I should modify the nnUNetTrainerV2 file?

Some questions:

1. should my augmentation happen before or after running nnUNet_plan_and_preprocess? Is nnUNet_plan_and_preprocess even aware of which network trainer we use?
2. Is there a way of testing these in an interactive environment such as jupyter notebook? I will be building fairly complex augmentations.

Even if you cant answer those questions, just giving me a direction to go would be a big help, thanks!

[FabianIsensee]

Augmentation should never be done offline. This just wastes disk space. You can do this on the fly during training. Thus, the trainer must be aware of it and implement it. The configuration of the data augmentation can depend on the plans and is thus part of the rule-based system nnU-Net is based on.
Augmentation is set in

nnUNet/nnunet/training/network_training/nnUNetTrainerV2.py

Line 55 in fd58e25

def initialize(self, training=True, force_load_plans=False):

. You should be able to follow the code from there.
Best,
Fabian