SGTC: Semantic-Guided Triplet Co-training for Sparsely Annotated Semi-Supervised Medical Image Segmentation

In the following, we first define some related preliminaries of this work. Specifically, given a training dataset $D$, we have a labeled set containing $M$ labeled cases, represented as $D_{L}=\{(X_{i}^{l},Y_{i})\}_{i=1}^{M}$, where $X_{i}^{l}$ denotes the input images and $y_{i}$ denotes the corresponding ground truth. Additionally, the unlabeled set $D_{U}$ containing $N$ unlabeled cases, represented as $D_{U}=\{X_{i}^{u}\}_{i=1}^{N}$, where $N\gg M$. Our method annotates three slices from three orthogonal planes, including $X_{is}^{lp}$ ($s$ means sagittal), $X_{ic}^{lq}$($c$ means coronal), and $X_{ia}^{lr}$($a$ means axial). Overall, there are $p^{th}$ slice in plane $s$, $q^{th}$ slice in plane $c$, and $r^{th}$ slice in plane $a$. Their corresponding annotations are $Y_{is}^{p}$, $Y_{ic}^{q}$, and $Y_{ia}^{r}$, respectively.

Figure 2 illustrates our proposed Semantic-Guided Triplet Co-Training (SGTC) framework, which consists of two main components: semantic-guided auxiliary learning mechanism and triple-view disparity training strategy. Benefiting from these, our SGTC introduces text representations to enhance semi-supervised learning to exploit more discriminative semantic information. Then, we will elaborate on the technical details of each component step by step.

Semantic-Guided Auxiliary Learning Mechanism (SGAL): While some methods have incorporated textual semantics as auxiliary information to guide segmentation, these methods are still fully supervised, requiring a large amount of labeled data, thus limiting their applicability in clinical settings (Li et al. 2023a; Liu et al. 2023; Shin et al. 2022). To fill this blank, we propose a novel semantic-guided auxiliary learning mechanism, utilizing the text representations from the pre-trained CLIP to enable semantic-aware and fine-granular semi-supervised medical image segmentation, while further enhancing the quality of pseudo labels.

Specifically, the designed medical prompts are processed through the pre-trained text encoder of CLIP to obtain the text embedding $w$. Since the CLIP is pre-trained on natural images (i.e., the domain gap between natural images and medical images (Ye et al. 2022)), it may prevent the CLIP text encoder from fully capturing the clinical semantics contained in medical prompts. Therefore, we freeze the pre-trained CLIP and fine-tune the Adapter module followed by the frozen CLIP encoder. The Adapter module consists of a dimension reduction projection layer followed by an activation function layer, and an up projection layer. The calculation process of $w$ can be written as:

$$w=\text{Adapter}(\text{CLIP}_{Enc}(\text{Text Prompt})).$$

(1)

In this paper, the text prompts are defined as “An image containing the [CLS], with the rest being background”, where [CLS] is a concrete class name. It is well known that the template of medical prompts is crucial, and therefore the effectiveness of different prompt templates will be verified in the following studies. After obtaining the text embedding $w$, we concatenate $w$ with the global image feature $f_{img}$ extracted through the encoder path of the segmentation network to better align the image-text modality. Subsequently, this concatenated representation is directly fed into a multi-layer perception (MLP) to obtain the cross-modal parameters $\gamma_{s/c/a}^{1,2}$, which can be formulated as follows:

$$\gamma_{s/c/a}^{1,2}=\text{MLP}(w~{}\copyright~{}f_{img}^{s/c/a}),$$

(2)

where $\copyright$ represents concatenation. Then, we perform element-wise addition of the cross-modal parameters $\gamma$ and the features extracted before the final classification layer $\Theta_{s/c/a}$ of each sub-networks, and pass it through a convolution layer to obtain the predictions $P_{s/c/a}$ of each branch.

$$P_{s/c/a}=\text{Conv3D}(\Theta_{s/c/a}\oplus\gamma_{s/c/a}^{1,2}),$$

(3)

where $\oplus$ represents element-wise addition, and Conv3D denotes the 3D convolution layer.

Triple-View Disparity Training Strategy (TVDT): The previous sparse annotation approach, which annotated slices on only one or two planes, failed to fully preserve spatial information, leading to a suboptimal segmentation results. To address this issue, we proposed a novel triple-view disparity training strategy, which significantly improves the robustness by maintaining the disparity of sub-networks during training as well as allowing the sub-networks to learn complementary knowledge from each other. More importantly, it just needs the clinician to annotate three orthogonal slices of a few volumetric samples.

Specifically, for a volume $X^{l}$, we first employ three distinct sparse labels $Y_{s}$, $Y_{c}$, and $Y_{a}$ to supervise the three sub-networks respectively. This ensures better consistency training while maintaining the disagreement among different sub-networks. For each sub-network, we select two orthogonal annotated slices from the three orthogonal planes as supervision signals. For the sparse label $Y_{s}$, we choose $Y_{c}^{q}$ slice and $Y_{a}^{r}$ slice, which can be formulated as follows:

$$Y_{s}=Y\otimes W_{s},$$

(4)

$$W_{s}^{i}=\begin{cases}1,&\quad\text{if voxel $i$ is on the selected slices,}\\ 0,&\quad\text{otherwise},\end{cases}$$

(5)

where $\otimes$ denotes element-wise multiplication. $W_{s}$ means the weight matrix. Similarly, for $Y_{c}$, we choose $Y_{s}^{p}$ and $Y_{a}^{r}$ slice, and for $Y_{a}$, $Y_{s}^{p}$ and $Y_{c}^{q}$ slice are allocate.

Our SGTC comprises three 3D segmentation networks (i.e., Vnet (Milletari, Navab, and Ahmadi 2016)), denoted as $F_{s}(\cdot)$, $F_{c}(\cdot)$, and $F_{a}(\cdot)$. As mentioned above, every volume $X_{i}^{l}$ ($i\leq M$) merely has three orthogonal labels $Y_{s}$ ,$Y_{c}$ and $Y_{a}$. Sub-network $F_{s}(\cdot)$ is trained with $Y_{s}$, $F_{c}(\cdot)$ is trained with $Y_{c}$, and $F_{a}(\cdot)$ is trained with $Y_{a}$, respectively. Under this setting, all three sub-networks can learn knowledge from different planes of the others while maintaining key spatial information of 3D medical volumes.

For abundant unlabeled data, the triplet sub-networks guide each other in turn, where the predictions generated by model $F_{s}(\cdot)$ serve as pseudo labels for $F_{c}(\cdot)$ and $F_{a}(\cdot)$. Here, following (Yu et al. 2019), we compute uncertainty using Monte Carlo dropout, selecting voxels with uncertainties lower than a threshold as better pseudo labels for triplet co-training. Similarly, the predicted results of $F_{c}(\cdot)$ and $F_{a}(\cdot)$ serve as pseudo labels for the other two sub-networks. The loss is computed using weighted cross-entropy loss as:

$$\mathcal{L}_{TVDT}^{k}=-\frac{1}{\sum_{i=1}^{H\times W\times D}m_{i}}\sum_{i=1}^{H\times W\times D}m_{i}\hat{y}_{i}^{k}\log p_{i},k=1,2,$$

(6)

where $m_{i}$ denotes whether the $i^{th}$ voxel is selected. $p_{i}$ is the prediction result of the current model, and $\hat{y}_{i}^{1}$ and $\hat{y}_{i}^{2}$ are the pseudo labels generated by the other two sub-networks.

In addition, for limited labeled data, the supervised loss contains weighted cross-entropy loss and weighted dice loss, which can be defined as follows:

$$\mathcal{L}_{WCE}=-\frac{1}{\sum_{i=1}^{H\times W\times D}w_{i}}\sum_{i=1}^{H\times W\times D}w_{i}y_{i}\log p_{i},$$

(7)

$$\mathcal{L}_{Dice}=1-\frac{2\times\sum_{i=1}^{H\times W\times D}w_{i}p_{i}y_{i}}{\sum_{i=1}^{H\times W\times D}w_{i}(p_{i}^{2}+y_{i}^{2})},$$

(8)

where $w_{i}$ is the $i^{th}$ voxel value of the weight matrix. $p_{i}$ and $y_{i}$ respectively represent the predicted results and the ground truth labels of the network on voxel $i$. Therefore, the supervised learning loss is represented as follows:

$$\mathcal{L}_{Sup}=\frac{1}{2}\mathcal{L}_{WCE}+\frac{1}{2}\mathcal{L}_{Dice}.$$

(9)

Overall, the total loss during training is defined as:

$$\mathcal{L}_{SGTC}=(1-\alpha)\mathcal{L}_{Sup}+\alpha\mathcal{L}_{TVDT}^{k},$$

(10)

where the first term $\mathcal{L}_{Sup}$ is tailored for labeled data, and the second term $\mathcal{L}_{TVDT}^{k}$ is employed for unlabeled data. $\alpha$ denotes the proposed dynamic parameter, which enables the supervised loss from the three annotated slices to dominate the training at the first few epochs, and then gradually increase the weight of the unsupervised loss to stabilize the entire semi-supervised learning.

LA2018 Dataset (Xiong et al. 2021) contains 100 gadolinium enhanced MR imaging scans with labels. All scans have the same isotropic resolution of $0.625\times 0.625\times 0.625mm^{3}$. Following existing works (Yu et al. 2019; Luo et al. 2021), we utilize 80 training samples and 20 testing samples for fair comparison with other methods.

KiTS19 Dataset (Heller et al. 2019) is provided by the Medical Centre of Minnesota University and consists of 300 abdominal CT scans. The slice thickness ranges from 1mm to 5mm. The dataset is divided into 190 training samples and 20 testing samples.

LiTS Dataset (Bilic et al. 2023) is a CT dataset focused on liver and liver tumour segmentation. The dataset collects 201 abdominal scans. Among these, 131 scans with segmentation masks are publicly available. We use the same split of the dataset as done in (Cai et al. 2023), where 100 scans are used for training, and the remaining 31 scans for test.

To fairly compare our method with others, we adopt four commonly used metrics: Dice similarity coefficient (Dice), Jaccard similarity coefficient (Jaccard), 95% Hausdorff Distance (95HD), and Average Surface Distance (ASD).

To make sparse annotations more effective, the three selected slices should contain the foreground area of the segmentation target, and we choose slices as close to the center position as possible in all three planes. The entire training is conducted in the PyTorch framework. We set the batch size to 4, with each batch containing two volumes with labels and two volumes without labels. We train for 6000 iterations using the Stochastic Gradient Descent (SGD) optimizer. The initial learning rate is set to 0.01 and gradually decays to 0.0001. The value of parameter $\alpha$ is initialized to 0.1 and is increased every 150 iterations.

In this paper, we compare our SGTC with 7 recent SOTAs semi-supervise medical image segmentation methods, including MT (Tarvainen and Valpola 2017), UA-MT (Yu et al. 2019), SASSNet (Li, Zhang, and He 2020), CPS (Chen et al. 2021), DTC (Luo et al. 2021), BCP (Bai et al. 2023) and Desco (Cai et al. 2023) on three challenging benchmarks, i.e., LA2018 (Xiong et al. 2021), KiTS19 (Heller et al. 2019), and LiTS (Heller et al. 2019) datasets. In other methods and the SGTC (Dual) version, we annotate two slices per volume according to the previous settings (Cai et al. 2023). In the SGTC version, we annotate three slices per volume. The S, C, A indicates whether the annotated slices is sagittal, coronal or axial plane.

Comparison Results on LA2018 dataset: Table 1 shows the comparison between our method and these state-of-the-art models on LA2018 dataset under 10% (8 samples), which shows our method achieves the best performance (i.e., surpassing the second best by 6.3% on Dice). Due to the semantic guidance, the network can comprehensively perceive subtle changes in boundary regions and voxel-level semantic information, which makes the network perform better in dealing with complex anatomy structures. Additionally, methods using two orthogonal annotations generally perform poorly, whereas our SGTC with three orthogonal annotations achieves better shape-related performance by modeling the entire volume’s supervision signal distribution more comprehensively. Visual comparisons are shown in Figure 3, where our SGTC accurately delineates the intricate structures and weak boundaries. More visualizations can be found in the supplementary materials.

Comparison Results on KITS19 dataset: Table 2 demonstrates the comparison between our method and recent SOTAs on KITS19 dataset under 10% (19 samples) labeled data settings. We can find that our SGTC outperforms existing methods in all metrics. For example, compared with BCP (Bai et al. 2023) and DTC (Luo et al. 2021), SGTC shows 1.2% and 1.5% improvements on Dice. Besides, the qualitative comparison results are shown in Figure 4, which shows our method effectively handles complex boundaries and fine structures.

Comparison Results on LITS dataset: Table 3 presents the comparison between our method and others on the LITS dataset under 10% (10 samples) labeled data settings. It can be observed that our method’s advantage becomes more pronounced with fewer labeled data. The main reason behind this is that our method achieves semantic-aware and fine-granular segmentation. Moreover, visual comparison results in Figure 5 that show our SGTC effectively handles complex boundaries and fine structures.

Ablation study of each component: As shown in Table 4, we conduct an ablation experiment on the KiTS19 dataset with 9 volumes to better evaluate the components of our method. By incorporating semantic cues, our method is able to focus more effectively on boundary information, leading to improved performance.

Ablation study of triple-view disparity training strategy: As shown in Table 5, where all these methods use the same annotated strategy (i.e., CAC or SCA). The results indicate that our method outperforms others under the same settings, and SCA provides better performance improvement, demonstrating that the performance gains are due to the more effective triple-view disparity training strategy.

Ablation study of text prompts in CLIP: Table 6 illustrates the results of using different text prompts. It can be observed that, compared to solely using image features, combining visual and textual modalities brings large performance gains. Our tailored text prompt, which effectively integrates semantics into the network by describing more details, resulting in more promising outcomes.

Ablation study of the different annotation method: Figure 6 illustrates segmentation performance by using our proposed three orthogonal annotations, the two orthogonal annotations (i.e., the third slice is parallel to one of the first two) and the parallel annotations (i.e., three parallel slices in the same plane). It is shown that our annotation scheme achieves superior results at different iterations. Additionally, Figure 7 shows t-SNE visualization of the extracted features using different annotation methods. Specifically, we trained five networks with just one single annotated slice per volume: three orthogonal slices for training $s$, $c$, and $a_{1}$, and three parallel slices for training $a_{1}$, $a_{2}$, and $a_{3}$. The features extracted using three orthogonal annotations are more concentrated, demonstrating the effectiveness of our triple-view disparity training strategy.

Parameter Analysis: Table 7 shows the analysis of the hyper-parameter $\alpha$ in Eq. 10. Specifically, we conducted ablation experiments on the KITS19 dataset by setting $\alpha$ to 0.1, 0.15, 0.2, 0.25, 0.3, and the proposed dynamic coefficient. It is observed that the proposed dynamic coefficient yields the best performance, as it makes the training more stable.