Volumetric memory network for interactive medical image segmentation

We propose a novel memory-augmented network named VMN for interactive segmentation of volumetric medical data.

Paper

This repository provides the official PyTorch implementation of VMN in the following papers:

Volumetric memory network for interactive medical image segmentation
Tianfei Zhou, Liulei Li, Gustav Bredell, Jianwu Li, and Ender Konukoglu
Biomedical Image Computing, CVL, ETH Zurich | Beijing Institute of Technology
Medical Image Analysis (MedIA) [Paper]
Elsevier-MedIA Best Paper Award

Quality-Aware Memory Network for Interactive Volumetric Image Segmentation
Tianfei Zhou, Liulei Li, Gustav Bredell, Jianwu Li, and Ender Konukoglu
Biomedical Image Computing, CVL, ETH Zurich | Beijing Institute of Technology
International Conference on Medical Image Computing and Computer Assisted Intervention (MICCAI) [Paper]

Preparation

Dataset Download

Download MSD and KiTS. This repo provides dataloaders for MSD, you can some modification to adapt them to other datasets.

Dataset Organization

To run the training and testing code, we require the following data organization format

${ROOT}--
        |--KiTS
        |--MSD
        │   ├── ImageSets06
        │   │   └── train.txt
        │   │   └── test.txt
        │   ├── ImageSest10
        │   ├── Task06_mask
        │   │   ├── lung_001
        │   │   │   ├── 0.png 
        │   │   │   ├── ...
        │   │   │   └── 199.png
        │   │   ├── lung_002
        │   │   ├── ...
        │   │   └── lung_060
        │   ├── Task06_origin
        │   │   ├── lung_001
        │   │   │   ├── 0.png 
        │   │   │   ├── ...
        │   │   │   └── 199.png
        │   │   ├── ...
        │   │   └── lung_060
        │   ├── ImageSets10
        │   ├── Task10_mask
        │   └── Task10_origin
        └──${DATASET3}

Download Pretrained Weights

• Download the weight pretrained on YouTube-VOS for VMN
• Update the initial attribution in option.py

Training and Testing

• 2D Interactive Network

    Mem3D/
    └── (train/test)_(dextr/hybrid/inter/scribble/two_point).py

• Volumetric Memory Network

    Mem3D/
    ├── (train/test)_STM.py.  # without Quality Assessment
    └── train_SAQ.py          # with Quality Assessment

• Round Based 3D Interactive Segmentation

    Mem3D/
    ├── eval_SAQ.py               # w QA
    └── eval_IOG_refine_dextr.py  # w/o QA

• Volume-wise Dice Evaluation

    Mem3D/
    └── eval.py

Acknowledgements

• This code is based on DEXTR and STM.

Citation

If you use VMN for your research, please cite our papers:

@article{zhou2022volumetric,
  title={Volumetric memory network for interactive medical image segmentation},
  author={Zhou, Tianfei and Li, Liulei and Bredell, Gustav and Li, Jianwu and Konukoglu, Ender},
  journal={Medical Image Analysis},
  year={2022},
  publisher={Elsevier}
}

@inproceedings{zhou2021quality,
  title={Quality-aware memory network for interactive volumetric image segmentation},
  author={Zhou, Tianfei and Li, Liulei and Bredell, Gustav and Li, Jianwu and Konukoglu, Ender},
  booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
  pages={560--570},
  year={2021},
  organization={Springer}
}