Colorectal Polyp Segmentation Combining Pyramid Vision Transformer and Axial Attention

doi:10.3778/j.issn.1002-8331.2203-0110

Abstract

Abstract: To address the challenges of automatic and accurate segmentation of colorectal polyps, a colorectal segmentation network is proposed：PVTA-Net. The network consists of PVTv2, feature pyramid network（FPN）, spatial pyramid（ASPP）, multi-headed self-attentive mechanism（MHSA）, and parallel axial attention module（PAA-d）：extracting feature maps at different scales by PVTv2; using FPN different levels of features are fused to obtain the enhanced feature maps; ASPP is used to aggregate the feature maps obtained by FPN; MHSA is used to obtain the perceptual field containing all input images; and PAA-d is used to generate features with global relationships. Five datasets, including ColonDB, are used to test the comparison between PVTA-Net and mainstream polyp segmentation networks, and the results show that PVTA-Net outperforms the existing mainstream baseline networks. To verify the generalization performance of PVTA-Net, it is used for COVID-19 lung CT image segmentation, and the results show that, PVTA-Net outperforms the mainstream baseline network.

Key words: colorectal polyp segment, COVID-19, PVTv2, axial attention, feature pyramid network

摘要： 针对结直肠息肉自动准确分割存在的挑战，提出了一种结直肠分割网络PVTA-Net。该网络由PVTv2、特征金字塔网络（FPN）、空间金字塔（ASPP）、多头自注意机制（MHSA）以及并行轴向注意模块（PAA-d）组成：通过PVTv2提取不同尺度的特征图；利用FPN不同层次特征进行融合，得到增强的特征图；利用ASPP聚合由FPN得到的特征图；通过MHSA获得包含所有输入图像的感受野；利用PAA-d生成具有全局关系的特征。采用ColonDB等5个数据集对PVTA-Net和主流息肉分割网络进行对比测试，结果表明，PVTA-Net优于现有主流基线网络。为了验证PVTA-Net的泛化性能，将其用于COVID-19肺部CT图像分割，结果表明，PVTA-Net优于主流基线网络。

关键词: 结直肠分割, COVID-19, PVTv2, 轴向注意, 特征金字塔网络

ZHOU Xue, BAI Zhengyao, LU Qianjie, FAN Shenglan. Colorectal Polyp Segmentation Combining Pyramid Vision Transformer and Axial Attention[J]. Computer Engineering and Applications, 2023, 59(11): 222-230.

周雪, 柏正尧, 陆倩杰, 樊圣澜. 融合视觉Transformer和轴向注意的结直肠息肉分割[J]. 计算机工程与应用, 2023, 59(11): 222-230.

References

[1] 刘明浩，马晓冰，张纪娟，等.结直肠息肉生活方式相关危险因素研究进展[J].中国临床医生杂志，2019，47（5）：519-521.
LIU M H，MA X B，ZHANG J J，et al.Advances in the study of lifestyle-related risk factors for colorectal polyps[J].Chinese Journal of Clinicians，2019，47（5）：519-521.
[2] MAMONOV A V，FIGUEIREDO I N，FIGUEIREDO P N，et al.Automated polyp detection in colon capsule endoscopy[J].IEEE Transactions on Medical Imaging，2014，33（7）：1488-1502.
[3] TAJBAKHSH N，GURUDU S R，LIANG J.Automated polyp detection in colonoscopy videos using shape and context information[J].IEEE Transactions on Medical Imaging，2016，35（2）：630-644.
[4] ARMATO S G，PETRICK N A，BRANDAO P，et al.Fully convolutional neural networks for polyp segmentation in colonoscopy[C]//Medical Imaging 2017：Computer-Aided Diagnosis，Orlando，Feb 11-16，2017.
[5] SUN X Z，ZHANG P F，WANG D C，et al.Colorectal polyp segmentation by U-Net with dilation convolution[C]//2019 18th IEEE International Conference on Machine Learning and Applications（ICMLA），Orlando，Dec 16-19，2019.Piscataway，NJ：IEEE，2019：851-858.
[6] FAN D P，JI G P，ZHOU T，et al.PraNet：parallel reverse attention network for polyp segmentation[C]//International Conference on Medical Image Computing and Computer-Assisted Intervention，Shenzhen，Oct 4-8，2020.Berlin，Heidelberg：Springer，2020：263-273.
[7] WANG W H，XIE E Z，LI X，et al.PVTv2：improved baselines with pyramid vision transformer[J].arXiv：2106. 13797，2021.
[8] LIN T Y，DOLLAR P，GIRSHICK R B，et al.Feature pyramid networks for object detection[C]//2017 IEEE Conference on Computer Vision and Pattern Recognition（CVPR），Honolulu，Jul 22-25，2017.Piscataway，NJ：IEEE Computer Society，2017：936-944.
[9] TAEHUN K，LEE H，KIM D J.UACANet：uncertainty augmented context attention for polyp segmentation[J].arXiv：2107.02368，2021.
[10] PETIT O，THOME N，RAMBOUR C，et al.U-Net transformer：self and cross attention for medical image segmentation[J].arXiv：2103.06104，2021.
[11] DOSOVITSKIY A，BEYER L，KOLESNIKOV A，et al.An image is worth 16x16 words：transformers for image recognition at scale[J].arXiv：2010.11929，2021.
[12] YUAN L，CHEN Y P，WANG T，et al.Tokens-to-token vit：training vision transformers from scratch on imagenet[J].arXiv：2101.11986，2021.
[13] CHU X X，TIAN Z，ZHANG B，et al.Conditional positional encodings for vision transformers[J].arXiv：2102. 10882，2021.
[14] WANG W H，XIE E Z，LI X，et al.Pyramid vision transformer：a versatile backbone for dense prediction without convolutions[J].arXiv：2102.12122，2021.
[15] LIU Z，LIN Y T，CAO Y，et al.Swin transformer：hierarchical vision transformer using shifted windows[J].arXiv：2103.14030，2021.
[16] JHA D，PIA H S，RIEGLER M，et al.Kvasir-SEG：a segmented polyp dataset[J].Lecture Notes in Computer Science（LNCS），2020，11962：451-462.
[17] BERNAL J，SANCHEZ F J，FERNAND E G，et al.WM-DOVA maps for accurate polyp highlighting in colonoscopy：validation vs.saliency maps from physicians[J].Computerized Medical Imaging and Graphics，2015，43：99-111.
[18] TAJBAKHSH N，GURUDU S R，LIANG J.Automated polyp detection in colonoscopy videos using shape and context information[J].IEEE Transactions on Medical Imaging，2015，35（2）：630-644.
[19] VAZQUEZ D，BERNAL J，SANCHEZ F J，et al.A benchmark for endoluminal scene segmentation of colonoscopy images[J].arXiv：1612.00799，2016.
[20] SILVA J，HISTACE A，ROMAIN O，et al.Toward embedded detection of polyps in WCE images for early diagnosis of colorectal cancer[J].International Journal of Computer Assisted Radiology and Surgery，2014，9（2）：283-293.
[21] FAN D P，LIU J J，GAO S H，et al.Salient objects in clutter：bringing salient object detection to the foreground[J].arXiv：1803.06091，2018.
[22] FAN D P，GONG C，CAO Y，et al.Enhanced-alignment measure for binary foreground map evaluation[J].arXiv：1805.10421，2018.
[23] FAN D P，CHENG M M，LIU Y，et al.Structure-measure：a new way to evaluate foreground maps[C]//2017 IEEE International Conference on Computer Vision（ICCV），Venice，Oct 22-29，2017.Piscataway，NJ：IEEE，2017：4558-4567.
[24] RONNERBERGER O，FISCHER P，BROX T.U-Net：convolutional networks for biomedical image segmentation[J].arXiv：1505.04597，2015.
[25] ZHOU Z，SIDDIQUEE M M R，TAJBAKHSH N，et al.Unet++：a nested U-net architecture for medical image segmentation[J].arXiv：1807.10165，2018.
[26] FANG Y，CHEN C，YUAN Y，et al.Selective feature aggregation network with area-boundary constraints for polyp segmentation[C]//International Conference on Medi-cal Image Computing and Computer-Assisted Intervention，Shenzhen，Oct 13-17.Berlin，Heidelberg：Springer，2019：302-310.
[27] HUANG C H，HUNG Y W，LIN Y.HarDNet-MSEG：a simple encoder-decoder polyp segmentation neural network that achieves over 0.9 mean dice and 86 FPS[J].arXiv：2101.07172，2021.
[28] YIN Z J，LIANG K M，MA Z Y，et al.Duplex contextual relation network for polyp segmentation[J].arXiv：2103.06725，2021.
[29] ZHANG R，LI G，LI Z，et al.Adaptive context selection for polyp segmentation[C]//International Conference on Medical Image Computing and Computer-Assisted Intervention，Lima，Oct 4-8.Berlin，Heidelberg：Springer，2020：253-262.
[30] PATEL K，BUR A M，WANG G.Enhanced U-Net：a feature enhancement network for polyp segmentation[J].arXiv：2105.00999，2021.
[31] WEI J，HU Y，ZHANG R，et al.Shallow attention network for polyp segmentation[J].arXiv：2108.00882，2021.
[32] OKTAY O，SCHLEMPER J，FOLGOC L L，et al.Attention U-Net：learning where to look for the pancreas[J].arXiv：1804.03999，2018.
[33] SCHLEMPER J，OKTAY O，SCHAAP M，et al.Attention gated networks：learning to leverage salient regions in medical images[J].Medical Image Analysis，2019，53：197-207.
[34] LI X，CHEN H，QI X，et al.H-DenseUNet：hybrid densely connected UNet for liver and tumor segmentation from CT volumes[J].IEEE Transactions on Medical Imaging，2018，37（12）：2663-2674.
[35] FAN D P，ZHOU T，JI G P，et al.Inf-Net：automatic COVID-19 lung infection segmentation from CT images[J].IEEE Transactions on Medical Imaging，2020，39（8）：2626-2637.