search

srinidhiPY / SSL_CR_Histo

Official code for "Self-Supervised driven Consistency Training for Annotation Efficient Histopathology Image Analysis" Published in Medical Image Analysis (MedIA) Journal, Oct, 2021.
https://doi.org/10.1016/j.media.2021.102256
MIT License
62 stars 21 forks source link
annotation-efficient breastpathq camelyon16 deep-learning digital-pathology histopathology self-supervised-learning semi-supervised-learning teacher-student-training

readme

Self-Supervised Driven Consistency Training for Annotation Efficient Histopathology Image Analysis

by Chetan L. Srinidhi, Seung Wook Kim, Fu-Der Chen and Anne L. Martel

Overview

We propose a self-supervised driven consistency training paradigm for histopathology image analysis that learns to leverage both task-agnostic and task-specific unlabeled data based on two strategies:

  1. A self-supervised pretext task that harnesses the underlying multi-resolution contextual cues in histology whole-slide images (WSIs) to learn a powerful supervisory signal for unsupervised representation learning.

  2. A new teacher-student semi-supervised consistency paradigm that learns to effectively transfer the pretrained representations to downstream tasks based on prediction consistency with the task-specific unlabeled data.

We carry out extensive validation experiments on three histopathology benchmark datasets across two classification and one regression-based tasks:

We compare against the state-of-the-art self-supervised pretraining methods based on generative and contrastive learning techniques: Variational Autoencoder (VAE) and Momentum Contrast (MoCo), respectively.

1. Self-Supervised pretext task

2. Consistency training

Results

Pre-requisites

Core implementation:

Additional packages can be installed via: requirements.txt

Datasets

Training

The model training consists of three stages:

  1. Task-agnostic self-supervised pretext task (i.e., the proposed Resolution sequence prediction (RSP) task)
  2. Task-specific supervised fine-tuning (SSL)
  3. Task-specific teacher-student consistency training (SSL_CR)

1. Self-supervised pretext task: Resolution sequence prediction (RSP) in WSIs

From the file "pretrain_BreastPathQ.py / pretrain_Camelyon16.py", you can pretrain the network (ResNet18) for predicting the resolution sequence ordering in WSIs on BreastPathQ & Camelyon16 dataset, respectively. This can be easily adapted to any other dataset of choice.

python pretrain_BreastPathQ.py    // Pretraining on BreastPathQ   
python pretrain_Camelyon16.py    // Pretraining on Camelyon16

2. Task specific supervised fine-tuning on downstream task

From the file "eval_BreastPathQ_SSL.py / eval_Camelyon_SSL.py / eval_Kather_SSL.py", you can fine-tune the network (i.e., task-specific supervised fine-tuning) on the downstream task with limited label data (10%, 25%, 50%). Refer to, paper for more details.

python eval_BreastPathQ_SSL.py  // Supervised fine-tuning on BreastPathQ   
python eval_Camelyon_SSL.py    // Supervised fine-tuning on Camelyon16
python eval_Kather_SSL.py    // Supervised fine-tuning on Kather dataset (Colorectal)

Note: we didn't perform self-supervised pretraining on the Kather dataset (colorectal) due to the unavailability of WSI's. Instead, we performed domain adaptation by pretraining on Camelyon16 and fine-tuning on the Kather dataset. Refer to, paper for more details.

3. Task specific teacher-student consistency training on downstream task

From the file "eval_BreastPathQ_SSL_CR.py / eval_Camelyon_SSL_CR.py / eval_Kather_SSL_CR.py", you can fine-tune the student network by keeping the teacher network frozen via task-specific consistency training on the downstream task with limited label data (10%, 25%, 50%). Refer to, paper for more details.

python eval_BreastPathQ_SSL_CR.py  // Consistency training on BreastPathQ   
python eval_Camelyon_SSL_CR.py    // Consistency training on Camelyon16
python eval_Kather_SSL_CR.py    // Consistency training on Kather dataset (Colorectal)

Testing

The test performance is validated at two stages:

  1. Self-Supervised pretraining followed by supervised fine-tuning

    • From the file "eval_BreastPathQ_SSL.py / eval_Kather_SSL.py ", you can test the model by changing the flag in argument: '--mode' to 'evaluation'.

  2. Consistency training

    • From the file "eval_BreastPathQ_SSL_CR.py / eval_Kather_SSL_CR.py", you can test the model by changing the flag in argument: '--mode' to 'evaluation'.

The prediction on Camelyon16 test set can be performed using "test_Camelyon16.py" file.

License

Our code is released under MIT license.

Citation

If you find our work useful in your research or if you use parts of this code please consider citing our paper:

@article{srinidhi2022self,
  title={Self-supervised driven consistency training for annotation efficient histopathology image analysis},
  author={Srinidhi, Chetan L and Kim, Seung Wook and Chen, Fu-Der and Martel, Anne L},
  journal={Medical Image Analysis},
  volume={75},
  pages={102256},
  year={2022},
  publisher={Elsevier}
}

Extended work

Acknowledgements

Questions or Comments

Please direct any questions or comments to me; I am happy to help in any way I can. You can email me directly at srinidhipy@gmail.com.