Skip to content
Implementation of "Data augmentation using learned transforms for one-shot medical image segmentation"
Python Shell
Branch: master
Clone or download


This repository contains the authors' implementation from "Data Augmentation using Learned Transformations for One-shot Medical Image Segmentation", which will be presented as an oral at CVPR 2019. We provide code for training spatial and appearance transform models, and for using the transform models to synthesize training examples for segmentation.

If you use our code, please cite:

Data augmentation using learned transforms for one-shot medical image segmentation
Amy Zhao, Guha Balakrishnan, Fredo Durand, John Guttag, Adrian V. Dalca
CVPR 2019. eprint arXiv:1902.09383

Getting started

Download dependencies

Run the script This will automatically pull the following dependencies and place them in the correct subdirectories:

Setting up your dataset

We have included a few sample MRI scans (including volumes and segmentations) in the data/ folder. If you wish to use the datasets mentioned in the paper, you should download them directly from the respective dataset sites.

If you wish to use your own dataset, place your volume and segmentation files in the data/ folder. The data loading code in src/ expects each example to be stored as a volume file {example_id}_vol.npz and, if applicable, a corresponding {example_id}_seg.npz file, with the data stored in each file using the keys vol_data and seg_data respectively. The functions load_dataset_files and load_vol_and_seg in src/ can be easily modified to suit your data format.

Training transform models

This repo does not include any pre-trained models. You may train your own spatial and appearance transform models by specifying the GPU ID, dataset name, and the model type.

python trans --gpu 0 --data mri-100unlabeled --model flow-fwd
python trans --gpu 0 --data mri-100unlabeled --model flow-bck
python trans --gpu 0 --data mri-100unlabeled --model color-unet

The results will be placed in experiments/. Note that in order to train an appearance/color transform model, you will need to edit to point at your trained forward/backward spatial transform models.

As described in the paper, each model is implemented using a simple architecture based on U-Net. You can change hyperparameters by modifying transform_model_arch_params in We encourage you to experiment with your favorite model architecture, and to adjust the model parameters to suit your dataset.

Training a segmentation network

You may train a segmentation model by specifying the GPU ID and dataset name.

python fss --gpu 0 --data mri-100unlabeled

Again, results will be placed under experiments/.

You can use additional flags:

  • --aug_rand will apply random augmentation to each training example consisting of a random smooth deformation and a random global multiplicative intensity factor.
  • --aug_sas will pseudo-label any unlabeled examples in the training set using the specified spatial registration model.
  • --aug_tm will synthesize training examples using our method.

If you wish to use --aug_sas or --aug_tm, it is important to specify the spatial and appearance transform models to use in seg_model_arch_params in


To evaluate trained segmenters, look at the code in You will have to modify the code to point at your trained models.

Repo name inspired by Magic: The Gathering.


You can’t perform that action at this time.