Semi-supervised semantic segmentation of cell nuclei with diffusion model and collaborative learning

Automated segmentation and classification of the cell nuclei in microscopic images is crucial for disease diagnosis and tissue microenvironment analysis. Given the difficulties in acquiring large labeled datasets for supervised learning, semi-supervised methods offer alternatives by utilizing unlabeled data alongside labeled data. Effective semi-supervised methods to address the challenges of extremely limited labeled data or diverse datasets with varying numbers and types of annotations remain under-explored.

Unlike other semi-supervised learning methods that iteratively use labeled and unlabeled data for model training, we introduce a semi-supervised learning framework that combines a latent diffusion model (LDM) with a transformer-based decoder, allowing for independent usage of unlabeled data to optimize their contribution to model training. The model is trained based on a sequential training strategy. LDM is trained in an unsupervised manner on diverse datasets, independent of cell nuclei types, thereby expanding the training data and enhancing training performance. The pre-trained LDM serves as a powerful feature extractor to support the transformer-based decoder's supervised training on limited labeled data and improve final segmentation performance. In addition, the paper explores a collaborative learning strategy to enhance segmentation performance on out-of-distribution (OOD) data.

Extensive experiments conducted on four diverse datasets demonstrated that the proposed framework significantly outperformed other semi-supervised and supervised methods for both in-distribution and OOD cases. Through collaborative learning with supervised methods, diffusion model and transformer decoder-based segmentation (DTSeg) achieved consistent performance across varying cell types and different amounts of labeled data.

The proposed DTSeg framework addresses cell nuclei segmentation under limited labeled data by integrating unsupervised LDM training on diverse unlabeled datasets. Collaborative learning demonstrated effectiveness in enhancing the generalization capability of DTSeg to achieve superior results across diverse datasets and cases. Furthermore, the method supports multi-channel inputs and demonstrates strong generalization to both in-distribution and OOD scenarios.