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 our MICCAI 2022 paper iSegFormer in 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}
}

isegformer's Issues

训练中断

你好你好，我这还有个问题

训练如果中断了，我想在原有的基础上继续训练

我看到你有两个参数
--resume-exp 和 --resume-prefix
这两个该如何使用呢，比如我想跑100个epoch ,在80断了，保存了last.pth ,这两个参数如何设置

Model weights finetuned by COCO_Lvis dataset

It seems that in the readme the model weight loaded for evaluation is imagenet21k_pretrain_cocolvis_finetune_swin_base_epoch_54.pth, however, the weights provided in the https://drive.google.com/drive/folders/1KG6QVwlydyEzmcKNHgvCZjKhsaycIFiV is fine-tuned by OAI-ZIB, could you provide the weight fine-tuned by COCO_Lvis?
Besides, I test the model imagenet21k_pretrain_oaizib_finetune_swin_base_epoch_54 using ssTEM, and all samples fail (IOU is less than 0.1 with 20 clicks), but it works well for OAI-ZIB, is this phenomenon normal?

Hi, does this support test now?

I try to run python demo_iseg2d.py --checkpoint=stcn_no_finetune --cpu, with below error happen
File "D:\projects\iSegFormer\demo_iseg2d.py", line 61, in
main()
File "D:\projects\iSegFormer\demo_iseg2d.py", line 18, in main
model = utils.load_is_model(checkpoint_path, args.device, args.eval_ritm, cpu_dist_maps=True)
File "D:\projects\iSegFormer\iseg2d\SimpleClick\isegm\inference\utils.py", line 34, in load_is_model
return load_single_is_model(state_dict, device, eval_ritm, **kwargs)
File "D:\projects\iSegFormer\iseg2d\SimpleClick\isegm\inference\utils.py", line 38, in load_single_is_model
model = load_model(state_dict['config'], eval_ritm, **kwargs)
KeyError: 'config'

Can't find the dataset and weight

try to run the following piece

# Download OAI-ZIB test set
URL_ISEGFORMER = "https://github.com/qinliuliuqin/iSegFormer/releases/download/v0.1"
DATA_FOLDER = "./datasets" 

!mkdir -p {DATA_FOLDER}

for dataset in ['OAI-ZIB-test']:
  dataset_url = f"{URL_ISEGFORMER}/{dataset}.zip"
  dataset_path = f"{DATA_FOLDER}/{dataset}.zip"
  !wget -q -O {dataset_path} {dataset_url}
  !unzip -q {dataset_path} -d {DATA_FOLDER}
  !rm {dataset_path}

# Download weights 
WEIGHTS_FOLDER = "./weights"
!mkdir -p {WEIGHTS_FOLDER}

MODEL_NAME_SWINB = "imagenet21k_pretrain_oaizib_finetune_swin_base_epoch_54"
MODEL_NAME_HR32 = "oai_pretrain_oaizib_finetune_hr32_epoch_54"

WEIGHTS_URL_SWINB = f"{URL_ISEGFORMER}/{MODEL_NAME_SWINB}.pth"
WEIGHTS_URL_HR32 = f"{URL_ISEGFORMER}/{MODEL_NAME_HR32}.pth"
!wget -q -P {WEIGHTS_FOLDER} {WEIGHTS_URL_SWINB}
!wget -q -P {WEIGHTS_FOLDER} {WEIGHTS_URL_HR32}

https://github.com/uncbiag/iSegFormer/blob/5dbda2446a96e61abe4ed1734b9d7d3c05899489/notebooks/colab_test_isegformer.ipynb

代码入口

您好，我一直不知道您这个代码入口在哪，我想用您的网络在我自己的医学数据集上跑一下，但是不知道该从哪入手，所以想请教一下。方便的话可以加个微信吗

NoC

你好，还有个问题需要请教一下，NoC的计算方式是按照RITM的来的吗，具体的代码是在哪里啊。感谢