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:

@article{molinier2026one,
  title={One Sequence to Segment Them All: Efficient Data Augmentation for CT and MRI Cross-Domain 3D Spine Segmentation},
  author={Molinier, Nathan and M{\"o}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},
  journal={arXiv preprint arXiv:2605.03098},
  year={2026}
}

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.

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.

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

Citation

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

@article{molinier2026one,
  title={One Sequence to Segment Them All: Efficient Data Augmentation for CT and MRI Cross-Domain 3D Spine Segmentation},
  author={Molinier, Nathan and M{\"o}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},
  journal={arXiv preprint arXiv:2605.03098},
  year={2026}
}

Available augmentations

Which augmentations exist, and which backends implement each one. An empty cell means no implementation on that backend yet. Regenerate with smauglab matrix --write.

Augmentation	Group	GPU	CPU	MONAI
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	—