UACANet: Uncertainty Augmented Context Attention for Polyp Segmentation

Fully Convolutional Network (FCN) (Long et al., 2015) introduced a fully convolutional neural network architecture to efficiently train a model to classify each pixel. They also utilized a multi-scale scheme with aggregating low-level feature maps from the early stages in the neural network. Noh et al. proposed deconvolution network (Noh et al., 2015) to compensate a spatial information degradation due to the pooling operation by leveraging transposed convolution and unpooling methods. Pyramid Spatial Pooling Network (PSPNet) utilized grid-wise pooling and comprise multiple sizes of the grid to deal with multi-scale objects. Deeplab (Chen et al., 2017) constructed multiple convolution layers with different dilation rates to diversify receptive fields within a single module to extract a multi-scale context information. Dual Attention Network (Fu et al., 2019) brought a self-attention mechanism to the semantic segmentation network and comprise a dual path network for spatial and channel dimensions which incorporates local features with their global dependencies. Object Contextual Representation (OCR) (Yuan et al., 2019) expanded a non-local operation to consider semantic region by aggregating the representation of the pixels from each class and compute a similarity with each pixel from the feature map.

Salient object detection (SOD) resembles semantic segmentation but has its own criterion. Rather than predicting the entire region in the given image and classifying an object class label, SOD focuses on the object itself, identifying more important regions from the surrounding areas. Since salient object does not have any specific object class, even if the task might seem superficially easier than multi-class segmentation, it is much harder when it comes to the accurate object detection and assuming their priorities among other surrounding objects. Unlike semantic segmentation, it is hard to say that there is a strong baseline model with a high performing yet simple architecture in SOD.

Current state-of-the-art methods exploit the boundary region of objects for complementary information by multi-task learning strategy which can be incorporated to enhance the quality of saliency map. Edge Guidance Network (EGNet) (Yang et al., 2017) used a bottom-up stream for edge detection branch and side-out fusion strategy which aggregates each path from the top-down stream for salient object branch. To clarify the terms that we use as bottom-up stream means from high-level feature level to low-level feature level from the backbone network which is commonly used for segmentation or heat-map related tasks such as pose estimation. Also, side-out fusion strategy aggregates multiple feature maps from the backbone with concatenation or summation. Boundary-aware Network (BANet) (Su et al., 2019) oppositely used a side-out fusion for boundary branch and single-stream for object branch, but rather than regarding edge detection as a separate task, they combine edge result and object result to generate saliency map. These methods qualitatively showed competitive results consistently (Wang et al., 2021) which proves that edge guidance helps better representation of objects in terms of edge-preserving results. However, acquiring additional edge dataset is often expensive and even with image processing techniques such as Canny edge detection (Canny, 1986), it contains redundant edges which are usually unrelated to the object.

Reverse attention (Chen et al., 2018) explicitly multiply a reverse area of prediction in order to capture residual details for the saliency refinement. However, based on both their experimental results and our experiments, the performance with or without reverse attention wasn’t very different. However, their suggestion provided an intuition that even without explicit edge guidance, we can access to the edge related context with saliency map. Based on this idea, we broaden a reverse saliency map with additional uncertain area, an ambiguous area that saliency score is neither biased to the foreground nor the background. Also, they suggested an efficient yet practical network architecture. The decoder at the end of the backbone feature map produces a global saliency map, and in the bottom-up stream, only the saliency map from the previous level is used unlike U-Net shape networks concatenated feature maps. The final prediction thus only requires a low-level feature map and saliency map from the previous level as a guidance. While U-Net like networks predicts saliency map at the last bottom-up layer, proposed architecture from Reverse attention (Chen et al., 2018) which computes saliency map at the intermediate level feature map and saliency map from the bottom-up stream only.

Polyp segmentation aims to precisely segment polyps from the given colonoscopy image. One can say that polyp segmentation is a semantic segmentation with binary class problem. While we can adopt such fully convolutional network architecture to solve the problem, since colonoscopy has different image domain compared to the general images, researchers focus on extracting semantic features with detail information. Famous architectures like PSPNet or Deeplab adopt multi-scale scheme to the backbone network which helps to capture detail information with multiple receptive field sizes, but they usually deploy such methods at the end of the backbone network which is spatially insufficient for recovering precise spatial information. DeeplabV3+ (Chen et al., 2017) tried to compensate such problem by concatenate low-level feature map from the backbone network but wasn’t sufficient for recovering such details coming from the input image.

U-Net (Ronneberger et al., 2015) introduced incremental up-sampling of the feature maps alongside the corresponding scales of the low-level feature maps. Such ”skip-connections” also appears in Feature Pyramid Networks (FPN) (Lin et al., 2017) but they use element-wise addition while U-Net aggregate features with channel-wise concatenation. While FPN is also designed similar to the U-Net in order to extract multi-scale features, they do not need a fine-grained detail features since object detection needs to locate a bounding box, not a precise shape of an object. U-Net++ (Zhou et al., 2018) added extra layers and dense connectivity to reduce the gap between low-level and high-level features.

As the significance of polyp segmentation has been increased, studies that only dedicated to the polyp datasets have been done recently. ResUNet++ (Jha et al., 2019) construct a U-Net shape network with famous CNN modules including residual blocks from ResNet, Atrous Spatial Pyramid Pooling (ASPP) module from Deeplab and Squeeze and Excitation mechanism from SENet (Hu et al., 2018). Selective Feature Aggregation network (SFA) (Fang et al., [n.d.]) added another bottom-up path for boundary estimation, similar to SOD networks with edge guidances (Yang et al., 2017; Su et al., 2019). Parallel Reverse Attention Network (PraNet) (Fan et al., 2020) achieved state-of-the-art performance on five different polyp segmentation benchmark by adopting the majority of the network design and techniques from (Chen et al., 2018). They added parallel partial decoder from (Wu et al., 2019) to combine low-level feature maps from the backbone network because the architecture from (Chen et al., 2018) obviously lack of feature sharing for both top-down stream and bottom-up stream. In our network design, we locate intermediate encoder for low-level features and use the output of encoder and previous decoder features for bottom-up stream to incorporate high-level semantic features.

We design UACANet based on the overall architecture from PraNet(Fan et al., 2020) which is a modified version of (Chen et al., 2018). As shown in Figure 1, We add additional encoder network, Parallel Axial Attention encoder (PAA-e) for bottom-up stream and side-out fusion path. This helps to reduce the computational cost for both bottom-up and side-out fusion paths by reducing the number of channels for their input feature maps. Feature maps from three PAA-e modules are both used for side-out fusion path (purple arrow), Parallel Axial Attention decoder (PAA-d) and Uncertainty Augmented Context Attention (UACA). We concatenate three feature maps from PAA-e modules for PAA-d and it predicts the initial saliency map for polyps. Then the feature maps from PAA-e and PAA-d is concatenated for UACA, and the output saliency map from PAA-d is used for context guidance (yellow arrow). We describe detailed information of how UACA incorporates feature maps and context guidance in section 3.3. Then, the output saliency map from UACA is added with previously computed saliency map from the PAA-d. After the first UACA, we concatenate PAA-e feature map and previous UACA feature map for the next UACA (gray and black arrow with concatenation symbol). Also, the saliency map from the previous UACA is used as a context guidance for the next UACA (yellow arrows). The output of UACA is forwarded to the final point-wise convolution and then added with the previous context guidance for the current saliency map. After three consequent UACAs, the final output is computed with bi-linear up-sampling with scale factor of 4 and sigmoid function.

To sum up, the overall architecture shows that the backbone features are encoded with PAA-e, and encoded features are forwarded to PAA-d for initial saliency map which serves as a initial guidance map, which leads UACA to learn a residual saliency map apart from the initial map. This helps the consequent UACA can focus more on uncertain area like boundaries rather than fairly evident region.

We use both binary cross entropy (BCE) loss and intersection over union (IoU) loss. The loss function $\mathcal{L}$ is computed as follows,

where $i\in\mathcal{I}$ refers to a pixel in the output and ground truth, $y$ denotes ground truth and $\hat{y}$ denotes the output. As shown in Figure 1, we add four losses from four prediction from PAA-d and UACA with loss function in Equation 1 (red arrows in Figure 1).

In the era of deep learning, computer vision researchers who study on semantic segmentation and its related tasks like polyp segmentation are dedicated to find a better architecture to extract a fine-grained feature maps which have both high-level semantic information and low-level detail information because usually they are both hard to extract and to combine each other. Self-attention mechanism (Zhang et al., 2019) is one of the major context modules in computer vision but requires heavy computation. Axial attention (Ho et al., 2019) solve this problem with performing non-local operation with respect to the single axis and sequentially connected each operation.

We propose Parallel Axial Attention (PAA) for extracting both global dependencies and local representation. We adopt axial attention strategy by computing non-local operation for both horizontal axis and vertical axis but collocated in parallel. By locating vertical and horizontal attention, both contributes to the final output almost equally compared to the sequential method. While sequentially connected axial attention added trainable positional encoding, we do not use encoding scheme since positional encoding substantially not effective in relatively small scales. Also, we discover that using a parallel connection, element-wise summation is effective to aggregate feature maps rather than concatenation without performance degradation since the same input is used for both horizontal axis and vertical axis and they contributes to the output almost equally by the parallel connection. Also, since a single axis based attention cause unexpected deformation, element-wise summation can help to compensate such artifact. As shown in Figure 2, we compute two non-local operations with input feature map, one for horizontal axis and the other for vertical axis.

We choose to actively use this module for encoder and decoder modules for better representation in terms of globally refined feature. First, we design Parallel Axial Attention encoder (PAA-e) which aggregates the low-level feature maps from the top-down stream which will be used for bottom-up stream. Since U-Net structure use the low-level features without channel reduction, redundant information may hinder the performance and the number of channels are quite large since backbone networks are trained for image classification. Not to lose any detail information while reducing the number of channels, we design PAA-e with Receptive Field Block (RFB) (Liu et al., 2018) strategy. As shown in Figure 3(a), feature map from the backbone network (green box) is forwarded to each receptive field path. We add PAA for additional global refinement for each scale and concatenate the outputs, then forward to the consecutive convolution layers. As shown in Figure 1, the outputs of PAA-e are used for both decoder module and bottom-up stream. Also, we design Parallel Axial Attention decoder (PAA-d) with simple structure yet add additional PAA for final feature aggregation from different level of PAA-e features which are denoted as purple arrows in Figure 1 and Figure 3(b).

While reverse attention (Chen et al., 2018) in both SOD and polyp segmentation does not bring the large margin of performance gain, it was evident that it has shown some qualitatively better results. This phenomenon is highly related to the boundary guided SOD networks (Yang et al., 2017; Su et al., 2019), which show state-of-the-art performances in multiple SOD benchmarks. Boundaries of the object as a extra supervision in SOD networks tends to compensate false negative areas, in other words, object regions with low saliency scores. Reverse attention is thus potentially effective method to bring implicit edge guidance without explicit shape of boundary supervision.

We focus on both saliency and reverse saliency map from the reverse attention and found that usually the boundary region appears where saliency score shows ambiguity. In other words, boundary region is highly related to the saliency score around $0.5$. From this property of saliency map, we assume that both saliency and reverse saliency map has almost equal amount of edge information since reverse saliency map is obtained by simple subtraction from 1. Based on this assumption, extracting ambiguous region from the saliency map as well as foreground and background area would improve attentive methods such as self-attention.

Based on our research, we propose Uncertainty Augmented Context Attention (UACA) module, a novel self-attention mechanism which incorporates uncertain area for rich semantic feature extraction without extra boundary guidance. We denote previously computed input saliency map as m and generate corresponding foreground map $\textbf{m}_{f}$, background map $\textbf{m}_{b}$ and uncertain area map $\textbf{m}_{u}$ as follows,

We compute foreground and background map with max operation in order to disentangle not only from each other but from the uncertain area map since uncertain area map already represents their joint region which makes a redundant information and may diminish the role of uncertainty.

While reverse attention applies explicit channel-wise multiplication to the feature map which resembles Convolutional Block Attention Module (CBAM) (Woo et al., 2018), we compose our context module with a set of non-local operations. Similar to OCR (Yuan et al., 2019) we first compute representative vectors for foreground map, background map and uncertain area map by aggregating the pixel representation with each area map from the input feature map x as follows,

where $i\in\mathcal{I}$ denotes pixels in spatial dimension. We implement Equation 2 with matrix multiplication as shown in Figure 4. Each vector stands for the representative feature vector, so $\textbf{v}_{f}$ represents foreground feature and $\textbf{v}_{u}$ represents the uncertain area. Then we compute the similarity between each representation vector ($\textbf{v}_{f}$, $\textbf{v}_{b}$ and $\textbf{v}_{u}$) and each pixel from the input feature map $\textbf{x}_{i}$ as follows,

Finally, we compute context feature map by weighted summation of representation vector $\textbf{v}_{f}$, $\textbf{v}_{b}$ and $\textbf{v}_{u}$ by similarity score $s_{f}$, $s_{b}$ and $s_{c}$ as follows,

Note that $\psi(\cdot)$, $\phi(\cdot)$, $\omega(\cdot)$ and $\delta(\cdot)$ are point-wise convolution. Each pixel in context feature map, $\textbf{t}_{i}$, can be interpreted as a weighted average of three representation vectors $\textbf{v}_{f}$, $\textbf{v}_{b}$ and $\textbf{v}_{u}$. The context feature map t and input feature map x are concatenated with respect to the channel axis and feed forward to the point-wise convolution for final output feature map as shown in Figure 4.

We describe most of our model architecture description in Section 3 and the number of channels in convolution layers which appear outside of the backbone network is unified as 32 for small model and 256 for large model. We denote the small model with 32 channels as UACANet-S and the large model with 256 channels as UACANet-L. We use Res2Net (Gao et al., 2019) with $26w\times 4s$ settings as a backbone network. Intermediate backbone feature maps are acquired from each stage’s last residual block (green box in Figure 1. For UACANet-L, similar to the DeeplabV3+ (Chen et al., 2017), we modified strides and dilation rates to increase the spatial size of the feature map. We resize images to $352\times 352$ for both training and inference and resize back to its original size. Unlike PraNet (Fan et al., 2020) and other state-of-the-art methods, we adopt additional data augmentation techniques which are fairly popular in semantic segmentation including random flipping on both horizontal and vertical axis, random image scaling from $0.75$ to $1.25$. We conduct random rotation from $0$ to $359$ degrees since colonoscopy images may rotate during examination. We also add additional random dilation and erosion for ground truth label to enhance generalization. We use Adam optimizer (Kingma and Ba, 2014) and the initial learning rate is set to $10^{-4}$ with polynomial learning rate decay (Chen et al., 2017) with factor $(1-(\frac{iter}{iter_{max}})^{0.9})$. Compared to PraNet which trained only 60 epochs, we increase training epoch to 240 since we apply diverse data augmentation. We use Pytorch (Paszke et al., 2019) to implement our model and single Titan RTX GPU for train the model.

Following (Fan et al., 2020), images which have been randomly selected from Kvasir and CVC-ClinicDB is used for training, but we use same training data for fair comparison which has already been extracted from Kvasir and CVC-ClinicDB and it contains 1450 images total. For benchmark dataset, we use five different datasets.

CVC-ClinicDB:

CVC-ClinicDB (Bernal et al., 2015), also known as CVC-612 contains 612 images from 25 colonoscopy videos and selected 29 sequences from them. The size of images is $384\times 288$. 62 images from this dataset are used for test and remaining images are used for training.

CVC-300:

CVC-300 is a test dataset from EndoScene (Vázquez et al., 2017). EndoScene contains 912 images from 44 colonoscopy sequences which were acquired from 36 patients total. Since EndoScene dataset is a combination of CVC-ClinicDB and CVC-300, following D.-P. Fan et al, we use CVC-300 as a test dataset which are 60 samples total.

CVC-ColonDB:

CVC-ColonDB (Bernal et al., 2012) dataset is collected from 15 different colonoscopy sequences and sampled 380 images from these sequences.

ETIS:

ETIS (Silva et al., 2013) dataset contains 196 images which are collected from 34 colonoscopy videos. The size of images is $1225\times 966$ which is the largest among other datasets. Unless polyps in this dataset are vary in size and shape, they are mostly small and hard to find, which makes this dataset more challenging.

Kvasir:

Kvasir (Jha et al., 2020) Kvasir dataset consists of 1000 polyp images and corresponding annotations. Unlike the other datasets, images vary in size, from $332\times 487$ to $1920\times 1072$ and also the size of polyps which appear in the images vary in its size and shape. There are 700 large polyps which is larger than $160\times 160$, 48 small polyps smaller than $64\times 64$ and 323 medium polyps within the large and small scale. 900 images are used for training and 100 images are used for test.

To exhibit the validity of PAA module, we conduct an experiment to evaluate UACANet without PAA modules. We design another model whose specific model architecture is identical to the UACANet except PAA modules (yellow box in Figure 3) is excluded. We choose three metrics to evaluate our methods, mean Dice (mDice), mean intersection over union (mIoU) and mean absolute error (MAE). We choose these two datasets for ablation study since CVC-ClinicDB is sampled for training while ETIS isn’t. As shown in Table 1, UACANet-S with PAA module shows better.

We conduct an experiment to demonstrate the effectiveness of UACA. We leave all the details identical to the UACA except for excluding the uncertainty map ($\textbf{m}_{u}$ in Figure 4), namely Context Attention (CA). We substitute UACA with CA in Figure 1 to make CANet. We design CANet with same small and large version, namely CANet-S and CANet-L respectively. Table 2 shows quantitative results on CANet and UACANet on CVC-ClinicDB and ETIS. We also choose these two datasets for same reason as section 4.3. In terms of performance measure, UACANet consistently outperforms CANet on three major metrics.

We also visualize the output feature map of attention modules in both settings in Figure 5 to verify the effectiveness of uncertain area qualitatively. UACANet substantially produce more precise results than CANet in terms of quality as well. In first and third row, while visualized feature map and the output of CANet-L shows that even though it detects substantial region of polyps, it fails to detect the ambiguous region which may seems hard to discriminate from mucosa, surface of colon. On the other hand, UACANet consistently segment polyp regions precisely. We also visualize uncertain area of UACA ($\textbf{m}_{u}$ in Figure 4), it is easy to recognize that uncertain area is closely related to the boundary of polyps. Especially for the third row, $\textbf{m}_{u}$ also helps to detect the ambiguous area denoted as red boxes. Also, in second row, since ground truth miss another polyp on left bottom corner, UACANet detected the missing polyp.

As mentioned above, we compare our method with previous state-of-the-art methods on five different polyp segmentation benchmarks. Since we train on sampled data from Kvasir and CVC-ClinicDB, even the test results are still unseen, the domain of images from these two datasets is still similar to the training data. Thus, we first demonstrate our results on first two datasets on Table 3. On both datasets, our UACANet achieve the best performance among other methods. Especially, our UACANet-L achieves 92.6% mean Dice on CVC-Clinic DB which is 2.7% improvement over PraNet, the latest state-of-the-art method. In Table 4, we evaluate our method with three completely unseen datasets. We mentioned above that ETIS is the most challenging datasets among other four, nevertheless UACANet-L achieve 76.6% mean Dice which is 13.8% improvement over PraNet.

We also demonstrate qualitative results on five benchmarks of previous state-of-the-art methods and our method (Figure 5). On Kvasir and CVC-ClinicDB (first and second row), since two datasets are similar to the training dataset, all methods are able to segment the location of polyps, but our method show the most similar results compared to the ground truth. On ETIS dataset (third row), which is the most challenging dataset among other four benchmarks, both UACANet-S and UACANet-L are able to detect a small polyp even if the size of the polyp is very small and hard to notice while other methods has failed to detect.