C³PS: Context-aware Conditional Cross Pseudo Supervision for Semi-supervised Medical Image Segmentation

We sample a labeled patch $(_{l}X,_{l}Y)$ from $\mathcal{D}_{l}$. Given a conditional label $c_{l}$ (details on generating $c_{l}$ will be described below). We can get predictions from $\mathcal{F}^{R}(.)$ and $\mathcal{F}^{C}(.)$ as follows:

The supervised loss $\mathcal{L}^{r}_{sup}$ for RNet and $\mathcal{L}^{c}_{sup}$ for CNet can be defined respectively as follows:

where $\mathcal{L}_{ce}$ and $\mathcal{L}_{dice}$ denote respectively the cross entropy loss and the Dice loss. $\mathbbm{1}\{_{l}Y=c_{l}\}$ aims to generate a binary mask depending on whether the label of a voxel is equal to $c_{l}$ or not. Finally, the overall supervised loss $\mathcal{L}_{sup}$ can be defined as follows:

When patch-based strategy is used to train a network, predictions of a voxel may be different when the context is different (e.g., in different patches). Liu et al. [10] sampled two overlapped patches in each training iteration to make the model more robust to the context change. Followed by this work, we incorporate context-awareness to increase the quality of pseudo label generation. Specifically, for each iteration, we require that two sampled patches ${}_{u}X_{1}$ and ${}_{u}X_{2}$ have an overlapping region $\Omega_{o}$ (shown in $\mathcal{D}_{u}$ of Fig. 1-(a)). We further require that the CCPS loss defined above is only computed in the region $\Omega_{o}$. Then we feed ${}_{u}X_{1}$ and ${}_{u}X_{2}$ to $\mathcal{F}^{R}(.)$ and $\mathcal{F}^{C}(.)$ and get predictions using Eq. (5):

The predictions of the overlapped region in ${}_{u}X_{1}$ and ${}_{u}X_{2}$ from the RNet are ${}^{R}_{u}\mathbf{P}_{o1}$ and ${}^{R}_{u}\mathbf{P}_{o2}$, respectively. Similarly, the predictions of the overlapped region in ${}_{u}X_{1}$ and ${}_{u}X_{2}$ from the CNet are ${}^{C}_{u}\mathbf{P}_{o1}$ and ${}^{C}_{u}\mathbf{P}_{o2}$, respectively. We generate pseudo labels of voxels in $\Omega_{o}$ by combining two predictions with different context as follows:

Finally, the loss $\mathcal{L}_{C^{3}PS}$ for the unlabeled data after incorporating context-awareness is computed as follows:

The overall training objective $\mathcal{L}$ of our proposed approach is:

where $\lambda$ is a weight factor defined as: $\lambda(t)=0.1\times e^{(-5(1-\frac{t}{t_{total}})^{2})}$, where $t$ denotes the current iteration and $t_{total}$ is the total iteration number.

Datasets. We evaluated our method on two public datasets: Beyond the Cranial Vault (BCV) dataset[16] and The MICCAI’17 Multi-Modality Whole Heart Segmentation challenge (MMWHS) dataset[17]. The BCV dataset contains 30 CT images with annotation for 13 organs. We chose to segment five abdominal organs including liver, spleen, pancreas, left kidney and right kidney. We used 24 samples (4 labeled and 20 unlabeled) for training and the remaining 6 samples for testing. The MMWHS consists of 20 cardiac CT samples with annotations for seven structures: left ventricle (LV), right ventricle (RV), left atrium (LA), right atrium (RA), pulmonary artery (PA), my-ocardium (MYO) and ascending aorta (AA). We took 14 samples (2 labeled and 12 unlabeled) for training and used the remaining 6 samples for testing.

Implementation details. We chose 3D U-Net[18] as the RNet and the conditional 3D U-Net (see Fig. 1-(b) for an illustration) as the CNet [19]. Please note, for a fair comparison with other SOTA SSMIS methods, we replaced the nnUNet used in [19] by 3D U-Net. We used a patch size of 160 $\times$ 160 $\times$96 for training. In total we trained our network 20000 iterations. We used SGD optimizer and set the initial learning rate to 0.01. Then at each iteration the initial learning rate was multiplied by $(1.0-\frac{t}{20000})^{0.9}$, where $t$ is current iteration. We implemented our method based on PyTorch framework and conducted the evaluation on a NVIDIA Tesla V100 GPU. Starting from ($t_{s}=15000$)th iteration, we conducted hard organ learning. On the BCV dataset, we chose left kidney, right kidney and pancreas as the hard organs while on the MMWHS dataset, we chose MYO, RA, AA and RV as the hard organs. We empirically set the thresholds $\gamma_{1}$ and $\gamma_{2}$ to 0.95 and 0.9, respectively. We use Dice score coefficient (DSC; %) and Average Surface Distance (ASD; Voxel) as the evaluation metrics. Paired t-test is used to check whether a difference is statistically significant or not. We set the significant level as 0.05.

We compared our methods with six SOTA SSMIS methods [11, 12, 13, 14, 4, 15]. Baseline and Upper Bound are obtained by training the RNet in a supervised manner.

BCV dataset. Table 1 presents the results on the BCV dataset. It can be found that the proposed C³PS method achieves better DSC performance than other 6 SOTA methods with a large margin when 4 labeled data were used. Additionally, from this table, one can also find that our method outperforms CPS with a large margin in terms of DSC (on average an increase of 5.2%). Paired t-test showed that the difference between our method and the CPS method [4] was statistically significant ($p$=0.011).

MMWHS dataset. Table 1 also presents the results on the MMWHS dataset. The proposed method achieved a superior performance over other SOTA methods with an average DSC of 81.7%. Paired t-test showed that the difference between our method and the CPS method [4] was statistically significant ($p$=1e-4).

Fig. 2 shows a visual comparison of the top-5 methods as well as our own method when applied to these two datasets. From this figure, one can see that our method predicts more consistent segmentation results with the ground truth than other methods.

Ablation Study. We conducted ablation studies on the BCV dataset to show the effectiveness of each individual components. As shown in Table 2, each component contributes to the superior performance of the proposed approach. We further investigate the effectiveness of HOL. The results are presented in Table 3. With HOL, we observed an improved performance for left and right kidney as well as pancreas.