通过跨多样背景修复与伪掩码精炼实现息肉分割泛化
Generalize Polyp Segmentation via Inpainting across Diverse Backgrounds and Pseudo-Mask Refinement
- CUHK-Shenzhen(香港中文大学(深圳))
- South China Hospital, Medical School, Shenzhen University(深圳大学医学院华南医院)
机构由 AI 辅助整理,请以论文原文为准。
AI总结:
本文结合 Stable Diffusion Inpaint、ControlNet、伪掩码精炼和困难样本选择生成多样息肉训练数据,显著提升分割模型跨数据集泛化性能。
AI中文摘要:
在不同正常背景中修复病变是解决泛化问题的一种潜在方法,而泛化对于息肉分割模型至关重要。然而,在将息肉无缝引入复杂内镜环境的同时生成准确的伪掩码,对现有修复方法仍是挑战。为解决这些问题,我们首先利用预训练的 Stable Diffusion Inpaint 和 ControlNet,提出一种鲁棒生成模型,能够跨不同背景修复息肉。其次,我们利用合成息肉被限制在修复区域内这一先验,构建了修复区域引导的伪掩码精炼网络。我们还提出一种样本选择策略,优先选择对齐良好且困难的合成案例,用于进一步的模型微调。实验表明,我们的修复模型在修复质量上于定性和定量方面均优于基线方法。此外,我们的数据增强策略显著提升了息肉分割模型在外部数据集上的性能,达到或超过该领域全监督训练基准的水平。代码位于 https://github.com/497662892/PolypInpainter。
英文摘要:
Inpainting lesions within different normal backgrounds is a potential method of addressing the generalization problem, which is crucial for polyp segmentation models. However, seamlessly introducing polyps into complex endoscopic environments while simultaneously generating accurate pseudo-masks remains a challenge for current inpainting methods. To address these issues, we first leverage the pre-trained Stable Diffusion Inpaint and ControlNet, to introduce a robust generative model capable of inpainting polyps across different backgrounds. Secondly, we utilize the prior that synthetic polyps are confined to the inpainted region, to establish an inpainted region-guided pseudo-mask refinement network. We also propose a sample selection strategy that prioritizes well-aligned and hard synthetic cases for further model fine-tuning. Experiments demonstrate that our inpainting model outperformed baseline methods both qualitatively and quantitatively in inpainting quality. Moreover, our data augmentation strategy significantly enhances the performance of polyp segmentation models on external datasets, achieving or surpassing the level of fully supervised training benchmarks in that domain. Our code is available at https://github.com/497662892/PolypInpainter.