uncbiag / iSegFormer

iSegFormer: Interactive Image/Volume Segmentation using Vision Transformers (MICCAI 2022)

Exploring Cycle Consistency Learning in Interactive Volume Segmentation

PyTorch implementation for paper Exploring Cycle Consistency Learning in Interactive Volume Segmentation.

Qin Liu¹, Meng Zheng², Benjamin Planche², Zhongpai Gao², Terrence Chen², Marc Niethammer¹, Ziyan Wu^2
1UNC-Chapel Hill, ²United Imaging Intelligence

Paper | Demo Videos

This repository also contains the following work (branch v1.0):

iSegFormer: Interactive Image Segmentation via Transformers with Application to 3D Knee MR Images. (MICCAI 2022)
Qin Liu, Zhenlin Xu, Yining Jiao, Marc Niethammer
UNC-Chapel Hill

Installation

The code is tested with python=3.9, torch=1.12.0, and torchvision=0.13.0 on an A6000 GPU.

git clone https://github.com/uncbiag/iSegFormer
cd iSegFormer

Now, create a new conda environment and install required packages accordingly.

conda create -n isegformer python=3.9
conda activate isegformer
pip3 install -r requirements.txt

Getting Started

First, download AbdomenCT-1K dataset and model weights. AbdomenCT-1K will be saved in the data folder; model weights will be saved in the saves folder.

python download.py

Unzip the AbdomenCT-1K in the data folder accordingly. Then run a demo:

./run_demo.sh

You will get a GUI as below:

Training

To finetune an STCN model on AbdomenCT-1K with cycle consistency loss:

./run_train_stcn_with_cycle.sh

To finetune an STCN model on AbdomenCT-1K without cycle consistency loss:

./run_train_stch_without_cycle.sh

Evaluation

To evaluate a trained model on AbdomenCT-1K:

./run_eval_stcn.sh

Acknowledgement

We sincerely thank STM, STCN, MiVOS, AbdomenCT-1K for providing their wonderful code to the community!

Citation

@article{liu2023exploring,
  title={Exploring Cycle Consistency Learning in Interactive Volume Segmentation},
  author={Liu, Qin and Zheng, Meng and Planche, Benjamin and Gao, Zhongpai and Chen, Terrence and Niethammer, Marc and Wu, Ziyan},
  journal={arXiv preprint arXiv:2303.06493},
  year={2023}
}