[论文解读] Prototype Learning for Out-of-Distribution Polyp Segmentation

Nikhil Kumar Tomar, Debesh Jha+1 more|arXiv (Cornell University)|Aug 7, 2023

Colorectal Cancer Screening and DetectionMedicine被引用 3

一句话总结

本文提出PrototypeLab，一种基于原型学习的分割模型，通过利用粗分割掩码生成模块（CMGM）、原型生成模块（PGM）和原型掩码生成模块（PMGM），在多种结肠镜成像模式（WLI、BLI、LCI、FICE）下显著提升了分布外（OOD）息肉分割的鲁棒性。该模型学习到捕捉形状、纹理和颜色特征的类别特定原型，在OOD数据集上实现了DSC ≥ 90%和mIoU ≥ 85%，并具备接近实时的推理速度。

ABSTRACT

Existing polyp segmentation models from colonoscopy images often fail to provide reliable segmentation results on datasets from different centers, limiting their applicability. Our objective in this study is to create a robust and well-generalized segmentation model named PrototypeLab that can assist in polyp segmentation. To achieve this, we incorporate various lighting modes such as White light imaging (WLI), Blue light imaging (BLI), Linked color imaging (LCI), and Flexible spectral imaging color enhancement (FICE) into our new segmentation model, that learns to create prototypes for each class of object present in the images. These prototypes represent the characteristic features of the objects, such as their shape, texture, color. Our model is designed to perform effectively on out-of-distribution (OOD) datasets from multiple centers. We first generate a coarse mask that is used to learn prototypes for the main object class, which are then employed to generate the final segmentation mask. By using prototypes to represent the main class, our approach handles the variability present in the medical images and generalize well to new data since prototype capture the underlying distribution of the data. PrototypeLab offers a promising solution with a dice coefficient of $\geq$ 90\% and mIoU $\geq$ 85\% with a near real-time processing speed for polyp segmentation. It achieved superior performance on OOD datasets compared to 16 state-of-the-art image segmentation architectures, potentially improving clinical outcomes. Codes are available at https://github.com/xxxxx/PrototypeLab.

研究动机与目标

• 解决现有息肉分割模型在不同临床中心获取的分布外（OOD）数据集上泛化能力差的问题。
• 应对小尺寸、扁平、部分可见或伪装性息肉在边界不清晰及成像条件多变情况下面临的挑战。
• 开发一种在多种成像模式（WLI、BLI、LCI、FICE）和多中心数据集上具有良好泛化能力的鲁棒分割模型。
• 通过集成原型学习捕捉数据内在分布特征，提升模型在临床环境中的泛化能力与可靠性。
• 实现在实际计算机辅助诊断（CADx）系统中可部署的高精度与接近实时的推理速度。

提出的方法

• 使用金字塔视觉Transformer（PVT）编码器从输入结肠镜图像中提取多尺度特征。
• 实施粗分割掩码生成模块（CMGM）以生成初始粗分割掩码，用于引导原型学习。
• 设计原型生成模块（PGM），通过捕捉潜在数据分布，为每类学习多个原型，增强特征表示能力。
• 引入编码器特征融合模块（EFFM），实现多尺度特征融合，提升特征丰富度。
• 采用原型掩码生成模块（PMGM），将学习到的原型与解码器特征融合，生成最终分割掩码。
• 采用交叉熵损失与Dice损失联合训练端到端架构，以优化在分布内和分布外数据上的分割性能。

Figure 1 : The overall architecture of the proposed PrototypeLab. The input image is fed to the PVT-encoder to generate a coarse mask, which is combined with encoder features in the prototype generation module to generate various prototypes. Subsequently, the decoder is employed, which produces a fe

实验结果

研究问题

• RQ1 原型学习是否能提升息肉分割模型在不同临床中心获取的分布外（OOD）数据集上的泛化能力？
• RQ2 所提出的PrototypeLab在处理多种成像模式下具有挑战性的息肉类型（如微小、扁平或部分可见病灶）时效果如何？
• RQ3 与基线架构相比，CMGM、PGM和PMGM各组件在提升分割精度与鲁棒性方面的贡献程度如何？
• RQ4 多尺度特征融合（EFFM）与基于原型的推理相结合，是否能提升在噪声或低对比度内镜图像上的性能？
• RQ5 该模型是否能在保持高精度的同时实现接近实时的推理速度，满足临床部署需求？

主要发现

• 在BKAI-IGH OOD数据集上，PrototypeLab的DSC达到0.9243，mIoU达到0.8744，显著优于16种SOTA方法。
• 在Kvasir-SEG数据集上，PrototypeLab的DSC为0.9086，mIoU为0.8544，HD为3.71，DSC和mIoU分别较SSFormer-L高出0.26%和0.44%。
• 在Kvasir-SEG上训练并在多中心PolypGen数据集上测试时，PrototypeLab的DSC为0.7583，mIoU为0.6957，F2为0.7563，优于Polyp-PVT和SSFormer-L。
• 消融实验证实，完整架构（包含CMGM、PGM、PMGM、EFFM及LKDC模块）相比基线模型，DSC提升1.15%，mIoU提升1.45%。
• PrototypeLab以23.85 FPS（≈30 FPS）的速度处理图像，展现出接近实时的推理速度，同时保持高精度。
• 定性结果表明，与DuAT和Polyp-PVT等SOTA基线相比，PrototypeLab在微小、扁平及噪声息肉上显著减少了过分割与欠分割现象。

Figure 2 : Qualitative results of models trained BKAI-IGH and tested on PolypGen. It can be observed that PrototypeLab produces a more accurate segmentation map in all the centers from C1 to C6.