MEGANet: Multi-Scale Edge-Guided Attention Network for Weak Boundary Polyp Segmentation (WACV 2024)

Nhat-Tan Bui · Dinh-Hieu Hoang · Quang-Thuc Nguyen · Minh-Triet Tran . Ngan Le

arXiv . open access

Introduction

Efficient polyp segmentation in healthcare plays a critical role in enabling early diagnosis of colorectal cancer. However, the segmentation of polyps presents numerous challenges, including the intricate distribution of backgrounds, variations in polyp sizes and shapes, and indistinct boundaries. Defining the boundary between the foreground (i.e. polyp itself) and the background (surrounding tissue) is difficult. To mitigate these challenges, we propose Multi-Scale Edge-Guided Attention Network (MEGANet) tailored specifically for polyp segmentation within colonoscopy images. This network draws inspiration from the fusion of a classical edge detection technique with an attention mechanism. By combining these techniques, MEGANet effectively preserves high-frequency information, notably edges and boundaries, which tend to erode as neural networks deepen. MEGANet is designed as an end-to-end framework, encompassing three key modules: an encoder, which is responsible for capturing and abstracting the features from the input image, a decoder, which focuses on salient features, and the Edge-Guided Attention module (EGA) that employs the Laplacian Operator to accentuate polyp boundaries. Extensive experiments, both qualitative and quantitative, on five benchmark datasets, demonstrate that our MEGANet outperforms other existing SOTA methods under six evaluation metrics.

Note that our model has been renamed. This means EGANet in the code file is now MEGANet in the paper.

Prerequisites

• Pytorch
• Torchvision
• Numpy
• Imageio
• SciPy
• Pillow

Datasets and Trained Models

• Both training and testing datasets are the same as PraNet.
• The MEGANet-ResNet version weights can be downloaded at Google Drive.
• The MEGANet-Res2Net version weights can be downloaded at Google Drive.
• The Res2Net weights can be downloaded at Google Drive.

Usage

Training

python train.py --trainsize "training size" --train_path "path to train dataset" --train_save "path to save checkpoint"

Testing

python test.py --testsize "testing size" --pth_path "path to checkpoint"

Evaluating

python predict_score.py

Predictions

• The pre-computed maps and scores of the MEGANet-ResNet version can be downloaded at Google Drive.
• The pre-computed maps and scores of the MEGANet-Res2Net version can be downloaded at Google Drive.

Citation

@InProceedings{Bui_2024_WACV,
    author    = {Bui, Nhat-Tan and Hoang, Dinh-Hieu and Nguyen, Quang-Thuc and Tran, Minh-Triet and Le, Ngan},
    title     = {MEGANet: Multi-Scale Edge-Guided Attention Network for Weak Boundary Polyp Segmentation},
    booktitle = {Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (WACV)},
    month     = {January},
    year      = {2024},
    pages     = {7985-7994}
}

Acknowledgment

A part of this code is adapted from these previous works: PraNet, CCBANet and UACANet.

FAQ

If you have any questions, please feel free to create an issue on this repository or contact us at [email protected] / [email protected].