This repo is the official implementation of 'Lesion-aware Knowledge Distillation for Diabetic Retinopathy Lesion Segmentation'.
The manuscript is now under review.
Install from the requirements.txt using:
pip install -r requirements.txt
IDRiD and DDR datasets can be downloaded in the following links:
Then prepare the datasets in the following format for easy use of the code:
├── data
│ ├── idrid
│ ├── ddr
The DDR dataset is only required to the segmentation part without including the grading part.
The first step is to train the teacher network. You can select your desired type of teacher network in train_single_network.py.
Run:
python train_single_network.py
Also different segmentation networks can be trained by this script, they are 'MCA-UNet', 'U²Net', 'UNet', 'UNet++', 'UCtransnet', 'UDtransnet', 'Attention-UNet', 'Res-UNet++', 'UNet3+', 'TransUNet' and 'ENet'.
Once the teacher network is well-trained, it can be used to guide the student network. Just set teacher_checkpoints to the path of your teacher network weights in train_student_lkd.py and then Run:
python train_student_lkd.py
Run:
python test.py
You can get the Dice、IoU and AUPR scores.
Yaqi Wang (wangyaqicv@gmail.com)