TCCNet: Temporally Consistent Context-Free Network for Semi-supervised Video Polyp Segmentation

1 Introduction

Automatic video endoscopic polyp segmentation is of great research value for the prevention of Colorectal Cancer (CRC). However, most previous methods are fully supervised and trained on static images. They have three major limitations. Firstly, recent methods like [Fan et al., 2020; Kim et al., 2021] only relied on static images to train and evaluate their models, while ignoring the temporal information in an endoscopic sequence. As shown in Fig. 1, the images are from the same endoscopic sequence and focus on the same polyp object. Their trajectories and appearance changes are temporally correlated. It is insufficient to only focus on independent static images for the Video Polyp Segmentation (VPS) task. Secondly, limited annotated data is the bottleneck for the VPS task. Due to the fuzzy boundary of polyp and its similarity to the background tissue, even skilled clinicians may fail to reach an agreement on the annotations for successive frames, as shown in the first row of Fig. 1. Thirdly, although an endoscopic video focuses on the same polyp tissue, the context of the polyp (i.e., lumen, specularity and mucosa tissue) changes due to angles of camera or lights, which may impact on the prediction results of adjacent frames. E.g., in the second row of Fig. 1, the contexts of the 2nd and 3rd frames are different from that of the 1st one.

To tackle the aforementioned limitations, we propose a novel Temporally Consistent Context-Free Network (TCCNet) for the semi-supervised VPS task, which contains a segmentation branch and a propagation branch. For the first limitation, we design a Sequence-Corrected Reverse Attention (SC-RA) module and a Propagation-Corrected Reverse Attention (PC-RA) module to fully exploit the temporal information and keep the prediction temporally consistent among successive frames. Reverse Attention [Fan et al., 2020] is employed to refine the boundary of saliency map. To correct the map's prediction error, we introduce an error correction mechanism. Positional information [Hu et al., 2021] is obtained from the endoscopic sequence and effectively prevents a suspected area from being misclassified. For the second limitation, we design a co-training scheme to train the network in a semi-supervised manner. A sequence clip with one labeled reference frame and several unlabeled frames is fed into the two parallel branches. For the unlabeled frames, the outputs of one branch are supervised by the pseudo labels [Chen et al., 2021] generated by the other branch. Different from existing semi-supervised methods, we focus on the temporal consistency between the ground truth of the reference frame and the pseudo labels of the unlabeled frames. For the third limitation, we propose a Context-Free Loss to mitigate the impact of varying contexts within successive frames.

The contributions are three-fold: (i) We propose a novel TCCNet for the semi-supervised VPS. To the best of our knowledge, this is the first work to study the VPS task using semi-supervised learning. (ii) We design the SC-RA and PC-RA modules to keep predictions temporally consistent and introduce the Context-Free Loss to alleviate the impact of varying contexts. (iii) We conduct experiments on three VPS datasets. Results show that even trained under 1/15 label ratio, TCCNet is comparable to the state-of-the-art fully supervised methods for VPS. Specially, TCCNet exhibits obvious superiority over the existing semi-supervised methods for natural image and other medical image segmentation. 1

2 RelatedWork

Polyp Segmentation. Deep Convolutional Neural Networks (DCNN) are widely used in polyp segmentation. UNet++ [Zhou et al., 2018] designed dense connections and enabled feature fusion of varying scales. ResUNet++ [Jha et al., 2019] merged residual blocks, attention blocks, Atrous Spatial Pyramidal Pooling and Squeeze-and-Excitation (SE) blocks into a U-shaped architecture. PraNet [Fan et al., 2020] fused features from varying scales with a Parallel Partial Decoder (PPD) [Chen et al., 2018] and recovered the boundary cues with a Reverse Attention (RA) module. Inspired by the overall structure of PraNet, UACANet [Kim et al., 2021] integrated the features from the foreground, background and uncertain area to refine the boundary cues. However, these methods only work on static images and cannot capture the temporal information in the endoscopic videos. PNSNet [Ji et al., 2021] proposed a Normalized Self-attention (NS) block to dynamically update the receptive field of the network and obtain the temporal representation. Limited by the scale of video polyp datasets, PNSNet was pre-trained on static images and fine-tuned on video polyp images. Such a training strategy requires abundant annotations and does not suit the semi-supervised segmentation. By contrast, our semisupervised framework fully exploits the temporal information. It keeps the prediction temporally consistent and breaks the bottleneck of limited annotations.

Semi-supervised Medical Image Segmentation. Semisupervised Learning (SSL) improves the performance of medical image segmentation by utilizing a large set of unlabeled data. Adversarial learning [Zhang et al., 2017] and consistency regularization [Tarvainen and Valpola, 2017] are the two most common semi-supervised methods. To avoid the inefficient perturbation, ICT [Verma et al., 2019] introduced a consistency strategy which encouraged the consistency between the prediction of two unlabeled images' interpolation and the interpolation of those two images' predictions. UA-MT [Yu et al., 2019] took the reliability of predictions into consideration and designed a consistency loss with the guidance of an estimated uncertainty. Apart from the input perturbation, feature perturbation [Ouali et al., 2020] and model perturbation [Chen et al., 2021] are also used for consistency regularization. URPC [Luo et al., 2021] presented an uncertainty rectified pyramid consistency loss which encouraged the prediction at different scales to be consistent. In our framework, we introduce the model perturbation [Chen et al., 2021] method into the VPS task and preserve the temporal consistency between the ground truth of the reference frame and the pseudo label of the unlabeled frame.

1 Code is available at https://github.com/wener-yung/TCCNet

3 Methodology

3.1 Temporally Consistent Context-Free Network

Fig. 2(a) shows the overall architecture of the proposed TCCNet. It consists of a segmentation branch and a propagation branch. In the training phase, a training clip is parallelly fed into the two branches. The SC-RA module in the segmentation branch and the PC-RA module in the propagation branch are designed to correct the error of saliency maps from the previous layer and keep the temporal consistency of the predictions. The outputs of these two branches are supervised with a co-training scheme. From each input image, we synthesize two different images S1 and S2 with the predicted polyp locations, which are fed into the two branches again to obtain the global maps and calculate the CFLoss in Sec. 3.4.

The encoder structures of the two branches are the same. For a given endoscopic clip with T frames, {It}Tt=1 , five levels of features {F l t }T ,5 t=1,l=1 ∈ RHl×W l×C are extracted from the encoder. For fair comparison, we use the same backbone in PraNet [Fan et al., 2020] as the encoder (i.e., Res2Net pre-trained on ImageNet). Following [Fan et al., 2020], we integrate the high-level features to obtain the global features {F 6 t }Tt=1 and the global maps {M 6 t }Tt=1 . To improve the quality of feature embedding and reduce computing resources, we update the encoder weights in the propagation branch as an Exponential Moving Average (EMA) [Tarvainen and Valpola, 2017] of the weights in the segmentation branch.

The decoder structures of these two branches are different, so are the prediction processes. Specifically, in the segmentation branch, taking the lth layer as an example, {F l s,t }Tt=1 , {F l+1 s,t }Tt=1 and {M l+1 s,t }Tt=1 are sent to the SC-RAmodule to obtain {M ls,t }Tt=1 . The predictions of the T frames are calculated as a whole. While in the propagation branch, the predictions are calculated frame by frame. Taking the tth frame as an example, the feature maps of the previous t-1 frames are stored as memory features in a memory pool. In the lth layer, the memory features, F lp,t and M l+1 p,t are sent to the PC-RA module, andM lp,t is therefore obtained.

TCCNet focuses on the VPS task with a semi-supervised method, where a training clip {It}Tt=1 consists of one labeled frame and T - 1 unlabeled frames. We denote the labeled frame as the reference frame (Ir, Yr) and the unlabeled frames as Du = {It}Tt=2 . The training clip is sent to the two parallel branches. The predictions of segmentation branch are denoted by {Ps,t}Tt=1 and the predictions of propagation branch are denoted by {Pp,t}Tt=2 . It should be highlighted that the segmentation branch predicts the entire clip and outputs T masks, while the propagation branch outputs T - 1 masks of the unlabeled frames, excluding the mask of the reference frame.

The supervised loss on the reference frame is formulated using the Binary Cross Entropy Loss (BCELoss) and the Intersection over Union Loss (IoULoss). It is calculated as:

Lsup(Ps,r, Yr) = LBCE(Ps,r, Yr) +LIoU(Ps,r, Yr), (1) where LBCE(·) and LIoU(·) denote the BCELoss and IoULoss functions, respectively.

The unsupervised loss on the unlabeled frames is defined using the Cross Pseudo Supervision Loss (CPSLoss) [Chen et al., 2021]. We get the pseudo segmentation labels {Y ′ s,t }Tt=2 and {Y ′ p,t }Tt=2 by thresholding the corresponding predictions. The CPSLoss is bidirectional. Specifically, the predictions {Ps,t}Tt=2 of the segmentation branch are supervised by the pseudo labels {Y ′ p,t }Tt=2 generated from the propagation branch and vice versa. It is calculated as:

3.2 Sequence-Corrected Reverse Attention

The SC-RA module is described in Fig. 2(b). Taking the lth layer as an example, the SC-RA module receives feature maps {F l s,t }Tt=1 , {F l+1 s,t }Tt=1 and segmentation maps

{M l+1 s,t }Tt=1 to calculate the sequence-corrected position maps {Mpos,t}Tt=1 , which are used to correct {M l+1 s,t }Tt=1 . {Mpos,t}Tt=1 is the average of Mpos and M ′pos . Specifically, we use {F l s,t }Tt=1 to compute the position map Mpos and {F l+1 s,t }Tt=1 to compute M ′pos . The calculation processes of Mpos and M ′pos are similar. Taking Mpos as an example, we complement the feature maps {F l s,t }Tt=1 with 2D positional encoding [Carion et al., 2020] and calculate the query vector Q and key vector K with the 1×1×1 convolutional layer. We reshape Q and K, and apply matrix inner product to obtain the similarity map Sim.

The segmentation maps {M l+1 s,t }Tt=1 from the previous layer are used to reduce the response of background region. We get the local response by:

where σ is the sigmoid function; and U(·) is the up-sampling operation. g is then reshaped to g′l ∈ R1×THlW l .

We apply element-wise product between Sim and g′ to suppress the response of non-polyp region. Then we select the top-K values on the key dimension and average them to get position map M lpos .

The sequence-corrected position maps {Mpos,t}Tt=1 are used to compute the sequence-corrected segmentation maps {M lSC,t }Tt=1 by:

The current segmentation maps are then calculated using the Reverse Attention [Fan et al., 2020],

where convs(·) denotes the convolution layers.

We adopt supervision on the segmentation predictions from the middle layers and the sequence-corrected segmentation maps. The calculation is:

where PSC = σ(U(MSC)) and Ps = σ(U(Ms)).

3.3 Propagation-Corrected Reverse Attention

The PC-RA module is shown in Fig. 2(c). The Space-Time Memory (STM) [Oh et al., 2019] module aims to utilize all the available temporal cues. For the tth frame, the feature map F lp,t along with the segmentation map M l+1 p,t are embedded as a pair of 2D key and value maps, which are defined as:

where ϕq(·) and gq(·) are two parallel 3 × 3 convolutional layers, and conp(·) is a 7 × 7 convolutional layer.

Each of the previous frame is independently embedded into a pair of key and value maps by following the above strategy. We concatenate them in the temporal dimension and store them in the memory pool. KlM ∈ RT ′×Hl×W l×C/8 and V lM ∈ RT ′×Hl×W l×C/2 are obtained, where T ′ is the number of memory frames. The memory embedding and the current feature embedding are fed into the STM module to obtain a memory map F lmem,t .

The formulation of the propagation-corrected map MPC,t is similar to that of the sequence-corrected map in Sec. 3.2. Similarity map Sim is calculated with the feature map of the tth frame and that of the reference frame. The ground truth of the reference frame Yr is used to calculate the local response by gr = exp(Yr)/e. We then get the position map Mpos,t by applying element-wise product between Sim and gr, selecting top-K values on the dimension and reshaping operation.

The propagation-corrected segmentation map M lPC,t is then calculated similarly to the sequence-corrected segmentation map in Eq. (4). We apply the Reverse Attention on F lmem,t andM lPC,t to get the segmentation mapM lp,t by Eq. (5). Similar to Sec. 3.2, we adopt the supervision upon the segmentation maps. The loss function is defined as:

3.4 Context-Free Loss

In an endoscopic video, the context around one polyp may vary due to angles of camera or lights, which may affect the prediction results. For our semi-supervised VPS task, the network may overfit on the limited labeled data and context. Therefore, we propose a Context-Free Loss (CFloss) to reduce the dependency on varying contexts and improve the robustness of the network.

For a training clip, we obtain the locations of polyps by applying average, erosion and dilation operations on {Ps,t}Tt=2 and {Pp,t}Tt=2 . For each frame in this sequence, two patches are randomly cropped with an overlapping region, which must contain the polyp tissue. To further increase the diversity of context, we randomly sample two frames from two different sequences as the background images and synthesize two images by overlaying the two patches onto the background images, thus obtaining two synthesized images S1and S2.

The synthesized images are fed into the segmentation branch and the propagation branch, and two global maps are obtained. Denoting the two segmentation maps of the overlapping region as Ωs,1 and Ωp,2, our CFLoss is bidirectional and formulated as:

where i ∈ Ω indicates a pixel in the overlapping region.

3.5 Training Strategy

Pre-training on pseudo sequences. We generate a pseudo sequence dataset with static frames to pre-train the network, which is usually used in Video Object Segmentation (VOS) task [Oh et al., 2019; Hu et al., 2021]. For a training clip, the first frame is sampled from the labeled frames and the rest are generated by applying random affine transforms on the first frame, such as translation, zooming, cropping, flip and rotation. In the pre-training phase, only the labeled frames are used and our network is fully supervised.

Main-training on real sequences. In the main-training phase, our network is semi-supervised. For a training clip, the first frame is sampled from the labeled frames and works as the reference frame. The rest unlabeled frames are randomly selected from the same sequence in temporal order.

Total loss. The total loss of the network is formulated as:

where λcps, λs, λp and λcf are the hyper-parameters to balance the loss terms.

4 Experiments

4.1 Experimental Setup

Datasets. We conduct experiments on three video-based polyp datasets: CVC-300 [Tajbakhsh et al., 2016], CVC612 [Bernal et al., 2015] and ETIS [Silva et al., 2014]. CVC-300/CVC-612/ETIS contains 300/612/196 images from 13/29/26 video clips. The image resolutions are 384 × 288, 500×574 and 1225×966 pixels, respectively. Following [Ji et al., 2021], 60% images from CVC-612 and 60% images from CVC-300 are used for training. For the semi-supervised training, we label an image every 15 frames for each sequence. That is, the label ratio is 1/15 in each sequence. We test the performance of TCCNet on the test datasets in [Ji et al., 2021], including CVC-300-TV, CVC-612-V and CVC-612T. Moreover, to verify the generalization ability of our TCCNet, we evaluate our network on ETIS, of which the image domain is unseen in the training set.

Evaluation metrics. We employ the same evaluation metrics in [Fan et al., 2020], including mean Dice (mDice), mean IoU (mIoU), weighted Fβ measure (wFm), S-measure (Sm), Enhanced-alignment measure (Em) and Mean Absolution Error (MAE).

Implementation details. The clip of input sequence T is set to 3 and the batch size is set to 2. In the pre-training phase, the network is trained for 200 epochs with Adam optimizer and a learning rate of 10-4 . In the main-training phase, the network is trained for 40 epochs with Adam optimizer. The initial learning rate is set to 10-4 with polynomial decay [Kim et al., 2021]. All images are resized to 352 × 352 and normalized into [-0.5, 0.5]. Random data augmentation is performed. In the testing phase of TCCNet, we only use the output from the segmentation branch. C, K, λcps, λs, λp and λcf are empirically set to 32, 8, 1, 1, 1 and 2, respectively.

4.2 Comparison with State-of-the-art Methods

Comparison with methods for fully supervised polyp segmentation. We compare our semi-supervised TCCNet with the state-of-the-art (SOTA) fully supervised polyp segmentation models, including U-Net++ [Zhou et al., 2018], ResUNet++ [Jha et al., 2019], PraNet [Fan et al., 2020], UACANet [Kim et al., 2021] and PNSNet [Ji et al., 2021]. We re-train the SOTA models with the released code under their default settings for further epochs to ensure fairness. The comparison results are shown in Table 1. It can be observed that our semi-supervised model trained under 1/15 label ratio is comparable to the fully supervised models. Under 1/2 label ratio, our TCCNet outperforms the fully supervised models in almost all cases. All the SOTA models perform poorly on the EITS dataset, since the image domain of ETIS is unseen in the training set and it is easy for the models to overfit on the CVC-300 and CVC-612 datasets. Thanks to the consistency regularization, our network under 1/15 label ratio gains 3.3% mDice improvement over the best SOTA fully supervised polyp segmentation model PraNet.

Comparison with methods for semi-supervised segmentation. We compare our method with the existing semisupervised models, including UA-MT [Yu et al., 2019] and URPC [Luo et al., 2021] for other medical image segmentation tasks and ICT [Verma et al., 2019], CCT [Ouali et al., 2020] and CPS [Chen et al., 2021] for natural image segmentation task. For fair comparison, we change the backbones of these SOTA networks to our segmentation network in the segmentation branch and train the networks with the same data augmentation and training strategy. The comparison results are listed in Table 2. It can be seen that our method achieves superior performance to other methods in almost all cases.

4.3 Ablation Studies

Ablation studies for different modules. The ablation results of different modules on the CVC-300-VT and CVC612-V datasets are shown in Table 3. From the first four rows, we can observe that the proposed modules, SC-RA module, PC-RA module and CFLoss gain improvement over Baseline (the first row). The combination of SC-RA module and PCRAmodule introduces the sequence information into the two branches and further improves the performance, as shown in the fifth row. The full model (the last row) achieves better results, which is attributed to the context-free constraint.

Ablation studies for training strategy. Table 4 lists the effect of our pre-training strategy. It can be seen that the pretraining strategy brings 2.4%/2.8% of mDice and 2.7%/3.3% of mIoU increase on CVC-300-VT/CVC-612-V, respectively.

4.4 Qualitative Results

Temporal consistency of predictions. The quantitative comparison results between our model and other fully supervised VPS models are shown in Fig. 3. It can be observed that our model well maintains the temporal consistency of the predictions for a testing clip. Specifically, for the third frame, other models are likely to be influenced by the optical artifact region (blue square), leading to the misclassification as polyp tissue. Our model can learn the position and texture of polyps from previous frames, thus substantially alleviating false positives of segmentation.

Analysis on various label ratios. We visualize the improvement over the supervised baseline under different label ratios, as shown in Fig. 4. The time intervals for annotations are set to 50, 15, 5, 2 and 1. Our network consistently surpasses Baseline. To be specific, the mDice gains of our network over Baseline are 18.1%/8.6%, 9.1%/6.5%, 3.1%/2.5%, 2.6%/1.7% and 2.1%/1.4% on CVC-300-VT/CVC-612-V.

Qualitative results of the error correction mechanism for segmentation map. Fig. 5(a) shows the qualitative results of the error correction mechanism for segmentation map. Using SC-RA module as an example, a suspected region is misclassified as the polyp tissue in the global segmentation map M 6 s but later corrected toM 5 SC by the SC-RA module.

Qualitative results of the proposed CFLoss. Fig. 5(b) illustrates the qualitative results of the proposed CFLoss. Given the three consecutive frames, the context of polyp in the 2nd and 3rd frames are quite different from that in the 1st one. As a result, the network without CFLoss misclassifies the background region as polyp tissue for the second and third frames. It can be observed that our network effectively reduces such false positives with the aid of CFLoss.

5 Conclusion

In this paper, we introduced a novel semi-supervised video polyp segmentation network, i.e., TCCNet. It consists of a segmentation branch and a propagation branch, which are trained with a co-training scheme. A SC-RA module and a PC-RA module are designed to keep the predictions temporally consistent for consecutive frames. A CFLoss is proposed to reduce the impact of varying contexts on the predictions of adjacent frames. Experiments demonstrate that our TCCNet gains better results than the SOTA methods on both fully supervised polyp segmentation and semisupervised medical segmentation. In the future, we would like to extend our method to other challenging tasks, such as instrument segmentation from robotic surgical videos.

Acknowledgments

This work is supported by National Science and Technology Innovation 2030 - Major Project (No. 2021ZD0114001; No. 2021ZD0114000), National Natural Science Foundation of China (No. 61976057; No. 62172101), the Science and Technology Commission of Shanghai Municipality (No. 21511101000; No. 20511101203; No. 20511101403), the Science and Technology Major Project of Commission of Science and Technology of Shanghai (No. 2021SHZDZX0103), Shanghai Natural Science Foundation (No. 19ZR1417200). Yuejie Zhang and Tao Zhang are corresponding authors.