Tensorflow based implementation of BEDs as described in BEDs: Bagging Ensemble Deep Segmentation for Nucleus Segmentation with Testing Stage Stain Augmentation by Xing Li, Haichun Yang, Jiaxin He, Aadarsh Jha, Agnes B. Fogo, Lee E. Wheless, Shilin Zhao and Yuankai Huo.
The proposed method is shown in figure below:
Training | Testing |
---|---|
![]() |
![]() |
The performance overview for the proposed method is shown in figure below:
Add options to skip stain augmentation and do 5 model fusion, please pull both BEDs_e2e_infer.py
and eval/utils.py
to enable this update. Turn -BEDs_5
on to run 5 models inference. Turn -auto
on to use 5 models selected randomly or to use the best 5 models we selected by default.
python BEDs_e2e_infer.py --model_dir models/deep_forest/ --target_dir SPECIFY/IF/DO/STAIN/AUG/ --annot_dir SPECIFY/IF/EVAL/ --output_dir experiments/BEDs_e2e_infer/BEDs_5_Model/ --ext tif -BEDs_5 datasets/Test/Images/
Add an End-to-end inference script to test your custom dataset(image) using BEDs All 33 with our pre-trained models and pre-defined stain targets(./stain_template
). Our default input size is 1000x1000.
python BEDs_e2e_infer.py --model_dir models/deep_forest/ --target_dir stain_template/ --annot_dir SPECIFY/IF/EVAL/ --output_dir experiments/BEDs_e2e_infer/ --ext tif IM_OR_FOLDER
CUDA9-cuDNN7
python>=3.5.2
tensorflow-gpu>=1.12.0
opencv-python>=3.4.0.14
pytorch>=1.4.0+cu92
torchvision>=0.5.0+cu92
spams==2.6.1
staintools==2.1.2
To setup the python environment for this project, run:
pip install -r requirements.txt
Follow the procedure below to setup and download necessary data to reproduce the results in the paper.
- Clone this repository.
- Download the tools for image feature vector extraction with Resnet18 pre-trained model here. This useful tool was done by Christian Safka. Put this repo under
data_proecss
directory to enablegenerate_stainTemplate.py
for stain template selection among training data. - Create directory for datasets.
mkdir datasets
cd datasets
mkdir Train Val Test
- Download training dataset from manually labeled dataset under the tab
Datailed Description
: Dataset of segmented nuclei in hematoxylin and eosin stained histopathology images of ten cancer types. Extract todatasets/Train
. Put*_crop.png
todatasets/Train/Images
and*_labeled_mask_corrected.png
todatasets/Train/Annotations
. - Download validation dataset from MoNuSeg Challenge 2018 (Training). Extract to
datasets/Val
. Rename the directoryTissue Images
toImages
- Download testing dataset from MoNuSeg Challenge 2018 (Testing). Extract to
datasets/Test
. Put.tif
files todatasets/Test/Images
and.xml
files todatasets/Test/Annotations
. - (Optional) Download our stain augmented testing images here. Extract to
datasets/Test
- (Optional) Download our pre-trained models here. Extract to
BEDs
.
After following the steps in Setup, run following script to:
cd BEDs/
- Generate benchmark and random split training data for U-net training:
python data_process/prepare_datasets.py --annot-dir datasets/Train/Annotations/ --output-dir datasets/Train/deep_forest/ --stage train --subset-num 33 datasets/Train/Images/
- Generate validation data:
python data_process/prepare_datasets.py --annot-dir datasets/Val/Annotations/ --output-dir datasets/Val/Val/ --stage val datasets/Val/Images/
- Cluster training dataset based on extracted feature vectorwith Resnet18 pre-trained model and select cluster center as stain template targets:
python generate_stainTemplate.py --input_dir ../datasets/Train/Images/ --output_dir ../datasets/Train/Stain_template/
generate_stainTemplate.py
used k-means to cluster training images into 12 classes (17, 20, 185, 255, 350, 412, 497, 926, 1041, 1094, 1210 _crop.png
). The templates are further evaluated on validation dataset to select the best 6 templates as the final stain augmentation targets. (This evaluation on validation dataset takes some time for inference) The final stain templates are 255, 350, 412, 926, 1041, 1043 _crop.png
.
- You can skip the previous step and use the template provided in
stain_template
to continue stain augmentation for testing images toward stain template:
python stain_norm_dir.py --input_dir ../datasets/Test/Images/ --target_dir stain_template/ --output_dir ../datasets/Test/Test_pairs_final/ --ext tif
- Generate testing data for each type of stain augmentation (0 is the original stain, 1-6 is the augmented stain):
python data_process/prepare_datasets.py --annot-dir datasets/Test/Annotation/ --output-dir datasets/Test/Test_pairs/0/ --stage test datasets/Test/Images_stainNormed/0/
Or can simply run:
. generate_test_stain_pair.sh
to prepare all testing images.
If Step 6 in Setup and Step 3 in Data Processing was successfully done, run the following command to do inference for BEDs:
cd eval
python BEDs_inference.py --model-dir ../models/deep_forest/ --output-dir ../experiments/BEDs_inference_results/ ../datasets/Test/Test_pairs/
The system would take a while to inference 7 difference stain augmentation with 33 models.
The ablation study, in the figure below, showed that the self-ensemble learning and testing stage stain augmentation were mutually complementary. Herein, the holistic model achieved the highest mean and median DSC, without using any extra training data:
Ablation Study | Distribution Boxplot |
---|---|
![]() |
![]() |
After obtained the BEDs inference results in experiments/BEDs_inference_results/
, following steps are followed to evaluate the performance reported in paper.
- Benchmark:
python benchmark_eval.py --nuclei-model ../models/benchmark/frozen_model.pb --output-dir ../experiments/benchmark/ ../datasets/Test/Test_pairs/0/
- Model 1, 2, 3 (Single Model):
python benchmark_eval.py --nuclei-model ../models/deep_forest/RANDOM_MODEL/frozen_model.pb --output-dir ../experiments/RANDOM_MODEL/ ../datasets/Test/Test_pairs/0/
- BEDs 5:
python BEDs_eval.py --model-dir ../models/deep_forest/ --output-dir ../experiments/BEDs_inference_results/ --fusion-dir ../experiments/fusing_results/ --experiments BEDs_Model --model_num 5 ../datasets/Test/Test_pairs/0/
- BEDs 33:
python BEDs_eval.py --model-dir ../models/deep_forest/ --output-dir ../experiments/BEDs_inference_results/ --fusion-dir ../experiments/fusing_results/ --experiments BEDs_Model --model_num 33 ../datasets/Test/Test_pairs/0/
- BEDs 33 Model-Stain:
python BEDs_eval.py --model-dir ../models/deep_forest/ --output-dir ../experiments/BEDs_inference_results/ --fusion-dir ../experiments/fusing_results/ --experiments BEDs_Model-Stain --model_num 33 ../datasets/Test/Test_pairs/0/
- BEDs 33 Stain-Model:
python BEDs_eval.py --model-dir ../models/deep_forest/ --output-dir ../experiments/BEDs_inference_results/ --fusion-dir ../experiments/fusing_results/ --experiments BEDs_Stain-Model --model_num 33 ../datasets/Test/Test_pairs/0/
- BEDs 33 All:
python BEDs_eval.py --model-dir ../models/deep_forest/ --output-dir ../experiments/BEDs_inference_results/ --fusion-dir ../experiments/fusing_results/ --experiments BEDs_All --model_num 33 ../datasets/Test/Test_pairs/0/
To compute the objectwise F1 for each experiment, run:
python objectwise_DSC_eval.py --ref-dir ../datasets/Test/Test_GT/ --input-dir ../experiments/fusing_results/EXPERIMENT_DIR/ --output-dir ../experiments/objectwise_F1/EXPERIMENT_DIR/
By following the procedure above, we obtained the following experiment results:
The U-net implementation used in this project is modified based on the encoder-decoder network with skip connection in pix2pix. A validation during training option is added to the original work. If a validation path is defined, the training script will save the best model based on evaluation on the validation dataset. Otherwise, only the lastest model will be saved. If Step 1 & 2 in Data Processing was successful, the model can be trained using train_unet.py
.
For example, to train the benchmark:
python train_unet.py --input_dir datasets/Train/deep_forest/benchmark/train/ --val_dir datasets/Val/Val/ --mode train --output_dir checkpoints/unet_ckpts/benchmark/ --max_epochs 30 --summary_freq 1356 --save_freq 1356 --display_freq 5424 --scale_size 256
To train the model with sub-dataset:
python train_unet.py --input_dir datasets/Train/deep_forest/RANDOM_MODEL/train/ --val_dir datasets/Val/Val/ --mode train --output_dir checkpoints/unet_ckpts/RANDOM_MODEL/ --max_epochs 30 --summary_freq 904 --save_freq 904 --display_freq 4520 --scale_size 256
Or can simply run:
. deep_forest_train.sh
During training, the process can be monitored with tensorboard
:
tensorboard --logdir checkpoints/unet_ckpts/deep_forest/RANDOM_MODEL/
After all the training are done, one can export and freeze the model with the tools provided in this package:
python tools/reduce_model.py --model-input checkpoints/unet_ckpts/deep_forest/MODEL_DIR/ --model-output checkpoints/unet_ckpts/deep_forest/MODEL_DIR/
python tools/freeze_model.py --model-folder checkpoints/unet_ckpts/deep_forest/MODEL_DIR/
Or can simply run:
. freeze_model.sh