Mining Gold from the Sand: Weakly Supervised Histological Tissue Segmentation with Activation Relocalization and Mutual Learning (MICCAI'24)

Abstract

Class activation maps- (CAMs-) based image-level weakly supervised tissue segmentation has became a popular research topic due to the advantage of its low annotation cost. However, there are still two challenges exist in this task: (1) low-quality pseudo masks generation, and (2) training with noisy label supervision. To address these issues, we propose a novel weakly supervised segmentation framework with Activation Relocalization and Mutual Learning (ARML). First, we integrate an Activation Relocalization Scheme (ARS) into classification phase to more accurately cover the useful areas in initial CAMs. Second, to deal with the inevitably noisy annotations in pseudo masks generated by ARS, we propose a noise-robust mutual learning segmentation model. The model promotes peer networks to capture different characteristics of the outputs, and two noise suppression strategies namely samples weighted voting (SWV) and samples relation mining (SRM) are introduced to excavate the potential credible information from noisy annotations. Extensive experiments on BCSS and LUAD-HistoSeg datasets demonstrate that our proposed ARML exceeds many state-of-the-art weakly supervised semantic segmentation methods, which gives a new insight for tissue segmentation tasks.

Datasets

• BCSS and LUAD-HistoSeg datasets can be found in WSSS-Tissue.
• Download the datasets and unpack them into datasets folder, and make sure the formats are same with WSSS-Tissue.

Pretrained weights and checkpoints

• Download the pretained weight of classification stage from this link, and put it into init_weights folder.
• Download the checkpoint of classification stage from this link.
• Download the checkpoint of segmentation stage from this link.

Run the whole pipeline

1. Train the classification model

python 1_train_stage1.py --dataset bcss --trainroot datasets/BCSS-WSSS/train/ --testroot dataset/BCSS-WSSS/test/ --max_epoches 20

2. Generate pseudo masks

python 2_generate_PM.py --dataroot datasets/BCSS-WSSS --dataset bcss

3. Train the segmentation model

python 3_train_stage2.py --dataset bcss --dataroot datasets/BCSS-WSSS --epochs 10

Acknowlegment

Many thanks to WSSS-Tissue.

Contact

If you have any question, please contact [email protected].