A Dual-stage Polyp Segmentation Network with a Custom Attention-based U-net and Segment Anything Model for Enhanced Mask Prediction

Abstract

Early detection of colorectal cancer through the proper segmentation of polyps in the colonoscopy images is crucial. Polyps' complex morphology and varied appearances are the greatest obstacles for the segmentation approaches. The paper introduces SAMU-Net, a novel deep learning-based dual-stage architecture consisting of a custom attention-based U-Net and modified Segment Anything Model (SAM) for better polyp segmentation. In our model, we used the custom U-Net architecture with an attention mechanism to obtain polyp segmentation masks as the first stage. This mask is then used to generate a bounding box input for the second stage that contains the modified Segment Anything Model. The modified SAM relies on the use of High-Quality token-based architecture along with global and local properties to segment polyps accurately, even in cases where the shapes and sizes of polyps are diverse and the polyps have different appearances. The efficiency of SAMU-Net generated from four different datasets of colonoscopy images was examined. Our process produced a dice coefficient score of 0.94, which is very impressive and has a considerable improvement over the existing state-of-the-art polyp segmentation methods. Moreover, the qualitative results also visualize that the SAMU-Net is capable of accurately segmenting polyps of wide ranges, thus, it is a relevant tool for computer-aided detection as well as the diagnosis of colorectal cancer.