PSTNet: Enhanced Polyp Segmentation with Multi-scale Alignment and Frequency Domain Integration

Deep learning has revolutionized medical image analysis[15, 16, 17, 18, 19, 20, 21, 22], particularly in the domain of polyp detection. The MICCAI 2015 Colonoscopy Video Challenge showcased the superiority of CNN-based methods over traditional hand-crafted approaches, with the top-performing CNN achieving higher precision and recall values[23].

Encoder-decoder architectures, such as U-Net[6], UNet++[7], and ResUNet++[8], have gained popularity in medical image analysis due to their exceptional performance. DDA-Net[24], a two-decoder attention network built upon ResUNet++, demonstrated its effectiveness in polyp segmentation on the Kvasir-SEG dataset, achieving high dice coefficient and mIoU values. Various advancements have been made in CNN-based architectures for polyp image segmentation, including the use of parallel LSTM with DeepLab-v3[25], multi-task segmentation models[26], reverse attention modules[9], and dual-tree wavelet pool CNNs[27]. ADSNet[28] adapted to diverse semantic nuances in ambiguous polyp segmentation areas using a complementary trilateral decoder and continuous attention module. Further developments include real-time polyp segmentation with ColonSNet[29], the application of generative adversarial networks[30], and the use of contextual pixel relations and attention mechanisms in DCRNet[31]. MSNet[2] and EUNet[32] addressed polyp size variability and image noise, while PolySeg Plus[33] utilized active learning to ameliorate data limitations and false positives. DUCK-Net[34] employed residual downsampling and custom convolutional blocks to extract multi-resolution features.

Despite these advancements, CNN-based methods still face challenges in capturing long-range dependencies and efficiently processing high-resolution images, which are crucial for accurate polyp segmentation. While promising results have been achieved, there remains room for improvement in addressing the variability in polyp size, shape, and appearance, as well as dealing with image noise and artifacts.

The transformer architecture, initially designed for machine translation, has been adapted for vision tasks, achieving remarkable performance[35, 36, 37, 38, 39, 40]. The vision transformer (ViT)[41] partitions an image into patches, encodes them, and feeds them sequentially to the transformer encoder, performing image classification using a multi-layer perceptron. Compared to traditional CNNs, ViT models offer advantages such as handling larger input sizes, capturing global dependencies between pixels, and faster convergence with larger batch sizes.

ViT has also been applied to segmentation tasks[42, 43, 44, 45, 46]. To address intensive prediction tasks, pyramidal structures have been incorporated into transformers, as seen in models like PVT and shunted transformer, which utilize hierarchical transformers with multiple stages. Recently, the transformer architecture has been applied to polyp segmentation, with models like Polyp-PVT[10] combining multiscale features from PVTv2 and a CNN-based decoder for accurate polyp segmentation. Other works like Duplex[31], TGANet[47], PraNet[9], and SegT[11] also utilize attention for polyp segmentation tasks, exploring techniques like inverse attention, adversarial training, and edge guidance.

While transformer-based methods have shown promising results in polyp segmentation, they are still relatively new in this field and face challenges such as high computational costs, the need for large amounts of training data, and interpretability concerns. To address the key challenges of polyp segmentation, such as scale inconsistency, low contrast, and unclear boundaries, we propose the Polyp Segmentation Shunted Transformer Network (PSTNet). PSTNet introduces innovative modules to effectively align multi-scale features, integrate frequency information from polyp images, and improve polyp localization and segmentation accuracy.

The SU captures complementary information from adjacent feature layers and emphasizes their differences, providing the decoder with differential feature information. This can be expressed as:

where $\ominus$ denote element-wise subtraction,$\left|\cdot\right|$ means calculate the absolute value. We upsample $\mathbf{X}_{4}$ to the same size as $\mathbf{X}_{3}$, and then pass $f(\cdot)$ , the results are noted as ${\mathbf{X}_{4}^{\prime}}$, ${\mathbf{X}_{3}^{\prime}}$. Then using the SU unit to fuse ${\mathbf{X}_{4}^{\prime}}$ and ${\mathbf{X}_{3}^{\prime}}$ to obtain $S_{1}$ can be expressed as $S_{1}=\left|f(\mathbf{X}_{4}^{\prime})\ominus f(\mathbf{X}_{3}^{\prime})\right|$. Then, $\mathbf{S}_{1}$ is upsampled to the same size as $\mathbf{X}_{2}$, which is recorded as $S_{1}^{\prime}$, $\mathbf{X}_{2}$ is passed through $f(\cdot)$, the result is recorded as ${\mathbf{X}_{2}^{\prime}}$, and use the SU to fuse $\mathbf{S}_{1}^{\prime}$ and ${\mathbf{X}_{2}^{\prime}}$, the result is recorded as $\mathbf{S}_{1}$, which can be expressed as $\mathbf{S}_{2}=\left|f(\mathbf{S}_{1}^{\prime})\ominus f(\mathbf{X}_{2}^{\prime})\right|$. Similarly, the final supplemental fusion result $S_{3}$ can be obtained, which can be expressed as $\mathbf{S}_{3}=\left|f(\mathbf{S}_{2}^{\prime})\ominus f(\mathbf{X}_{1}^{\prime})\right|$.

In the multi-scale feature fusion process, a natural spatial offset exists between the pixel positions of the upper feature $\mathbf{F}_{i}$ and its lower feature $\mathbf{F}_{i-1}$. This offset cannot be eliminated either through concatenation or elemental operations[51]. In Figure 4, $\mathbf{C}_{i}$ is initially upsampled to obtain $\tilde{\mathbf{C}}_{i}$, which is then concatenated with $P$. These concatenated features are subsequently processed through DCNv2[52] with a kernel size of $3\times 3$, resulting in the generation of two offset maps, namely $\Delta_{C}$ and $\Delta_{P}$.

These offset maps play a pivotal role in the calibration of the low-resolution features, $\mathbf{C}_{i}$, and the high-resolution features, $\mathbf{P}$[53], respectively. Once the offset maps are obtained, our feature alignment aggregation can be difined as follows:

where upsample denotes a bilateral interpolation function, while ${u}(\cdot,\cdot)$ represents the alignment function. It can be assumed that the spatial coordinates of each position on the feature map $F$ to be aligned are given as ${(1,1),(1,2),\ldots,(H,W)}$, with the offset map represented by $\Delta\in{R}^{2\times H\times W}$. The $U_{hw}$ is proposed in SFSegNet[54], which is the output of the alignment function ${u}\left(F,\boldsymbol{\Delta}\right)$ and alignment function is defined as follows:

which samples feature on position $\left(h+\Delta_{1hw},w+\Delta_{2hw}\right)$ of $F$, using the bilinear interpolation kernel, where $\Delta_{1hw}$, $\Delta_{2hw}$ indicate the learned 2D transformation offsets for position $(h,w)$.

Our approach underwent evaluation using five challenging colonoscopic polyp image datasets: Kvasir-SEG[58], CVC-Clinic[3], ETIS[59], CVC-ColonDB[60], and EndoScene[1] datasets. The EndoScene dataset is a combination of CVC-612 and CVC-300. In our experiments, we solely utilized theCVC-300 test set, as a portion of the CVC-612 dataset might have been employed for training purposes.

We compared the performance of recent state-of-the-art (SOTA) models for polyp image segmentation on seven different metrics, including UNet++[7], PraNet[9], DCRNet[31], SANet[61], MSNet[2], ADSNet[28], and SegT[11]. For a fair comparison, we used their open-source code and default settings to evaluate the models on the same training and test sets, and generated the result maps.

During the training phase, we incorporated a multi-scale method[48] to handle the size variations among individual polyp images. The AdamW optimizer, a common choice for transformer networks, was employed to optimize the model parameters. Across all experiments, the model was trained for approximately 135 epochs, with a fixed learning rate of $1e-4$ and a decay rate of 0.1 every 45 epochs. For more specific parameter settings, please refer to Table 2. To ensure the stability and effectiveness of the results, we conducted five independent training and testing runs.

Our model was implemented using the PyTorch framework, and all training and testing experiments were conducted on a server equipped with four NVIDIA GeForce RTX 3090 GPUs. During the evaluation phase, images were scaled to a fixed size of $352\times 352$ without applying any post-processing optimization techniques.

We employed six standard evaluation metrics commonly used in image segmentation tasks: mDic, IoU (Intersection over Union), S-measures ($S_{\alpha}$), weighted F-measure $F_{\beta}^{\omega}$, E-measure ($E_{\epsilon}$), and mean absolute error (MAE) to comprehensively evaluate the model performance. For a rigorous and fair comparison, all models were trained, validated, and tested using identical data splits. Additionally, we utilized pre-trained backbones from the ImageNet dataset and incorporated the authors’ provided source code when available. This consistent methodology in metric selection and model evaluation enhances the reliability and accuracy of our comparative analysis.

Learning Capability: We selected two datasets, CVC-ClinicDB and Kvasir-SEG, as benchmarks. CVC-ClinicDB contains 612 images extracted from 31 colonoscopy videos, while Kvasir-SEG consists of 1000 polyp images collected from the polyp category of the Kvasir dataset. Following the practice of PraNet, we used 548 images from CVC-ClinicDB and 900 images from Kvasir-SEG as the training set, with the remaining 64 and 100 images serving as the test set for each dataset, respectively. As shown in Table 1, our model achieved the best results on both datasets, surpassing the state-of-the-art methods. On Kvasir-SEG, our model outperformed ADSNet and SegT by 1.5% and 0.8% in terms of mDice, and 2.4% and 1.5% in terms of mIoU, respectively. Similarly, on CVC-ClinicDB, our model matched the performance of TGANet in mDice and outperformed ADSNet and SegT by 0.7% and 0.5%, respectively, while achieving the highest mIoU of 90.1%, demonstrating its strong learning capability. This can be attributed to our proposed multi-scale feature fusion framework, which enhances the model’s ability to capture comprehensive feature information by combining information from both the RGB and frequency domains. Furthermore, the FCAM incorporates frequency cues from low-level features into global features, enriching the available feature information and aiding in the precise localization of polyp regions. These innovations enable our model to better learn the feature representations of polyps, outperforming other methods.

Generalization Capabilities: To comprehensively evaluate the generalization capability of the model, we employed three datasets that the model had never encountered before: CVC-ColonDB (380 images), CVC-300 (60 images), and ETIS (196 images). It is worth noting that this differs from the validation method used for the ClinicDB and Kvasir-SEG datasets, as the model had no exposure to these datasets during the training process. As shown in Table 3, our proposed method still achieved the best results. On the CVC-ColonDB dataset, our model outperformed ADSNet and SegT by 1.2% and 1.3% in terms of mDice, and 1.8% and 1.6% in terms of mIoU, respectively. Similarly, on the CVC-300 dataset, our model surpassed SANet and SegT by 2.2% and 1.5% in mDice, and 3.2% and 1.9% in mIoU, respectively. Particularly on the most challenging ETIS dataset, where most images contain small polyps, our method outperformed the state-of-the-art ADSNet and SegT by 0.2% and 1.0% in terms of mDice, and 1.1% and 0.6% in terms of mIoU, respectively, demonstrating its exceptional generalization capability. This can be primarily attributed to our proposed FSAM). FSAM employs feature alignment techniques to mitigate noise and utilizes multi-scale subtraction to extract the differences between features at various scales, enabling more precise delineation of polyp boundaries. This effectively addresses the issue of indistinct polyp boundaries, allowing our model to better adapt to unseen datasets. Moreover, the design of the Cross-Perception Module (CPM) also plays a crucial role, as it connects the features generated by all components of PSTNet and integrates the enhanced features from FCAM and FSAM, ensuring the effective utilization of comprehensive feature information and achieving superior generalization performance compared to other methods.

In Figs.6 and7, we present the visualization results obtained from our PSTNet and eight other models, showcasing their performance on challenging examples. Our approach outperforms other methods in several critical aspects: (a) Enhanced Detection of Small Polyps: Our method excels in capturing a more comprehensive range of small polyps. By effectively learning richer multi-scale information, it successfully identifies polyps of varying sizes, thus reducing the occurrence of missed detections. (b) Improved Noise Suppression: Our model demonstrates superior noise suppression capabilities, effectively excluding polyps camouflaged by normal tissue. Given the frequent resemblance of polyps to their background, our approach leverages frequency domain information to facilitate the accurate identification of polyps amidst normal tissue. (c) Enhanced Edge Prediction: The segmentation results generated by our model exhibit remarkable internal consistency and a closer alignment with ground truth data. This reflects our model’s ability to predict edges more effectively, contributing to the overall accuracy and reliability of the segmentation results.

Our baseline (Bas.) is the shunted transformer[48], and we assess the effectiveness of the modules by removing or replacing components from the complete PSTNet and comparing the variants with the standard version. The standard version is denoted as “PSTNet (ST+FCAM+FSAM+CPM)”, where “FCAM”, “FSAM” and “CPM” indicate the usage of the FCAM, FSAM and CPM, respectively.

Effectiveness of FCAM. We investigate the contribution of the FCAM module. We trained a version of ”PSTNet (w/o FCAM)”. Table 4 shows that the method without the FCAM module performs worse on the five datasets compared to the standard PSTNet. Notably, the mDic on the ClinicDB dataset drops by 1.2%, and the mIoU drops by 3.4%. Meanwhile, the absence of FCAM is undeniably associated with a substantial introduction of noise (Fig. 8).

Effectiveness of FSAM. o analyze the effectiveness of FSAM, a version of ”PSTNet (w/o FSAM)” is trained. Table 4 shows that the mIoU metric drops on all datasets after removing the FSAM module, with the most significant drop observed on the CVC-300 dataset (from 84.7% to 80.6%). Meanwhile, the absence of FSAM is undeniably associated with a substantial introduction of noise (Fig. 8).

Effectiveness of CPM. To investigate the contribution of the CPM module to the model, we removed CPM from PSTNet and used the element summing operation as a substitute. This version of the training is called ”PSTNet (w/o CPM)”. Table 4 shows that the standard version of PSTNet has better metrics than PSTNet (w/o CPM) on each dataset, with the most significant difference observed on the ClinicDB dataset ( mIoU: 90.1% vs. 86.9%). Fig. 8 shows the benefits of CPM more intuitively. It is observed that the absence of CPM results in more pronounced errors in detail, and in some cases, even leads to missed detections.

To assess the impact of each component of the loss function on model performance, we performed an ablation study. This study involved evaluating our model using three distinct loss function setups: 1) excluding the combined dice and focal loss (”w/o (dice+focal)”), 2) omitting the weighted binary cross-entropy loss (”w/o (wBCE)”), and 3) employing our ultimate combined loss function (”Final”). The results, presented in Table 5, indicate that our final loss function configuration persistently attains the highest mean Dice coefficient (mDic) and mean Intersection over Union (mIoU) scores across all datasets. This observation underscores the efficacy of our combined loss strategy in enhancing the robustness and accuracy of the segmentation outcomes.