Edge-Aware Vessel Segmentation using Scribble Supervision

This repository contains the official implementation of our paper:

Edge-Aware Vessel Segmentation using Scribble Supervision, IEEE International Symposium on Biomedical Imaging 2024

Zhanqiang Guo, Jianjiang Feng, Jie Zhou

Paper: tsinghua.edu.cn

Abstract: Accurate segmentation of blood vessels is critical for diagnosing various diseases. However, the complexity of manually labeling vessels impedes the practical adoption of fully supervised methods. To alleviate this challenge, we propose a weakly supervised vessel segmentation framework. Our approach leverages scribble annotation to train the Unet and identifies reliable foreground and background regions. Addressing the issue of insufficient boundary information inherent in scribble annotation, we incorporate a conventional approach specifically designed to leverage the innate structural attributes of vessels for edge detection, subsequently ensuring effective edge supervision. In addition, a bilateral filtering module is introduced to improve edge awareness of network. Furthermore, to augment the quantity of annotated pixels, we employ an image mixing strategy for data augmentation, thereby enhancing the network's segmentation capability. The experimental results on three datasets show that our framework outperforms the existing scribble-based methods.

Illustration of The Proposed Framework.

News

• [Feb 25 2024] 🔔 Code is coming.

Requirements

• Python = 3.8
• Pytorch = 1.10.0
• CUDA = 11.1
• Install other packages in requirements.txt

Data preparation

The file structure is as follows:

root_path/dataset/
├── size_img
├── size_label
├── train
├────── images
├───────── img1.png
├────── labels
├───────── img1_scribble.png
├────── annotations
├───────── img1_label.png
├── val
├────── images
├───────── img2.png
├────── labels
├───────── img2_label.png
├── TestSet
├────── images
├───────── img3.png
├────── labels
└───────── img3_label.png

Run

Stage1

• Training.

  python train_infer/stage1_main.py --mode 'train' --dataset 'SSVS_XRAY_Coronary'

• Pseudo label.

  python train_infer/stage1_main.py --mode 'pseudo' --dataset 'SSVS_XRAY_Coronary'

• Pseudo label Refinement.

  python train_infer/pseudo_label_refine.py --dataset 'SSVS_XRAY_Coronary'

Stage2

• Training.

  python train_infer/stage2_main.py --mode 'train' --dataset 'SSVS_XRAY_Coronary'

• Inference.

  python train_infer/stage2_main.py --mode 'infer' --dataset 'SSVS_XRAY_Coronary'
  

Citation

We hope you find our work useful. If you would like to acknowledge it in your project, please use the following citation:

@inproceedings{guo2024edge,
  title={Edge-Aware Vessel Segmentation using Scribble Supervision},
  author={Guo, Zhanqiang and Feng, Jianjiang and Zhou, Jie},
  booktitle={2024 IEEE 21th International Symposium on Biomedical Imaging (ISBI)},
  pages={1--5},
  year={2024},
  organization={IEEE}
}

Contact me

If you have any questions about this code, please do not hesitate to contact me.

Zhanqiang Guo: [email protected]