将拓扑特征融入Mamba模型,提升肠息肉分割边界精度
Topo-VM-UNetV2: Encoding Topology into Vision Mamba UNet for Polyp Segmentation
- 用持久性图生成拓扑注意力图,指导模型关注关键结构
- 在5个公开数据集上显著提升分割性能,边界更准确
- 适合需要精细医学图像分割的研究者和临床应用
卷积神经网络(CNN)和基于Transformer的架构是肠息肉分割的主流深度学习模型。然而,CNN难以建模长距离依赖,而Transformer存在二次计算复杂度。近期,状态空间模型如Mamba因其能有效建模长程交互且保持线性计算复杂度,成为有前景的解决方案。但现有Mamba模型仍难以捕捉拓扑特征(如连通分量、环路、空洞),导致边界分割不准确。为此,我们提出Topo-VM-UNetV2,将拓扑特征编码进最先进的基于Mamba的分割模型VM-UNetV2。方法分为两阶段:第一阶段,使用VM-UNetV2生成概率图(PMs),计算其持久性图,再将每个拓扑特征的持久性值(出生与死亡时间之差)映射到出生位置,经Sigmoid函数转换为注意力权重,生成拓扑注意力图;第二阶段,将拓扑注意力图融入VM-UNetV2的语义与细节融合(SDI)模块,构建拓扑引导的语义与细节融合(Topo-SDI)模块,以优化分割结果。在五个公开肠息肉分割数据集上的大量实验验证了该方法的有效性。代码将公开。
原文摘要 · Abstract (English)
Convolutional neural network (CNN) and Transformer-based architectures are two dominant deep learning models for polyp segmentation. However, CNNs have limited capability for modeling long-range dependencies, while Transformers incur quadratic computational complexity. Recently, State Space Models such as Mamba have been recognized as a promising approach for polyp segmentation because they not only model long-range interactions effectively but also maintain linear computational complexity. However, Mamba-based architectures still struggle to capture topological features (e.g., connected components, loops, voids), leading to inaccurate boundary delineation and polyp segmentation. To address these limitations, we propose a new approach called Topo-VM-UNetV2, which encodes topological features into the Mamba-based state-of-the-art polyp segmentation model, VM-UNetV2. Our method consists of two stages: Stage 1: VM-UNetV2 is used to generate probability maps (PMs) for the training and test images, which are then used to compute topology attention maps. Specifically, we first compute persistence diagrams of the PMs, then we generate persistence score maps by assigning persistence values (i.e., the difference between death and birth times) of each topological feature to its birth location, finally we transform persistence scores into attention weights using the sigmoid function. Stage 2: These topology attention maps are integrated into the semantics and detail infusion (SDI) module of VM-UNetV2 to form a topology-guided semantics and detail infusion (Topo-SDI) module for enhancing the segmentation results. Extensive experiments on five public polyp segmentation datasets demonstrate the effectiveness of our proposed method. The code will be made publicly available.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。