Towards Robust General Medical Image Segmentation

This repository provides a PyTorch implementation of the Benchmark and baseline method presented in Towards Robust General Medical Image Segmentation (MICCAI 2021). In this work, we propose a new benchmark to evaluate robustness in the context of the Medical Segmentation Decathlon (MSD) by extending the recent AutoAttack natural image classification framework to the domain of volumetric data segmentation, and we present a novel lattice architecture for RObust Generic medical image segmentation (ROG).

For mor information, please visit our project website

Paper

Towards Robust General Medical Image Segmentation,
Laura Daza¹, Juan C. Pérez^1,2, Pablo Arbeláez¹*
MICCAI 2021.

¹ Center for Research and Formation in Artificial Intelligence (CINFONIA), Universidad de Los Andes.
² King Abdullah University of Science and Technology (KAUST).

Installation

Cloning the repository

$ git clone [email protected]:BCV-Uniandes/ROG.git
$ cd ROG
$ python setup.py install

Dataset Preparation

1. Download the Medical Segmentation Decathlon (MSD) Dataset from here. Each task will be organized with the following structure:

TaskXX_TaskName
|_ imagesTr
|_ |_ *.nii.gz
|_ imagesTs
|_ |_ *.nii.gz
|_ labelsTr
|_ |_ *.nii.gz
|_ dataset.json

2. Set the data_root, out_directory and num_workers variables in the file data_preprocessing.py and run the following command:

python data_preprocessing.py

If you want to use ROG for a different dataset you can store it in the same data_root folder, making sure that it follows the same name formating and organization as the MSD tasks.

Training and evaluating the models

We train ROG on clean images and then fine-tune the models using Free AT:

# For the standard training
python main.py --task TASK_ID --gpu GPU_IDs --batch BATCH_SIZE --name OUTPUT_DIR

# For the Free AT fine tuning
python main.py --task TASK_ID --gpu GPU_IDs --batch BATCH_SIZE --name OUTPUT_DIR_FREE_AT --ft --pretrained OUTPUT_DIR --AT

Evaluating the models:

# Standard inference
python main.py --task TASK_ID --gpu GPU_IDs --batch BATCH_SIZE --name OUTPUT_DIR --test

# AutoAttack
python main.py --task TASK_ID --gpu GPU_IDs --batch BATCH_SIZE --name OUTPUT_DIR --test --adv

You can change the strength of the attacks by adjusting the magnitude (--eps) and number of iterations (--adv_iters).

Citation

If you find our paper useful, please use the following BibTeX entry for citation:

@inproceedings{daza2021towards,
  title={Towards Robust General Medical Image Segmentation},
  author={Daza, Laura and P{\'e}rez, Juan C and Arbel{\'a}ez, Pablo},
  booktitle={MICCAI},
  year={2021}
}