VesselSeg-Pytorch : Retinal vessel segmentation toolkit based on pytorch
Introduction
This project is a retinal blood vessel segmentation code based on python and pytorch framework, including data preprocessing, model training and testing, visualization, etc. This project is suitable for researchers who study retinal vessel segmentation.
Requirements
The main package and version of the python environment are as follows
# Name Version
python 3.7.9
pytorch 1.7.0
torchvision 0.8.0
cudatoolkit 10.2.89
cudnn 7.6.5
matplotlib 3.3.2
numpy 1.19.2
opencv 3.4.2
pandas 1.1.3
pillow 8.0.1
scikit-learn 0.23.2
scipy 1.5.2
tensorboardX 2.1
tqdm 4.54.1
The above environment is successful when running the code of the project. In addition, it is well known that pytorch has very good compatibility (version>=1.0). Thus, I suggest you try to use the existing pytorch environment firstly.
The current version has problems reading the .tif
format image in the DRIVE dataset on Windows OS. It is recommended that you use Linux for training and testing
Usage
0) Download Project
Runninggit clone https://github.com/lee-zq/VesselSeg-Pytorch.git
The project structure and intention are as follows :
VesselSeg-Pytorch # Source code
βββ config.py # Configuration information
βββ lib # Function library
β βββ common.py
β βββ dataset.py # Dataset class to load training data
β βββ datasetV2.py # Dataset class to load training data with lower memory
β βββ extract_patches.py # Extract training and test samples
β βββ help_functions.py #
β βββ __init__.py
β βββ logger.py # To create log
β βββ losses
β βββ metrics.py # Evaluation metrics
β βββ pre_processing.py # Data preprocessing
βββ models # All models are created in this folder
β βββ denseunet.py
β βββ __init__.py
β βββ LadderNet.py
β βββ nn
β βββ UNetFamily.py
βββ prepare_dataset # Prepare the dataset (organize the image path of the dataset)
β βββ chasedb1.py
β βββ data_path_list # image path of dataset
β βββ drive.py
β βββ stare.py
βββ tools # some tools
β βββ ablation_plot.py
β βββ ablation_plot_with_detail.py
β βββ merge_k-flod_plot.py
β βββ visualization
βββ function.py # Creating dataloader, training and validation functions
βββ test.py # Test file
βββ train.py # Train file
1) Datasets preparation
- Please download the retina image datasets(DRIVE, STARE and CHASE_DB1) from TianYi Cloud. Otherwise, you can download three data sets from the official address: DRIVE,STARE and CHASE_DB1.
- Unzip the downloaded
datasets.rar
file. The results are as follows:
datasets
βββ CHASEDB1
β βββ 1st_label
β βββ 2nd_label
β βββ images
β βββ mask
βββ DRIVE
β βββ test
β βββ training
βββ STARE
βββ 1st_labels_ah
βββ images
βββ mask
βββ snd_label_vk
- Create data path index file(.txt). running:
Please modify the data folder path:
data_root_path
(in thedrive.py
,stare.py
andchasedb1.py
) to the absolute path of the datasets downloaded above
python ./prepare_dataset/drive.py
In the same way, the data path files of the three datasets can be obtained, and the results are saved in the ./prepare_dataset/data_path_list
folder
2) Training model
Please confirm the configuration information in the config.py
. Pay special attention to the train_data_path_list
and test_data_path_list
. Then, running:
CUDA_VISIBLE_DEVICES=1 python train.py --save UNet_vessel_seg --batch_size 64
You can configure the training information in config, or modify the configuration parameters using the command line. The training results will be saved to the corresponding directory(save name) in the experiments
folder.
3) Testing model
The test process also needs to specify parameters in config.py
. You can also modify the parameters through the command line, running:
CUDA_VISIBLE_DEVICES=1 python test.py --save UNet_vessel_seg
The above command loads the best_model.pth
in ./experiments/UNet_vessel_seg
and performs a performance test on the testset, and its test results are saved in the same folder.
Visualization
- Training sample visualization
- Segmentation results
The original image, predicted probability image, predicted binary image and groundtruth
DRIVE:
STARE:
CHASE_DB1:
To DO
- Add other retinal vessel segmentation models and performances.
- Add SOTA loss function.