Self-Supervised Correction Learning for Semi-Supervised Biomedical Image Segmentation

In this work, we adopt a coarse-to-fine strategy and propose a self-supervised correction learning paradigm for semi-supervised biomedical image segmentation. Specifically, we introduce inpainting as the pretext task of self-supervised learning to take advantage of massive unlabeled data and thus construct a dual-task network, as shown in Fig. 1. Our proposed framework is composed of a shared encoder, two decoders and five GFF modules placed on each layer of both decoders. We utilize ResNet34 [8] pretrained on the ImageNet [5] as our encoder, which consists of five blocks in total. Accordingly, the decoder branch also has five blocks. Each decoder block is composed of two Conv-BN-ReLU combinations. For the convenience of expression, we use ${E}_{seg}$, ${D}_{seg}$ to represent the encoder and decoder of the segmentation branch, and ${E}_{inp}$ and ${D}_{inp}$ for those of the inpainting branch.

In the first step, given the image $x\in\mathbb{R}^{H\times W\times C}$, in which $H$,$W$,$C$ are the height, width and channels of the image respectively, we use the segmentation branch ${E}_{seg}$, ${D}_{seg}$ with skip-connections, the traditional U-shape encoder-decoder structure, to obtain a coarse segmentation map $\hat{y}_{coarse}$ and then mask the original input based on its binary result $\overline{y}_{coarse}$ by the following formulas:

$$\hat{y}_{coarse}=D_{seg}(E_{seg}(x))$$

(1)

$${x}_{mask}=x\times(1-\overline{y}_{coarse})$$

(2)

In the second phase, the original image $x$ and the masked image ${x}_{mask}$ are sent into ${E}_{seg}$ and ${E}_{inp}$ simultaneously to extract features ${e}_{seg}$ and ${e}_{inp}$. Obviously, ${e}_{seg}$ is essential for the inpainting task, and since the initial segmentation is usually inaccurate and incomplete, ${e}_{inp}$ may also contain important residual lesion features for the correction of the initial segmentation. In order to adaptively select the useful features of ${e}_{inp}$ and achieve complementary fusion of ${e}_{seg}$ and ${e}_{inp}$, we design the GFF modules ($G_{1}$-$G_{5}$) and place them on each decoder layer of both branches. Specifically, for the $i^{th}$ layer, the features ${e}_{seg}^{i}$ and ${e}_{inp}^{i}$ are delivered into ${G}_{i}$ through skip-connections to obtain the fusion ${e}^{i}$ = $G_{i}$(${e}_{seg}^{i}$, ${e}_{inp}^{i}$), and then sent to the corresponding decoder layer. Thus, both $G_{i}$ of the two branches shown in Fig. 1 actually share parameters, taking the same input and generating the identical output. To enhance the learning of the GFF modules, we adopt a deep supervision strategy and each layer of the two decoder branches generate a segmentation result and an inpainting result respectively by the following formulas:

$$\hat{y}_{fine}^{i}=\left\{\begin{aligned} D_{seg}^{i}([{e}^{i},d_{seg}^{i+1}])&,\quad i=1,2,3,4\\ D_{seg}^{i}({e}^{i})&,\quad i=5\end{aligned}\right.$$

(3)

$$\hat{x}^{i}=\left\{\begin{aligned} D_{inp}^{i}([{e}^{i},d_{inp}^{i+1}])&,\quad i=1,2,3,4\\ D_{inp}^{i}({e}^{i})&,\quad i=5\end{aligned}\right.$$

(4)

Where [$\cdot$ , $\cdot$] denotes the concatenation process, and $d_{seg}^{i+1}$, $d_{inp}^{i+1}$ represent the features from previous decoder layers. The deep supervision strategy can also avoid ${D}_{inp}$ directly copying the features of the low-level ${e}_{seg}$ to complete the inpainting task without in-depth lesion feature mining. The output of the last layer $\hat{y}_{fine}^{1}$ is the final segmentation result of our network.

To better incorporate complementary features and filter out the redundant information, we design the GFF modules placed on each decoder layer to integrate the features delivered from the corresponding encoder layer of two branches. The details are shown in Fig. 2. Our GFF module consists of a reset gate and a select gate. Specifically, for the $G_{i}$ placed on the $i^{th}$ decoder layer, the value of two gates is calculated as follows:

$$r_{i}=\sigma(W_{r}\left[{e}_{seg}^{i},{e}_{inp}^{i}]\right)$$

(5)

$$s_{i}=\sigma(W_{s}\left[{e}_{seg}^{i},{e}_{inp}^{i}]\right)$$

(6)

Where $W_{r}$, $W_{s}$ denote the convolution process, taking the concatenation of ${e}_{seg}^{i}$ and ${e}_{inp}^{i}$ as input. $\sigma$ represents the Sigmoid function. ${r}_{i}$ and ${s}_{i}$ represent the value of the reset gate and the select gate, respectively. Since the input of the inpainting branch is the masked image, the reset gate is necessary to suppress massive invalid background information. And then the select gate achieves adaptive and complementary feature fusion between the reintegrated features $\tilde{e}^{i}$ and the original segmentation feature ${e}_{seg}^{i}$ by the following operations:

$$\tilde{e}^{i}=W\left[{r}_{i}\times{e}_{inp}^{i},{e}_{seg}^{i}\right])$$

(7)

$${e}^{i}=s_{i}\times\tilde{e}^{i}+(1-s_{i})\times{e}_{seg}^{i}$$

(8)

where $W$ also represents the convolution process to make the reintegrated features $\tilde{e}^{i}$ have the same dimension with ${e}_{seg}^{i}$.

We only calculate loss in the second stage. The labeled dataset and unlabeled dataset are denoted as $D_{l}$ and $D_{u}$. For the labeled data $x_{l}\in D_{l}$, $y_{l}$ is the Ground Truth. Since we adopt the deep supervision strategy, the overall loss is the sum of the combination of Binary CrossEntropy (BCE) loss and Dice loss between each output and the Ground Truth:

$$\mathcal{L}_{seg}(x_{l})=\sum\limits_{i=1}^{5}{\mathcal{L}_{BCE}^{i}(\hat{y}_{l}^{i},y_{l}^{i})+\mathcal{L}_{Dice}^{i}(\hat{y}_{l}^{i},y_{l}^{i})}$$

(9)

where $\hat{y}_{l}^{i}$, $y_{l}^{i}$ denote the segmentation map $\hat{y}_{fine}^{i}$ of the $i^{th}$ decoder layer and the corresponding down-sampling Ground Truth $y_{l}$.

For unlabeled data $x_{u}\in D_{u}$, the inpainting loss is the sum of $L1$ loss between each inpainting image and the original image in the masked region:

$$\mathcal{L}_{inp}(x_{u})=\sum\limits_{i=1}^{5}{\overline{y}_{u}^{i}\times\left|\hat{x}_{u}^{i}-x_{u}^{i}\right|}$$

(10)

where $\hat{x}_{u}^{i}$, $x_{u}^{i}$ and $\overline{y}_{u}^{i}$ represent the inpainting image, down-sampling original image and binary segmentation result of the $i^{th}$ decoder layer, respectively. In the end, the total loss function is formulated as follows:

$$\mathcal{L}=\lambda_{1}\sum\limits_{x_{l}\in D_{l}}{\mathcal{L}_{seg}(x_{l})}+\lambda_{2}\sum\limits_{x_{u}\in D_{u}}{\mathcal{L}_{inp}(x_{u})}$$

(11)

where $\lambda_{1}$, $\lambda_{2}$ are weights balancing the segmentation loss and the inpainting loss. And we set $\lambda_{1}=2$ and $\lambda_{2}=1$ in our experiments.

We conduct experiments on a variety of medical image segmentation tasks to verify the effectiveness and robustness of our approach, including polyp, skin lesion and fundus optic disc segmentation, respectively.
Polyp Segmentation We use the publicly available kvasir-SEG [10] dataset containing 1000 images, and randomly select 600 images as the training set, 200 images as the validation set, and the remaining as the test set.
Skin Lesion Segmentation We utilize the ISBI 2016 skin lesion dataset [7] to evaluate our method performance. This dataset consists of 1279 images, among which 900 are used for training and the others for testing.
Optic Disc Segmentation The Rim-one r1 dataset [6] is utilized in our experiments, which has 169 images in total. We randomly split the dataset into a training set and a test set with the ratio of 8:2.
Evaluation Metric Referring to common semi-supervised segmentation settings [22, 13], for all datasets, we randomly use $20\%$ of the training set as the labeled data, $80\%$ as the unlabeled data and adopt five metrics to quantitively evaluate the performance of our approach and other methods, including “Dice Similarity Coefficient (Dice)”, “Intersection over Union (IoU)”, “Accuracy (Acc)”, “Recall (Rec)” and “Specificity (Spe)”.

Data pre-processing In our experiments, since the image resolution of all datasets varies greatly, we uniformly resize all images into a fixed size of $320\times 320$ for training and testing. And in the training stage, we use data augmentation, including random horizontal and vertical flips, rotation, zoom, and finally all the images are randomly cropped to $256\times 256$ as input.
Training details Our method is implemented using PyTorch [18] framework. We set batch size of the training process to 4, and use SGD optimizer with a momentum of 0.9 and a weight decay of 0.00001 to optimize the model. A poly learning rate policy is adopted to adjust the initial learning rate, which is $lr=init\_lr\times(1-\frac{iter}{max\_iter})^{power}$, where $init\_lr=0.001$, $power=0.9$). The total number of epochs is set to 80.

In our experiments, ResNet34 [8] based UNet [19] is utilized as our baseline, which is trained using all training set and our selected $20\%$ labeled data separately in a fully-supervised manner. Besides, we compare our method with other state-of-the-art approaches, including DAN [24], MASSL [2], MT [20] and its variants (UA-MT [22], TCSM_V2 [13]). All comparison methods adopt ResNet34UNet as the backbone and use the same experimental settings for a fair comparison. On the Kvasir-SEG dataset, Table. 1 shows that our method obtains the outstanding performance compared with other semi-supervised methods, with Dice of $87.14\%$, which is $2.74\%$ improvement over the baseline only using the 120 labeled data, outperforming the second best method by $0.69\%$. On the ISBI 2016 skin lesion dataset, we obtain a $90.95\%$ Dice score, which is superior to other semi-supervised methods and very close to the score of $91.38\%$ achieved by the baseline using all training set images. On the Rim-one r1 dataset, we can conclude that our method achieves the best performance over five metrics, further demonstrating the effectiveness of our method. Note that detailed results on the latter two datasets are listed in the supplementary material due to the space limitation. Some visual segmentation results are shown in Fig. 3 (col.1-8).

Effectiveness of our approach with different ratio of labeled data We draw the curves of Dice score under three settings in Fig. 4. To verify that our proposed framework can mine residual lesion features and enhance the lesion representation by GFF modules in the second stage, we conduct experiments and draw the blue line. The blue line denotes that our method uses the same labeled data with the baseline (the red line) to perform the two-stage process, without utilizing any unlabeled data. Note that we only calculate the segmentation loss for the labeled data. The performance gains compared with the baseline show that our network mines useful lesion information in the second stage. The green line means that our method introduces the remaining as unlabeled data for the inpainting task, further enhancing the feature representation learning of the lesion regions and improving the segmentation performance, especially when only a small amount of labeled data is used. When using $100\%$ labeled data, the green line is equivalent to the blue line since no additional unlabeled data is utilized to do the inpainting task, thus maintaining the same results.
Effectiveness of the GFF modules To verify the effectiveness of the GFF modules, we also design two variants, which merge features by directly addition and concatenation, denoting as Ours (add) and Ours (concat) respectively. In Table. 1, we can observe performance degradation by both approaches compared with our method, proving that the GFF module plays a significant role in filtering redundant information and improving the model performance.