Meta-Polyp: a baseline for efficient Polyp segmentation

Colorectal cancer, arising from polyp growth in the colon or rectum, is a significant health concern with severe implications [1]. Early identification of polyps is vital in improving prognosis and treatment outcomes [1]. Detecting and removing polyps prior to malignancy is crucial for preventing disease progression. Early diagnosis of polyps is paramount in averting the extensive consequences of colorectal cancer [10]. However, current black-box methods lack explanatory transparency, posing challenges in the medical imaging field [1].

Endoscopic image segmentation is a well-studied research field [10]. Early research relied on handcrafted descriptors and a machine learning (ML) classifier that distinguished lesions from the background based on attributes like color, shape, texture, and edges [11]. In recent years, deep learning and convolutional neural networks have led to many new segmentation techniques, such as UNet [3]. The UNet [3] model is considered groundbreaking as it was the first to introduce skip connections in the encoder-decoder architecture for medical segmentation tasks. This innovative technique allows for the combination of both shallow and deep features to improve the accuracy and reliability of the segmentation process. Since its inception, numerous studies and research have been conducted to further enhance the performance of this technique in the segmentation field. As a result, many advancements have been made, and the UNet [3] model has become a crucial tool for medical imaging professionals and researchers in their pursuit of more accurate and efficient segmentation methods [12, 13].

It has been observed that the abstracted architecture of the Transformer, known as MetaFormer [7], plays a crucial role in achieving high levels of performance. This innovative architecture has demonstrated its effectiveness in various applications, particularly in natural language processing (NLP) and image recognition. By leveraging the powerful capabilities of MetaFormer [7], researchers and developers have been able to achieve competitive results and make significant advancements in their respective fields.

We have developed a network that builds on Encoder-Decoder architecture with modifications such as the combination of the ConvFormer and Transformer blocks from the MetaFormer baseline in the encoder stage. In addition, we have proposed the use of the Level-up Upsampling stage and use the Multi-scale Upsampling block to improve the performance of the Up-Sampling layer in the decoder stage. The full architecture is described in the 1. The input $X\in R^{Width\times Height\times 3}$ of the architecture has shape $Width\times Height\times 3$, and the encoder extracts the feature $X_{i}\in R^{\frac{Width}{2^{i+1}}\times\frac{Height}{2^{i+1}}\times F_{i}}$ where $Filters_{i}\in\{64,128,320,512\};i\in\{1,2,3,4\}$ is the filter at step $i$ of the encoder and the decoder stage. Whereas, in the decoder stage, although the feature is decoded 2 times in each step through Convolution Transpose 2D, the feature at the $i$ step is also decoded 4 times by our Multi-scale Upsampling block and then it is merged to the $i+2$ step for enhancing the feature while Upsampling the previous features. From then, the decoder stage with generate the mask with the shape $Width\times Height\times Filters$, then a Convolution layer $1\times 1$ is applied to map the feature map from 64 filters to 1 filter. In the initial two stages, the emphasis lies on acquiring significant local features, which is why the Convformer Encoder is employed. Conversely, in the subsequent stages, the overall information pertaining to the object becomes more crucial. Therefore, the Transformer Encoder is incorporated in the last two stages of the model to capture the global context effectively.

The MetaFormer baseline [7] investigates how existing token mixers can achieve exceptional performance. Rather than inventing new token mixers, our work relies on the MetaFormer architecture. The ConvFormer encoder in the MetaFormer follows a four-step process. The first step involves generating token mixers, achieved through Depthwise Convolution and Separable Convolution for the creation of these mixers.

The equation.1 which is mentioned in the [7], the $Conv_{pw}$ at $i$ is the Convolution pointwise, while $Conv_{dw}$ denotes the Depth-wise Convolution. Then for the next stages, the output of the equation.1 is normalized before the skip connection is applied to the output, and the demonstration follows by the equation.2. The output then is learnable by $W1$ and $W2$ through the Channel MLP layer, which is demonstrated by the equation.3.In addition, in the equation.3, the $\sigma(.)$ denotes the activation function that is used in the ConvFormer block. The use of the Convformer concept in the encoder is helping the model focus on learning the important texture

In the Convformer block, by the idea of MetaFormer [7], we create our module Convformer block (in fig.1), which is different from the Convformer encoder, but have the idea of the previous Convformer block from MetaFormer [7], that can capture the global information, also we include the local feature by the Pointwise Convolution, then the Self-Attention mechanism [14] is applied, in this case, the weights $W\in R^{Width\times Height\times Filter}$ is kept, this weight helps to generate the attention mask, which helps model learn the crucial information from the local information, therefore, the local information is added with the attention mask, and a Channel MLP layer is applied to help model learn the fusion of the local information and the crucial information of the local information.

The Transformer encoder shares a similar concept with the ConvFormer encoder, but with a different token mixer. Instead of using Convolution Block to create the token mixer, the Transformer block uses a classic self-attention mechanism to create an attention mask, which is used as the token mixer. The self-attention mechanism allows the model to attend to different parts of the input sequence and identify relevant features. The attention mask is generated based on the similarity between the input tokens and is used to weight the contribution of each token to the final output. This mechanism allows the model to capture long-range dependencies and contextual information.

In the equation.4, Self Attention is presented as the self-attention mechanism, and the output will follow the equation.5, which is the skip connection applied. Then the output is also learnable by two parameters $W1$ and $W2$ from the Channel MLP layer.

For the Upsampling block, there are 2 stages in this block, the first stage is the Feature extraction stage, and the second stage is the Upsampling stage.

Equation.[6, 7] describes the Multi-scale block with $X\in R^{W\times H\times Filters}$ as the input tensor. In this block, we utilize the Convolution module feature by extracting the input tensor by convolution layers, this one can be used with the others convolution kernel size or the others convolution components, then, the skip connection is used, the $\sigma(.)$ is the activation for the output of the step.

In the experiment, we follow the merged dataset between the ClinicDB dataset and the Kvasir-SEG dataset [15] which is mentioned in the PraNet [4] experiment setup, and also this training set is widely used in various experiments on the later methods. This dataset contains 2 subsets: Kvasir-SEG [15] (900 train samples) and CVC-ClinicDB [16] (550 train samples).

For benchmarking, we choose 4 datasets: Kvasir-SEG [15], ColonDB [17], CVC300 [16] and the Etis [18] dataset for the benchmarking. In those 4 datasets, apart from the Kvasir-SEG [15] dataset, others are out of the distribution datasets.

For research and study purposes, we split the merged dataset into 3 parts, one for training, one for validating, and one for testing; and do experiments on all these 3 parts. The training, validating, and testing make up 60%, 20%, and 20%, respectively, this data split method is used for evaluating our model before benchmarking on various datasets.

While training the model, we also use the Augmentation technique to improve the number of data in the dataset. Augmentation also creates a beneficial impact on the domain of the dataset by making the distribution of the data more complex [19]. To enrich the dataset, we propose some augmentation methods. We use Center Crop [20], Random Rotate [20], GridDistortion [20], Horizontal [20], and Vertical Flip to improve the quantity of the dataset. Moreover, some advanced augmentation methods were applied to improve the distribution of the feature in the data sample.

CutOut augmentation [21] is also applied in our experiments, this advanced method adds the noise, which is the area with all pixel values are 0, to the image randomly and then the mask has also added that area in the same position. Moreover, the CutMix augmentation is used, in this case, the original image is added with a patch from another image, and their corresponding masks are also combined in the same way.

We utilize the Jaccard Loss Function [22] with the following formula

This loss function enables the segmentation process better and can control the model’s performance on the pitch of the tissues. The Jaccard Loss [22] is also known as the IOU metric, with y as the true label and the predicted label being $\hat{y}$, these two labels are demonstrated in the one-hot vector to present classes $C$ being their length. However, to prevent the exploding gradient, there is a smoothing factor called alpha $\alpha$, which helps stabilize the training result.

All architectures were implemented using the Keras framework with TensorFlow as the backend. The input volumes are normalized to [-1, 1]. All networks were trained with described augmentation methods. We used Adam optimization [23] with an initial learning rate of 1e-4. After that, we use the Cosine Annealing learning rate schedule to stable the training process. The smoothing factor alpha $\alpha$ in the Jaccard Loss is $0.7$. We performed our experiment on a single NVIDIA Tesla A100 40GB. The batch size is 128, and it takes around 6 hours to train the entire dataset. Finally, we trained all the models for 300 epochs.

We use IOU and Dice-Coefficient metrics to evaluate our method’s performance. The metrics evaluate the ground truth mask with the predicted mask from the test dataset.

The following is the formula of mIOU [24]:

The Area of overlap is the common area of two predicted masks, and the Area of Union is all of the areas of two masks.

The Dice Coefficient [25] which calculates the division between the common area of two masks and the union area of two masks has the following formula:

For the relevance of our results, we also benchmark our model with previous methods from UNet [3] in 2015 to the newest method which is on the top of the state-of-the-art in early 2023 is FCB-SwinV2 Transformer [26]. Below is the result on the Kvasir-SEG dataset[15] and CVC-300 dataset [16] with results that are competitive on state-of-the-art as shown in Table.I, and Table.II. In addition, to evaluate the deficiency of our approach, we propose to do experiments on the Etis dataset [18], and the Colon-DB dataset [17]. From the experiment, results are observed in Table.III and Table.IV.

Overall, the Table.II, Table.I, and Table.IV, our method achieves the state-of-the art on the Kvasir-SEG [15], CVC-ColonDB dataset [17], and CVC-300 [16] dataset. These results evaluate the effectiveness of our proposed Convformer and Multiscale Upsampling block. However, the weakness of the method is on the small object which is illustrated by the result from the Table.III from the Etis dataset [18].

The fig.4 visualizes masks that are generated from our method when compared with other methods each year. The dataset that is used for this visualization is the Kvasir dataset [32]. From the results, our method is demonstrated to improve the weakness of the previous method to identify the shape of some difficult polyp objects, however, with many polyps in the image, our method also shows a worse result than some methods in the previous year.

The experiments to compare the MetaFormer UNet with the combination of Meta-UNet and our proposed block are done to evaluate the effectiveness of our method. It seems that the Multi-scale Upsampling block has led to significant improvements in the performance of the MetaFormer UNet on the Kvasir dataset, while the Convformer block can help the model learn the local feature and also can learn the captured global feature. The increase in mean Intersection over Union (mIOU) score from 0.877 to 0.921 with the addition of the Multi-scale Upsampling block and Convformer block indicates that this technique has effectively improved the model’s ability to capture fine details and produce more accurate segmentation results. Furthermore, the best-proposed model achieving a 0.921 mIOU score on the Kvasir dataset suggests that the MetaFormer UNet with Multi-scale Upsampling block and Convformer block is capable of achieving state-of-the-art performance in this task. It would be important to consider the computational cost and other practical considerations of the Multi-scale Upsampling block, as well as potential trade-offs in other performance metrics when deciding whether to incorporate it into other models or applications.