yijingru / Vertebra-Landmark-Detection

[ISBI 2020] Vertebra-Focused Landmark Detection for Scoliosis Assessment
MIT License
93 stars 26 forks source link

Vertebra-Focused-Landmark-Detection-Pytorch

Vertebra-Focused Landmark Detection for Scoliosis Assessment [arXiv]

Accepted to ISBI2020.

Please cite the article in your publications if it helps your research:

@article{yi2020vertebra,
  title={Vertebra-Focused Landmark Detection for Scoliosis Assessment},
  author={Yi, Jingru and Wu, Pengxiang and Huang, Qiaoying and Qu, Hui and Metaxas, Dimitris N},
  booktitle={ISBI},
  year={2020}
}

Dependencies

Ubuntu 14.04, Python 3.6.4, PyTorch 1.1.0, OpenCV-Python 4.1.0.25

How to start

Prepare Dataset

To directly use dataset.py, you can arrange the dataset as follows:

/dataPath/data
    /train/*.jpg
    /val/*.jpg
    /test/*.jpg
/dataPath/labels/
    /train/*.mat
    /val/*.mat
    /test/*.mat

The source dataset is from [dataset16]. To adapt the code to your own dataset, you can modify the dataset.py, for example, change the 'load_gt_pts' function to adapt it to your own annotations. The pretrained weights can be downloaded here.

Train the model

python main.py --data_dir dataPath --epochs 50 --batch_size 2 --dataset spinal --phase train

Test the model

python main.py --resume weightPath --data_dir dataPath --dataset spinal  --phase test

Evaluate the model


python main.py --resume weightPath --data_dir dataPath --dataset spinal --phase eval