GuidedNet: Semi-Supervised Multi-Organ Segmentation via Labeled Data Guide Unlabeled Data

TL;DR

GuidedNet addresses the gap in semi-supervised multi-organ segmentation by enabling labeled-data knowledge to guide unlabeled-data training. It introduces a 3D Consistent Gaussian Mixture Model (3D-CGMM) that models class-distribution in feature space to rectify pseudo-labels, and a Knowledge Transfer Cross Pseudo Supervision (KT-CPS) strategy that re-weights unlabeled losses based on prior knowledge from labeled data. The method uses two parallel CPS-based networks and optimizes a combined objective that includes $\mathcal{L}_{sup}$, $\mathcal{L}_{kt-cps}$, and $\mathcal{L}_{cgmm}$, with $\lambda_u$ and $\lambda_g$ tuning the influence of each term. Experiments on FLARE22 and AMOS show state-of-the-art performance, particularly improving segmentation of small and complex organs under limited labeled data, demonstrating the practical impact for efficient, accurate semi-supervised medical image analysis. The work provides a concrete, plug-in framework with open-source code for researchers and clinicians aiming to leverage unlabeled clinical data more effectively.

Abstract

Semi-supervised multi-organ medical image segmentation aids physicians in improving disease diagnosis and treatment planning and reduces the time and effort required for organ annotation.Existing state-of-the-art methods train the labeled data with ground truths and train the unlabeled data with pseudo-labels. However, the two training flows are separate, which does not reflect the interrelationship between labeled and unlabeled data.To address this issue, we propose a semi-supervised multi-organ segmentation method called GuidedNet, which leverages the knowledge from labeled data to guide the training of unlabeled data. The primary goals of this study are to improve the quality of pseudo-labels for unlabeled data and to enhance the network's learning capability for both small and complex organs.A key concept is that voxel features from labeled and unlabeled data that are close to each other in the feature space are more likely to belong to the same class.On this basis, a 3D Consistent Gaussian Mixture Model (3D-CGMM) is designed to leverage the feature distributions from labeled data to rectify the generated pseudo-labels.Furthermore, we introduce a Knowledge Transfer Cross Pseudo Supervision (KT-CPS) strategy, which leverages the prior knowledge obtained from the labeled data to guide the training of the unlabeled data, thereby improving the segmentation accuracy for both small and complex organs. Extensive experiments on two public datasets, FLARE22 and AMOS, demonstrated that GuidedNet is capable of achieving state-of-the-art performance. The source code with our proposed model are available at https://github.com/kimjisoo12/GuidedNet.

Figures (5)

• Figure 1: (a) Previously developed pseudo-labeling methods separate the labeled and unlabeled data training flows, which does not reflect the interrelationship between labeled and unlabeled data (i.e., CPS chen2021semi, ARCO you2024rethinking, UCMT shen2023co). (b) In the FLARE22 dataset, rankings are provided for both the sizes of all foreground organs and the Dice. (c) Our GuidedNet, which is comprised of two components: 3D-CGMM and KT-CPS.
• Figure 2: Category-wise performance for the FLARE22 dataset. The orange bars represent the Dice for each organ, and the blue bars denote the Euclidean distances danielsson1980euclidean between the feature centers of the labeled and unlabeled data for each organ.
• Figure 3: The workflow of GuidedNet involves processing input data from $model~A$ and $model~B$ to yield predictions. The feature distributions of the labeled predictions are utilized to train the 3D-CGMM, and the generated CGMM predictions are used to rectify the initial pseudo-labels. The prior knowledge obtained from the labeled predictions are transferred to the unlabeled predictions using the KT-CPS strategy for cross pseudo supervised training.
• Figure 4: Visualization of the segmentation results on the FLARE22 dataset. (a-d)Segmentation results of one case on transverse section, coronal section, sagittal section, and 3D view, respectively. The regions enclosed by the dashed yellow boxes indicate misclassification executed by the model; our method corrects these errors within these regions.
• Figure 5: Quantitative comparisons between different hyperparameters for the FLARE22 dataset: mean Dice and mean Jaccard generated by GuidedNet when trained with various (a) ${\lambda_g}$ values and (b) layers.