Advice on integrating custom intensity augmentation into nnU-Net

[Le0p]

Hi everyone,

First of all, thank you very much for developing and maintaining nnU-Net — it has been a key component of our work.

TL;DR: I observe a consistent performance drop when moving a contrast-based MRI augmentation from offline generation to nnU-Net’s on-the-fly transform pipeline, despite matching normalization and randomness as closely as possible, and I am looking for guidance on what might be missing or how on-the-fly augmentations should ideally be implemented in nnU-Net.

This question comes as a continuation of a project using nnU-Net for cardiac MRI segmentation, in the context of automatic cardiac function assessment. Our work builds on a pipeline similar to the one described in:
Cardiac Function Assessment with Deep-Learning-Based Automatic Segmentation of Free-Running 4D Whole-Heart CMR

In this setting, robustness to contrast variations across MRI acquisitions is particularly important.

Problem description
I am currently working on making a segmentation model more contrast-agnostic for MRI data, using an intensity augmentation strategy that modifies image contrast conditioned on the segmentation (i.e. anatomy-informed contrast augmentation).
Concretely, the augmentation changes the contrast of different anatomical regions independently based on the segmentation map.

Offline augmentation (baseline)
Initially, I implemented this augmentation offline:
• Augmented images were generated beforehand using fixed random parameters
• A new dataset was created from these augmented images
• nnU-Net was trained on this fixed augmented dataset
With this setup, I obtained clearly improved performance, and the resulting model showed better robustness to contrast variations.

On-the-fly augmentation attempt
Based on these results (and following the implementation idea in here https://github.com/MIC-DKFZ/anatomy_informed_DA )), I then tried to integrate the same augmentation into the "on-the-fly" transform pipeline in nnU-Net (via batchgeneratorsv2), with the goal that:
• Each image is augmented just before being fed to the network
• The network sees a larger diversity of contrasts during training
• Iterating on the augmentation becomes easier without regenerating the dataset
However, when switching to on-the-fly augmentation, I consistently observed worse performance compared to the offline version.

Investigation and normalization
While investigating potential causes, I noticed that nnU-Net applies z-score normalization to MRI images as part of its preprocessing pipeline.
To better match this behavior online, I implemented an additional transform so that the contrast-augmented images are explicitly re-normalized (z-score) after the augmentation.
This improved the results slightly, but they still remained significantly below the offline setup.
Trying to match the offline setup as closely as possible
To further reduce differences, I also tried:
• Using a fixed set of augmentation parameters per image, effectively mimicking offline augmentation
• Keeping everything else identical (architecture, training schedule, dataset split, etc.)
Despite this, the performance gap persisted.

Observed difference
In practice, the on-the-fly version is typically 5–8% lower in Dice score compared to the offline-augmented training, which seems larger than what I would expect from stochasticity alone.

Question / request for insight
At this point, I am wondering whether there are aspects of the nnU-Net preprocessing or training pipeline that I may be overlooking, for example:
• Interactions between preprocessing, normalization, and intensity-based transforms
• Differences between how offline data and on-the-fly augmented data are handled internally
I would greatly appreciate any insight or advice on what I might have missed.
I can of course provide more details (code snippets, exact numbers, or a minimal reproducible example) if helpful.

Thanks a lot for your time and for your work on nnU-Net.

Leonard
MSc Student at CIBM MRI CHUV-UNIL