0liliulei / mem3d

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}
}

mem3d's Issues

关于𝑓IN输入通道数目的疑惑

对于2D interaction network(𝑓IN)的输入，论文中提到的输入是3通道，分别是灰度图片、先前的分割结果和cue map进行叠加得到3通道

但从代码（train_dextr.py、train_inter.py、train_scribble.py）中的nInputChannels来看，输入通道的数目是4。

请问怎么样理解这个4，这个4是由什么叠加而成的

Maybe someday I'll merge the files of five 2D-interactive-networks into one.

Mistakes in data using

As said in the paper, for the MSD dataset, lung(64/32 for train and validation), colon(126/64for train and validation). But in the code, the author only give ImageSets10(100/26 for train and validation), there is no ImageSets06.

The authors should respond to this ambiguity.

Manually marked mask for the val set of two tasks in MSD

I saw in the article that you manually annotated the val sets of two tasks in MSD. This is essential information to evaluate the effect of interactive segmentation model. Can you provide the mask you manually annotated for the val sets of two tasks in MSD? Thank you very much.

Manual scribble annotations

Hi,

In the journal paper you mentioned that you will make the manual scribble annotators for MSD and KiTS publicly available. Any updates on this?

Thanks!

Running the code on my dataset

Could you please clarify the components of each sample in the training batch: frame, mask, num_obj, and info? I'm unsure about their sizes and types, as it's unclear how they are generated in the provided dataloader. Some things in data.py are confusing to me, and I would appreciate further explanation.

What's the difference between task06_mask and task06_orignial?

Could you provide the data pre-processing code?

Thanks for sharing the code. It is such a solid and wonderful work.
I want to reproduce the code in my own datasets. Could you please share the pre-processing code to the repositiory?

pretrained models??

Hello, dear author
Thanks very much for opening source your code!

I tried to use your code to train 2D Interactive Network, while it required some pretrained models.
Would you please provide some models?

BTW, when we train the models of 2D Interactive Network and Memory Network, are they independent from each other?
Or they should be trained in order?