[论文解读] Shallow Attention Network for Polyp Segmentation

Jun Wei, Yiwen Hu+4 more|arXiv (Cornell University)|Aug 2, 2021

Advanced Neural Network Applications参考文献 22被引用 9

一句话总结

该论文提出了一种浅层注意力网络（SANet）用于息肉分割，解决了颜色差异、小息肉退化以及前景-背景类别不平衡的问题。通过引入颜色互换以解耦内容与颜色、采用浅层注意力模块增强高分辨率特征，并在推理阶段应用概率校正策略，SANet在五个基准数据集上实现了SOTA性能，推理速度达72 FPS。

ABSTRACT

Accurate polyp segmentation is of great importance for colorectal cancer diagnosis. However, even with a powerful deep neural network, there still exists three big challenges that impede the development of polyp segmentation. (i) Samples collected under different conditions show inconsistent colors, causing the feature distribution gap and overfitting issue; (ii) Due to repeated feature downsampling, small polyps are easily degraded; (iii) Foreground and background pixels are imbalanced, leading to a biased training. To address the above issues, we propose the Shallow Attention Network (SANet) for polyp segmentation. Specifically, to eliminate the effects of color, we design the color exchange operation to decouple the image contents and colors, and force the model to focus more on the target shape and structure. Furthermore, to enhance the segmentation quality of small polyps, we propose the shallow attention module to filter out the background noise of shallow features. Thanks to the high resolution of shallow features, small polyps can be preserved correctly. In addition, to ease the severe pixel imbalance for small polyps, we propose a probability correction strategy (PCS) during the inference phase. Note that even though PCS is not involved in the training phase, it can still work well on a biased model and consistently improve the segmentation performance. Quantitative and qualitative experimental results on five challenging benchmarks confirm that our proposed SANet outperforms previous state-of-the-art methods by a large margin and achieves a speed about 72FPS.

研究动机与目标

• 解决由于内镜条件差异导致的息肉图像颜色分布不一致所引起的过拟合问题。
• 改善深度网络中因重复下采样导致的小息肉特征退化问题。
• 缓解严重前景-背景像素不平衡问题，该问题会 bias 模型训练并降低小息肉的分割性能。
• 开发一种轻量化、高推理速度的模型，保持高精度且无需微调。
• 在多个公开息肉分割基准上实现SOTA性能，且计算开销极低。

提出的方法

• 提出一种颜色互换（CE）操作，将随机图像的颜色迁移到输入图像，解耦内容与颜色，减少基于颜色的虚假相关性。
• 引入一种浅层注意力模块（SAM），利用深层特征从高分辨率浅层特征中过滤背景噪声，保留小息肉的精细细节。
• 在推理阶段采用概率校正策略（PCS），对偏差预测进行重新校准，提升分割性能而不修改训练过程。
• 采用类似U-Net的编码器-解码器结构，结合残差连接和多尺度监督，实现鲁棒的特征学习。
• 应用标准数据增强（随机翻转、旋转、多尺度训练）和端到端训练，使用BCE + Dice损失进行优化。
• 利用高分辨率浅层特征，以保持小息肉分割所必需的空间细节。

实验结果

研究问题

• RQ1 解耦图像内容与颜色分布是否能减少息肉分割模型中的过拟合？
• RQ2 通过注意力机制增强浅层特征是否能改善小息肉的分割效果？
• RQ3 后处理校正策略是否能有效提升偏差模型的分割性能，而无需重新训练？
• RQ4 结合颜色互换、浅层注意力和概率校正是否能在多种数据集上带来一致的性能提升？
• RQ5 所提方法是否能在真实世界息肉分割基准上实现SOTA性能，并达到高推理速度（如72 FPS）？

主要发现

• 在Kvasir数据集上，SANet的mDice达到0.904，mIoU达到0.847，优于全部六种SOTA方法。
• 在CVC-ClinicDB数据集上，SANet的mDice为0.916，mIoU为0.859，显著超过PraNet（0.898和0.840）。
• 在ETIS数据集上，SANet的mDice为0.750，mIoU为0.654，优于此前SOTA方法PraNet（0.628和0.567）。
• 消融实验表明，所有组件——主干网络、SAM、CE和PCS——均不可或缺，完整SANet在ColonDB上达到0.753 mDice，在Kvasir上达到0.904 mDice。
• SANet在RTX 2080Ti上实现72 FPS的实时推理速度，快于PraNet的64 FPS，展现出高效率。
• 视觉对比显示，SANet生成的掩码更清晰、更准确，背景抑制效果更好，尤其在小息肉和难分类型息肉上表现更优。