Multi-scale Information Sharing and Selection Network with Boundary Attention for Polyp Segmentation

Recently, deep learning has demonstrated remarkable performance for precise medical image segmentation. UNet [4] achieves the segmentation process through a symmetric encoder-decoder structure. Following this approach, various types of improvement structures have been primarily designed for segmentation research.

The first improvement focuses on skip connections. UNet++ [5] enhances the fusion of multi-scale features by using a series of nested dense skip connections on both the encoder and the decoder. U2-Net [11] defines a nested U-shaped structure and introduces Residual U-blocks, achieving the capture of context information at different scales. The second enhancement involves utilizing different types of backbone. ResUNet [6] introduces ResNet [12] and enhances learning performance by fitting residuals by adding skip connections. ResUNet++ [6] improves ResUNet by integrating efficient components into the UNet structure. R2U-Net [13] combines the advantages of UNet, residual networks, and Recurrent CNNs(RCNNs) to design a recurrent residual convolutional neural network, which can perform the rich feature representation and segment the target successfully. Next is the incorporation of various mechanisms into UNet, such as the attention mechanism in Attention UNet [14]. It incorporates attention gates (AG) into skip connections, allowing the network to emphasize the salient features of the specific target for effective segmentation. Attention UNet++ [15] further enhances UNet++ by adding attention gates (AG) between nested convolution blocks, allowing features at different levels to selectively focus on their respective tasks. ERDUnet [16] addresses the challenge of extracting global contextual features by enhancing UNet, aiming to improve segmentation accuracy while saving parameters.

After that, transformer-based [17] methods have introduced new ideas for medical image segmentation. Medical Transformer (MedT) [18] based on Transformer introduces a gated axial self-attention mechanism and local global training strategy (LoGo) to learn image features automatically. SwinUNet [19] employs a hierarchical Swin Transformer [20] for feature extraction and designs a decoder based on the symmetric Swin Transformer with patch expansion layers to perform upsampling operations. TransUNet [21] combines the strengths of UNet and Transformer, exhibiting robust performance in medical image segmentation. MSCAF-Net [22] adopts the improved Pyramid Vision Transformer (PVTv2) model as its backbone, refining features at each scale and achieving a comprehensive interaction of multi-scale information for accurate segmentation.

Deep learning has been widely applied in polyp segmentation. Brandao et al. [1] are pioneers in utilizing a Fully Convolutional Neural Network [23] (FCN) to segment polyps in colonoscopy. Wicakam et al. [3] propose a fully compressed convolutional network by improving the FCN-8s network, enabling real-time polyp segmentation and improving segmentation performance. Wickstrm et al. [24] propose an advanced architecture that combines FCN-8s and SegNet [25], introducing batch normalization and dropout to improve model generalization and estimate model uncertainty.

However, the up-sampling process of FCN-based methods often results in the loss of detailed information. To solve this problem, UNet++ and ResUNet++ based on UNet serve for polyp segmentation successfully. PolypMixer [26] is a model based on MLP that flexibly handles input scales of polyps and models long-term dependencies for precise and efficient polyp segmentation. However, these approaches often pay insufficient attention to valuable boundary details.

Several approaches have actively explored solutions to this problem. SFANet [8] introduces a boundary-sensitive loss and employs a shared encoder with two mutually constrained decoders to select and aggregate polyp features at different scales. PsiNet [7] employs parallel decoders for joint training of three tasks and proposes a new joint loss function that achieves improved results after retaining more boundary information. However, it still exhibits limitations in handling boundaries due to the integration of multiple tasks. In addition, PraNet [9] utilizes reverse attention [10] to acquire additional boundary cues. CaraNet [27] proposes a contextual axial reverse attention network by adding axial attention to the reverse attention module and using the channel feature pyramid (CFP) module to improve the segmentation performance of small targets. ACSNet [28] leverages local and global contextual features to achieve layer-wise feature complementarity and refine predictions for uncertain regions. However, it extracts limited information from the last feature generated by the encoder for guidance and focusing attention entirely on boundaries might make it challenging for the model to segment the polyp region. CCBANet [29] introduces attention to the foreground, background, and boundary regions to ensure that the model covers the segmentation area with attention as much as possible. MSNet [30] introduces a multi-scale subtraction network and comprehensively supervises features to capture more details and structural cues for accurate polyp localization and edge refinement. BDG-Net [31] proposes a boundary-distribution-guided segmentation network to aggregate high-level features and generate a boundary distribution map, which successfully segments polyps at different scales. DCNet [32] explores candidate objects and additional object-related boundaries by constraining object regions and boundaries. It segments polyps in a coarse-to-fine manner.

The overall architecture of MISNet is illustrated in Figure 1. For the input colonoscopy image $I$, we first use Res2Net [33] to extract features at five different scales. To improve computation efficiency, we use Low-level Fusion Module (LFM) to integrate shallow features extracted from the first two blocks of backbone, and use High-level Fusion Module (HFM) to integrate deep features extracted from the last three blocks of backbone.

MISNet first use Selectively Shared Fusion Module (SSFM) to adaptively aggregate low-level and high-level features and generate an initial map $F_{fuse}$ for the following process. Then, we leverage Parallel Attention Modules (PAM) to emphasize local boundary information. Subsequently, the features of the current layer, the boundary attention from PAM, and the explicit boundary cues mined from low-level features are adaptively combined with Balancing Weight Module (BWM) to refine polyp segmentation in the bottom-up flow of the network. We elaborate each component in below.

With the significance of preserving valuable local fine-grained information for polyp boundary identification, we introduce a Low-level Fusion Module (LFM). This module facilitates the exploitation of low-level features, enhancing the model’s ability to extract precise boundary information.

The low-level fusion module is shown in Figure 2. Specifically, shallow features ${F_{1}}$ and ${F_{2}}$ generated by the first two blocks of the backbone are firstly processed by Receptive Field Block (RFB) [34] to expand the receptive field. Then, the channel number of ${F_{1}}$ and ${F_{2}}$ is adjusted by $1\times 1$ convolutional layers respectively, after which ${F_{2}}$ will be upsampled to the same resolution as that of ${F_{1}}$. Then, these two features are passed through $3\times 3$ convolutional layers then concatenated to generate a aggregated feature. The aggregated feature is then processed with two $3\times 3$ convolution layers and downsampled to match with the the selective shared fusion module.

The fused low-level features ${F_{lf}}$ serve two purposes within the network. Firstly, they contribute as inputs to SSFM, aiding in the sharing and adaptive selection of shallow features, thereby enhancing the accuracy of the initial guidance map. Secondly, these features explicitly guide the boundary refinement in the decoding process, facilitating clarity enhancement in the generation of segmentation results.

For the fusion of deep features, we adopt the Cascaded Partial Decoder [35]. Specifically, the high-level features $\{{F_{i}},i=3,4,5\}$ produced by the last three blocks of the backbone are first expanded in receptive field through RFB. Then, the high-level features are aggregated using the Cascaded Partial Decoder $cpd\left(\cdot\right)$, ${F_{hf}}=cpd({F_{3}},{F_{4}},{F_{5}})$.

As demonstrated by most image recognition research, high-level features capture abstract and high-level semantic information, providing a more global context, whereas low-level features contribute to capturing details and local information. For segmentation tasks, high-level features enable the model to better understand the complex structures and relationships within the image, enhancing the model’s ability to recognize and locate objects, while low-level features provide finer details, assisting the model in more accurately locating and segmenting the boundaries of objects.

Thus, appropriate integration of features at different scales is crucial for providing comprehensive information in image segmentation tasks. In conventional ”U-shaped” network structures, features of the same scale are typically fused through skip connections in a single stream, which is not conducive for information interaction. Motivated by these observations, we propose a novel Selectively Shared Fusion Module (SSFM) to enforce information sharing and active selection between features at different levels.

The structure of SSFM is illustrated in Figure 3. Firstly, the channels of fused low-level feature ${F_{lf}}$ and high-level feature ${F_{hf}}$ are reduced with $1\times 1$ convolution to obtain the squeezed feature $S_{lf}\in\mathbb{R}^{W\times H\times C}$ and $S_{hf}\in\mathbb{R}^{W\times H\times C}$. Then, we employ cross-fusion to share information between $S_{lf}$ and $S_{hf}$. Specifically, $S_{lf}$ and $S_{hf}$ are first cross-concatenated, and then passed through two separate convolutional layers with different kernels to obtain more comprehensive feature representation:

where ${C_{3\times 3}}$ and ${C_{5\times 5}}$ denote a $3\times 3$ convolutional layer and a $5\times 5$ convolutional layer followed by a BN layer and a ReLU activation function. Then we obtain ${S_{lhf\_3}}$ and ${S_{lhf\_4}}$ by repeating cross-fusion:

With two rounds of cross-fusion, the low-level and high-level features complete the circulation of information.

Subsequently, we allow the model to adaptively select the essential information from low-level features and high-level features. We first apply an element-wise summation operation to the features from these two branches:

Global average pooling is then employed to embed the global information. In the fused feature ${S_{lhf}}$, the channel-wise mean value $k$ is computed, $k\in\mathbb{R}^{C}$. For each channel, we have:

Then, we use a $C\times d$ fully connected layer $FC_{C\times d}$ to obtain a relatively dense feature $q\in\mathbb{R}^{d\times 1}$ by reducing the dimension of $k$:

where $\delta$ denotes the sigmoid function and $\beta$ denotes the BN.

where $r$ denotes the reduction ratio, $L$ denotes the minimal value of $d$ ($L$ is 16 in our experiments).

Following that, we compute the attention for low-level feature and high-level feature:

where $G,H\in\mathbb{R}^{C\times d}$ are learnable weights, and $g$, $h$ denote the soft attention for ${S_{lf}}$ and ${S_{hf}}$. $G_{c}\in\mathbb{R}^{1\times d}$, $H_{c}\in\mathbb{R}^{1\times d}$ is the c-th row of $G$ and $H$. $g_{c}$, $h_{c}$ is the c-th element of $g$ and $h$.

Finally, the information from low-level feature and high-level feature are adaptively selected by multiplying ${S_{lf}}$ and ${S_{hf}}$ with attention weight and obtain the initial guidance map $D$:

where $D=[{D_{1}},{D_{2}},...,{D_{C}}]$ , $D_{c}\in\mathbb{R}^{H\times W}$.

With information sharing and adaptive selection between low-level and high-level feature, SSFM can attend to both local details and global contextual information, thus generate more accurate initial guidance map.

The initial guidance map generated through SSFM can provide a coarse area and location of the polyp. Still, it lacks accurate identification and localization of object boundaries. To address this issue, we introduce a parallel attention module to progressively refine the boundary cues of three parallel high-level features during the decoding stage.

The structure of Parallel Attention Module is illustrated in Figure 4. PAM consist of two parallel attention branches: Parallel Axial Reverse Attention (PA-RA) and Parallel Axial Boundary Attention (PA-BA). PA-RA learns the relationship between background and target areas, and PA-BA further address the attention to boundary information. The structure of PA-RA and PA-BA is demonstrated in Figure 5 and 6.

In the bottom-up process of generating segmentation map, the upsampling operation can lead to blurriness. To address this problem, we introduce Balancing Weight Module to adaptively integrate low-level features ${F_{lf}}$, boundary attention $F_{i}^{rb}$ and high-level features $\{{F_{i}},i=3,4,5\}$, as shown in Figure 7. This allows the network to concentrate on both local details and global context.

To preserve local details in the decoding flow, it is crucial to incorporate low-level features. In contrast to ”U-shaped” network that directly decoding low-level features, We employ the BWM module to incorporate low-level features as guidance for recognizing uncertain regions during the generation of segmentation maps for each deep layer. Specifically, the fused low-level features will first be filtered by CBAM [37] (denoted as ${F_{clf}}$) to reduce noise and reinforce the boundary information from the low-level features. The parallel attention features $F_{i}^{rb}$ and the high-level features $\{{F_{i}},i=3,4,5\}$ from the current layer are resized to match the resolution of the ${F_{clf}}$. Then the three parts are concatenated to form the cascaded features.

The cascaded features are first processed by a $3\times 3$ convolutional layer to reduce dimensionality. Then the compressed features are passed through a global average pooling layer followed by element-wise multiplication to highlight contextual information for improving polyp segmentation accuracy. The segmentation map generated from the previous layer is added to the output of BWM to generate segmentation map of current layer.

The layer-by-layer generation of the segmentation map with BWM progressively refines and incorporates features from different levels, contributing to the accurate delineation of object boundaries in the final segmentation result.

We define the loss function as:

where ${\cal L}_{BCE}^{\omega}$ [38] and ${\cal L}_{IoU}^{\omega}$ [38] denote the weighted binary cross entropy (BCE) loss and the weighted IoU loss, respectively.

We found that in the training set, the sizes of polyps exhibit an obvious imbalance. Commonly used BCE loss can address the imbalance in positive sample segmentation. However, it primarily complements pixel-wise aspects at a microscopic level, which neglects the global structure within the image. This leads to certain deficiencies in learning challenging samples. Therefore, we supplement the IoU loss to constrain from the global aspect. Compared with standard BCE loss and IoU loss, ${\cal L}_{BCE}^{\omega}$ and ${\cal L}_{IoU}^{\omega}$ can emphasize the significance of small objects and boundary information by assigning larger weights to them. The effectiveness of these two methods has been confirmed in various studies [10] [39].

We employ deep supervision for the three high-level feature maps $\{{F_{i}},i=3,4,5\}$ and the initial guidance map $F_{fuse}$. These maps are upsampled (denoted as $F_{fuse}^{up}$ , $F_{3}^{up}$ , $F_{4}^{up}$ , $F_{5}^{up}$) to match the same size of the ground truth $G$.

Finally, the whole segmentation framework can be trained in an end-to-end manner with total loss function:

We utilize the same dataset, and training and testing splits with PraNet for polyp segmentation. Specifically, the dataset consists of 900 images from Kvasir-SEG and 550 images from CVC-ClinicDB, totaling 1450 samples. They are randomly divided into 80% for training, 10% for validation, and 10% for testing. We evaluate the performance of the proposed method on five benchmark datasets. Detailed information of benchmarks is described below.

(1) CVC-T (CVC-300) [40]: this dataset contains 60 samples from 44 colonoscopy sequences from 36 patients, where all images are $500\times 574$ in size. All of them are used for testing.

(2) CVC-ClinicDB (CVC-612) [41]: this dataset contains 612 images extracted from 29 different endoscopic video clips with very similar polyp targets. The size of all the images is $384\times 288$. 62 images of this dataset are used for test and the rest of the images are used for training.

(3) CVC-ColonDB [42]: this dataset contains 380 images from 15 different colonoscopy sequences with an image size of $500\times 570$, all of which are used for testing.

(4) ETIS-LaribPolypDB [43]: this dataset contains 196 images collected from 34 colonoscopy videos. The size of all the images is $1225\times 966$. This dataset is currently the most difficult in the field of polyp segmentation. Due to the smaller size and concealed locations of the polyps in this dataset, detection becomes challenging.

(5) Kvasir [44]: this dataset contains 1000 images of polyps with sizes ranging from $332\times 487$ to $1920\times 1072$. The images exhibit variations in the size and shape of the polyps. For training purposes, 900 images from this dataset are utilized, while 100 images are reserved for testing.

To comprehensively assess our model, we employ the following metrics:

(1) MeanDice and MeanIoU are used to measure the similarity between predicted segmentation results and ground truth. Larger values denote higher similarity between the predicted segmentation results and the ground truth.

(2) $F_{\beta}^{\omega}$ ($\omega Fm$) [45] intuitively generalized the F-measure [46] by calculating the accuracy and recall alternately, assigning different weights to various errors in different locations of neighborhood information. This approach is used to correct the ”equal-importance defect” in Dice [47]:

where $P$ denotes precision and $R$ denotes recall. $\beta$ and $\omega$ are used to adjust the relative importance of precision and recall, respectively.

(3) S-measure (${S_{m}}$) [48] is used to evaluate the structural similarity between the prediction segmentation results and the ground truth:

where ${S_{0}}$ and ${S_{r}}$ denote the object-aware and region-aware structural similarity, respectively. $\alpha$ is set to 0.5 in our experiments.

(4) E-measure ($E_{\phi}^{\max}$) [49] is an enhanced alignment method measured from the perspective of global average and local pixel matching with the ground truth:

where ${\phi_{s}}$ denotes the enhanced alignment matrix, which reflects the correlation between the prediction segmentation results and the ground truth.

(5) MAE [50] is the average pixel-wise absolute error, which is used to evaluate the pixel-level accuracy between the ground truth and the predicted segmentation results:

where $G$ denotes the ground truth and $S$ denotes the predicted segmentation results.

We use Res2Net-50 pretrained on ImageNet as the backbone for feature extraction. We use Adam optimizer with the initial learning rate of 1e-5, and weight decay of 1e-5 to optimize our model. The batch size is set to 16.

In training, all images are resized to 352$\times$352. To ensure better model convergence, the total number of training epochs is set to 300. To enhance model stability in later training stages, a Poly learning rate decay strategy is employed, represented as $lr=base\_lr\times{(1-\frac{{epoch}}{{total\_epoch}})^{power}}$ , where the power is set to 0.9. Our proposed network is implemented using PyTorch and trained on a Tesla V100 with 32GB of memory. During training, extensive data augmentation is applied on-the-fly to improve the generalization, including random scaling and cropping, flipping, Gaussian noise, contrast, brightness and sharpness variations. For the ground truth, dilation and erosion operations are applied with kernels vary from 2 to 5.

In order to validate the effectiveness of the proposed model, we compare our method to state-of-the-art segmentation methods: UNet [4], UNet++ [5], SegNet [24], ResUNet [51], ResUNet++ [6], U2Net [11], PraNet [9], C2FNet [52], MSNet [30], CaraNet [27], DoubleU-Net[58], FCBFormer[53], GMSRF-Net[54], HarDNet-MSEG[55], Polyp-PVT[56], TransFuse[57]. To ensure fair comparison, we use official implementations of these comparison models and apply the same data augmentation strategy and dataset splits. All experiments are conducted in the same environment. We report quantitative comparisons on test sets of Kvasir and CVC-ClinicDB in Table I to validate our model’s learning ability, and comparisons on unseen datasets CVC-300, CVC-ColonDB and ETIS in Table II to verify the model’s generalizability.