This repository investigates the influence of different data augmentation strategies on MRI training performance.

Citation

If you use SmaugLab, please make sure to cite the following paper:

@InProceedings{moliniermoller_MICCAI2026_SmaugLab,
        author = { Molinier, Nathan AND Möller, Hendrik AND Dagonneau, Thomas AND Curto-Vilalta, Anna AND Graf, Robert AND Atad, Matan AND Rueckert, Daniel AND Kirschke, Jan S. AND Cohen-Adad, Julien},
        title = { { One Sequence to Segment Them All: Efficient Data Augmentation for CT and MRI Cross-Domain 3D Spine Segmentation } },
        booktitle = {Medical Image Computing and Computer Assisted Intervention -- MICCAI 2026},
        year = {2026},
        publisher = {Springer Nature Switzerland},
        volume = {LNCS 16881},
        month = {September},
        page = {pending}
}

What is available ?

This repository contains:

• A nnUNet trainer with extensive data augmentations
• A basic Monai segmentation script incorporating data augmentations
• A script generating augmentations from input images and segmentations

How to install ?

1. Open a bash terminal in the directory where you want to work.

2. Create and activate a virtual environment using python >=3.10 (highly recommended):

   • venv

   python3 -m venv venv
   source venv/bin/activate

   • conda env

   conda create -n myenv python=3.10
   conda activate myenv
   

3. Clone this repository:

   • Git clone

   git clone git@github.com:neuropoly/SmaugLab.git
   cd SmaugLab

4. Install SmaugLab using one of the following commands:

   Note: If you pull a new version from GitHub, make sure to rerun this command with the flag --upgrade

   • nnunetv2 only usage (tested with nnunetv2==2.6.2)

   python3 -m pip install -e . nnunetv2==2.6.2

   • full usage (with Monai and other dependencies)

   python3 -m pip install -e .[all]

5. Install PyTorch following the instructions on their website. Be sure to add the --upgrade flag to your installation command to replace any existing PyTorch installation. Example:

python3 -m pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu118 --upgrade

Run nnunet training with SmaugLab trainer

To use the SmaugLab trainer with nnUNet, first add the trainer to your nnUNet installation by running:

smauglab_add_nnunettrainer --trainer nnUNetTrainerDAExt

Then, when you run nnUNet training as usual, specifying the SmaugLab trainer, for example:

nnUNetv2_train 100 3d_fullres 0 -tr nnUNetTrainerDAExtGPU -p nnUNetPlans

You can also specify your data augmentation parameters by providing a JSON file using the environment variable SMAUGLAB_PARAMS_GPU_JSON:

Note: By default smauglab/configs/transform_params_gpu.json is used if no file is specified.

SMAUGLAB_PARAMS_GPU_JSON=/path/to/your/params.json nnUNetv2_train 100 3d_fullres 0 -tr nnUNetTrainerDAExtGPU -p nnUNetPlans

⚠️ Warning : To avoid any paths issues, please specify an absolute path to your JSON file.

Mixed-precision gradient scale

The trainer floors the AMP gradient scale at 1.0. torch's GradScaler halves the scale on every step with a non-finite gradient and only doubles it back after 2000 consecutive clean steps, with no floor underneath, so a bad-batch rate above 1-in-2000 drives it to exactly 0.0 -- at which point every gradient becomes NaN, clean batches included, and the run cannot recover while still logging epochs. Below 1.0 the scaler can only shrink gradients, so the floor costs nothing.

SMAUGLAB_MIN_GRAD_SCALE overrides it, for the one case the default forecloses: gradients that genuinely overflow fp16 at scale 1, where halving is the correct response.

SMAUGLAB_MIN_GRAD_SCALE=0.25 nnUNetv2_train 100 3d_fullres 0 -tr nnUNetTrainerDAExtGPU -p nnUNetPlans

A value at or below zero is rejected, since that is the collapse the floor exists to prevent, and a value below 1.0 warns: a floor deep in the denormals trains on all-zero gradients without ever reporting a non-finite step. Resets are written to the training log and counted in the epoch summary either way.

Run Monai training with SmaugLab augmentations

To use SmaugLab augmentations in a MONAI training pipeline, refer to the example training script. Key implementation lines required for proper integration are marked with a 🐞 emoji in the comments.

SmaugLab's augmentations reach a MONAI pipeline as the GPU pipeline: MONAI loads the data and does the deterministic preprocessing, then AugTransformsGPU(...).to(device) augments each collated batch on device inside the training loop. Every augmentation in the GPU column of Available augmentations is therefore available here. There is no MONAI-native backend, and none is planned: the CPU column is batchgeneratorsv2 transforms for the nnU-Net dataloader, not MONAI MapTransforms.

To run the Monai training script directly, you need to provide a config JSON (config.json) file with paths to the images and labels (ground truth) for TRAINING, VALIDATION and TESTING sets like this:

{
   "TYPE": "LABEL",
   "TRAINING": [
      {
         "IMAGE": "/path/to/image1.nii.gz",
         "LABEL": "/path/to/label1.nii.gz"
      },
      {
         "IMAGE": "/path/to/image2.nii.gz",
         "LABEL": "/path/to/label2.nii.gz"
      }
   ],
   "VALIDATION": [
      {
         "IMAGE": "/path/to/image3.nii.gz",
         "LABEL": "/path/to/label3.nii.gz"
      },
      {
         "IMAGE": "/path/to/image4.nii.gz",
         "LABEL": "/path/to/label4.nii.gz"
      }
   ],
   "TESTING": [
      {
         "IMAGE": "/path/to/image5.nii.gz",
         "LABEL": "/path/to/label5.nii.gz"
      },
   ]
}

Then run the training script with the following command, specifying the path to your config JSON file and the path to your data augmentation parameters JSON file (if you want to use custom parameters, otherwise the default transform_params_gpu.json is used):

python scripts/train_monai.py --config <your_path>/config.json --transforms <your_path>/transform_params_gpu.json

Additional parameters can be specified—see python scripts/train_monai.py -h for details. If anything is unclear, feel free to open an issue.

Contributing

Development setup, the test suite, and the release process are documented in CONTRIBUTING.md. The short version:

pip install -e ".[dev]"
pre-commit install
pytest

Pull requests are gated on Ruff (lint + format) and the test suite across Python 3.10–3.12.

How to use my data ?

Scripts developped in this repository use JSON files to specify image and segmentation paths: see this example.

How do I specify my parameters ?

To track parameters used during data augmentation, JSON files are also used: see this example

Paper configs

Two shipped configs are the setups behind the paper's tables. smauglab hash <config> prints the identity below, which is how a run folder named ...-c-<config_hash> is traced back to the file that produced it.

shipped config	config hash	paper setup
smauglab/configs/transform_params_one-sequence-to-segment-them-all.json	c69b9872	Ours
smauglab/configs/transform_params_paper-synthseg.json	580beb13	SynthSeg

The SynthSeg config hashed to 0d5d802f until em_merge_prob: 0.0 was added to it; the two files generate identically, so a run folder named ...-c-0d5d802f is the same setup.

Despite the name, "SynthSeg" is not "Ours plus SynthSeg": it is Ours with all six transfer augmentations (RandomScharrGPU, RandomUnsharpMaskGPU, RandomRandConvGPU, RandomRedistributeSegGPU, RandomInverseGPU, RandomHistogramEqualizationGPU) at p=0 and RandomSynthSegGPU at p=1.0 in their place. The general-enhancement half is identical between the two.

Available augmentations

Which augmentations exist, and which of the two backends implements each one. An empty cell means no implementation on that backend yet. MONAI training draws from the GPU column -- see "Run Monai training with SmaugLab augmentations" above. Regenerate with smauglab matrix --write.

Augmentation	Group	GPU	CPU
flip	GEO	RandomFlipTransformGPU	—
affine	GEO	RandomAffineGPU	—
crop	GEO	RandomCropTransformGPU	—
spatial	GEO	—	SpatialTransform
gaussian_noise	GE	RandomGaussianNoiseGPU	GaussianNoiseTransform
gaussian_blur	GE	RandomGaussianBlurGPU	GaussianBlurTransform
brightness	GE	RandomBrightnessGPU	MultiplicativeBrightnessTransform
contrast	GE	RandomContrastGPU	ContrastTransform
gamma	GE	RandomGammaGPU	GammaTransform
inv_gamma	GE	RandomInvGammaGPU	InvertedGammaTransform
clamp	GE	RandomClampGPU	—
low_res	GE	RandomLowResTransformGPU	SimulateLowResolutionTransform
acq	GE	RandomAcqTransformGPU	—
zscore	GE	ZscoreNormalizationGPU	ZscoreNormalization
mirror	GEO	—	MirrorTransform
scharr	TA	RandomScharrGPU	ScharrConvTransform
laplace	TA	RandomLaplaceGPU	LaplaceConvTransform
unsharp_mask	TA	RandomUnsharpMaskGPU	—
rand_conv	TA	RandomRandConvGPU	—
bias_field	TA	RandomBiasFieldGPU	—
inverse	TA	RandomInverseGPU	—
histogram_equal	TA	RandomHistogramEqualizationGPU	HistogramEqualTransform
redistribute_seg	TA	RandomRedistributeSegGPU	RedistributeTransform
palette	TA	RandomPaletteGPU	—
palette_composed	TA	PaletteSynthesisGPU	—
domain_transfer	TA	RandomDomainTransferGPU	—
synthseg	TA	RandomSynthSegGPU	—
artifact	TA	—	ArtifactTransform
spatial_custom	GEO	—	SpatialCustomTransform
shape	GE	—	ShapeTransform
func_log1p	TA	RandomLog1pGPU	Log1pTransform
func_sqrt	TA	RandomSqrtGPU	SqrtTransform
func_sin	TA	RandomSinGPU	SinTransform
func_exp	TA	RandomExpGPU	ExpTransform
func_sigmoid	TA	RandomSigmoidGPU	SigmoidTransform