StenUNet: Automatic Stenosis Detection from X-ray Coronary Angiography

Introduction

This algorithm is for the stenosis detection task in ARCADE Challenge, which was held at MICCAI 2023. We are ranked ${\textsf{\color{red}3rd}}$ !

Our publication: StenUNet: Automatic Stenosis Detection from X-ray Coronary Angiography Arxiv

Please refer to MICCAI-ARCADE for the segmentation detection task.

Installation

python>=3.9 and torch>=2.0.0

  conda create -n stenunet_env python=3.9
  conda activate stenunet_env
  git clone https://github.com/HuiLin0220/StenUNet.git
  cd StenUNet
  pip install  -r ./requirements.txt

Prepare data

• The training data folder structure is like this:

     Raw_data/Dataset_train_val/  
      ├── imagesTr
      │   ├── sten_0000_0000.png
      │   ├── sten_0000_0001.png
      │   ├── ...
      │   ├── sten_0001_0000.png      
      │   ├── sten_0001_0001.png      
      │   ├── ... 
      │   ├── sten_0002_0000.png
      │   ├── sten_0002_0001.png
      │   ├── ...
      ├── labelsTr
      │   ├── sten_0000.png
      │   ├── sten_0001.png
      │   ├── sten_0002.png
      │   ├── ...
      ├── dataset.json
  
  (1) sten_0000_0000.png and sten_0000_0001.png are considered two different modalities for the same raw image (sten_0000).
  (2) We provide some preprocessing methods in [preprocess.py](pre_process/preprocess.py) You can do some preprocessing on the raw image and get several modalities for training.
  (3) Note that inference and training should use the same preprocessing strategies.
  

• Rename and put the training images in this folder "./nnNet_training/Raw_data/"

• Edit dataset.json ("numTraining" indicates the number of training samples in your dataset.)

Train

• Planning hyper_parameters

  python training_planning.py 
  

• Training from scratch

  CUDA_VISIBLE_DEVICES=0 python training.py 0
  #CUDA_VISIBLE_DEVICES=X python train.py fold_ID (can be 0,1,2,3,4)
  

• Finetune the pre-trained model on your own data

  CUDA_VISIBLE_DEVICES=0 python training.py 1 -pretrained_weights MODEL_WEIGHTS_PATH
  

  if you want to use Shared weights, you need to replace your nnUNetPlans.json with my nnUNetPlans.json and use the "foreground_intensity_properties_per_channel" in your dataset_fingerprint.json

Inference

1. Rename and put the test images in this folder'./dataset_test/raw';

2. Run

     python inference.py -chk MODEL_WEIGHTS_PATH
   

3.Sharing StenUnet's weight (Google drive).
4. You will get the preprocessed images, raw prediction after StenUNet, and post_prediction after postprocessing.

You can integrate your own preprocessing/postprocessing strategies in preprocess.py/post_process

The inference folder structure is like this:

  daset_test/
      ├── raw
      │   ├── sten_0000_0000.png
      │   ├── sten_0001_0000.png
      │   ├── ...
      ├── preprocessed
      │   ├── sten_0000_0000.png       # prerpocessing method0
      │   ├── sten_0000_0001.png       # prerpocessing method1
      │   ├── sten_0000_0003.png       # prerpocessing method2
      │   ├── ... 
      │   ├── sten_0001_0000.png
      │   ├── sten_0001_0001.png
      │   ├── sten_0001_0003.png
      │   ├── ...
      ├── raw_prediction
      │   ├── sten_0000.png
      │   ├── sten_0001.png
      │   ├── ...
      ├── post_prediction
      │   ├── sten_0000.png
      │   ├── sten_0001.png
      │   ├── ...

References

nnunet

Citation

Please cite the following paper when using SteUNet:

  @article{lin2023stenunet,
    title={StenUNet: Automatic Stenosis Detection from X-ray Coronary Angiography},
    author={Lin, Hui and Liu, Tom and Katsaggelos, Aggelos and Kline, Adrienne},
    journal={arXiv preprint arXiv:2310.14961},
    year={2023}
  }

Contact Us

Feel free to contact me at huilin2023@u.northwestern.edu

To-do List