Metadata-Version: 2.4
Name: smauglab
Version: 2.0.0
Summary: SmaugLab investigates the influence of different data augmentation strategies on MRI training performance.
License-File: LICENSE
Keywords: deep learning,image segmentation,nnU-Net,nnunet,magnetic resonance imaging,online augmentation,offline augmentation,mri
Author: Nathan Molinier
Author-email: nathan.molinier@polymtl.ca
Maintainer: Hendrik Möller
Requires-Python: >=3.10
Classifier: Development Status :: 5 - Production/Stable
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Programming Language :: Python :: 3.15
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Image Recognition
Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
Provides-Extra: all
Provides-Extra: dev
Provides-Extra: nnunetv2
Requires-Dist: batchgenerators
Requires-Dist: batchgeneratorsv2 (<0.3.6)
Requires-Dist: build ; extra == "dev"
Requires-Dist: coverage (>=7) ; extra == "dev"
Requires-Dist: kornia (>=0.7.3,<0.9)
Requires-Dist: monai[all] ; extra == "all"
Requires-Dist: mypy (==2.4.0) ; extra == "dev"
Requires-Dist: nibabel
Requires-Dist: nnunetv2 (>=2.5) ; extra == "nnunetv2"
Requires-Dist: numpy
Requires-Dist: pre-commit ; extra == "dev"
Requires-Dist: progress
Requires-Dist: pytest (>=8) ; extra == "dev"
Requires-Dist: pytest-cov ; extra == "dev"
Requires-Dist: ruff (==0.16.1) ; extra == "dev"
Requires-Dist: scipy
Requires-Dist: torchio
Requires-Dist: torchvision
Requires-Dist: tqdm ; extra == "all"
Requires-Dist: twine ; extra == "dev"
Requires-Dist: wandb ; extra == "all"
Project-URL: Homepage, https://github.com/neuropoly/SmaugLab
Project-URL: Repository, https://github.com/neuropoly/SmaugLab
Description-Content-Type: text/markdown

[![paper](https://img.shields.io/badge/Paper-miccai2026-orange)](https://papers.miccai.org/miccai-2026/paper/1360_paper.pdf)
[![PyPI](https://img.shields.io/pypi/v/smauglab)](https://pypi.org/project/smauglab/)
[![Python Versions](https://img.shields.io/pypi/pyversions/smauglab)](https://pypi.org/project/smauglab/)
[![tests](https://github.com/neuropoly/SmaugLab/actions/workflows/tests.yml/badge.svg)](https://github.com/neuropoly/SmaugLab/actions/workflows/tests.yml)
[![lint](https://github.com/neuropoly/SmaugLab/actions/workflows/lint.yml/badge.svg)](https://github.com/neuropoly/SmaugLab/actions/workflows/lint.yml)
[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
[![License: GPL v3](https://img.shields.io/badge/License-GPLv3-blue.svg)](https://www.gnu.org/licenses/gpl-3.0)

<h1 align="center">
<img src="smauglab_logo.png" width="500">
</h1>

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](https://github.com/neuropoly/SmaugLab/blob/bed6c1b5cf8ec3dbe6165daca507bf431cad65e5/smauglab/trainers/nnUNetTrainerDAExt.py) with extensive data augmentations
- A basic Monai segmentation [script](https://github.com/neuropoly/SmaugLab/blob/bed6c1b5cf8ec3dbe6165daca507bf431cad65e5/scripts/train_monai.py) incorporating data augmentations
- A [script](https://github.com/neuropoly/SmaugLab/blob/bed6c1b5cf8ec3dbe6165daca507bf431cad65e5/scripts/generate_augmentations.py) 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
   ```bash
   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
   ```bash
   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)
   ```bash
   python3 -m pip install -e . nnunetv2==2.6.2
   ```
   - full usage (with Monai and other dependencies)
   ```bash
   python3 -m pip install -e .[all]
   ```

5. Install PyTorch following the instructions on their [website](https://pytorch.org/). Be sure to add the `--upgrade` flag to your installation command to replace any existing PyTorch installation.
   Example:
```bash
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:
```bash
smauglab_add_nnunettrainer --trainer nnUNetTrainerDAExt
```

Then, when you run nnUNet training as usual, specifying the SmaugLab trainer, for example:
```bash
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](https://github.com/neuropoly/SmaugLab/blob/main/smauglab/configs/transform_params_gpu.json) is used if no file is specified.
```bash
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.

```bash
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](https://github.com/neuropoly/SmaugLab/blob/main/scripts/train_monai.py). 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](#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 `MapTransform`s.

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:
```json
{
   "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](https://github.com/neuropoly/SmaugLab/blob/main/smauglab/configs/transform_params_gpu.json) is used):
```bash
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](CONTRIBUTING.md). The short version:

```bash
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](https://github.com/neuropoly/SmaugLab/blob/16653a84e031c40e25a72e946c2724494606b21c/smauglab/configs/data/data.json).

## How do I specify my parameters ?

To track parameters used during data augmentation, JSON files are also used: see this [example](https://github.com/neuropoly/SmaugLab/blob/16653a84e031c40e25a72e946c2724494606b21c/smauglab/configs/transform_params.json)

### 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`.

<!-- BEGIN AUG MATRIX (generated by `smauglab matrix --write`; do not edit) -->
| 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` |
<!-- END AUG MATRIX -->

