HiDiff: Hybrid Diffusion Framework for Medical Image Segmentation

GenAI Diffusion CV MIAS
2024年07月03日
随着深度学习技术的快速发展,医学图像分割得到了显著的进展。现有的基于深度学习的分割模型通常是判别式的,即它们旨在从输入图像到分割掩模的映射。然而,这些判别式方法忽略了底层数据分布和内在类别特征,导致特征空间不稳定。在本文中,我们提出了一种将生成模型的底层数据分布知识与判别式分割方法相结合的方法。为此,我们提出了一种新的医学图像分割混合扩散框架,称为HiDiff,它可以协同利用现有的判别式分割模型和新的生成扩散模型的优势。HiDiff包括两个关键组件:判别式分割器和扩散细化器。首先,我们利用任何传统的训练过的分割模型作为判别式分割器,它可以为扩散细化器提供分割掩模先验。其次,我们提出了一种新的二元伯努利扩散模型(BBDM)作为扩散细化器,通过建模底层数据分布,可以有效、高效、交互式地细化分割掩模。第三,我们以交替协作的方式训练分割器和BBDM,相互提升。对腹部器官、脑肿瘤、息肉和视网膜血管分割数据集进行了广泛的实验,涵盖了四种广泛使用的模态,证明了HiDiff相对于现有的医学分割算法,包括最先进的转换器和扩散基础算法的卓越性能。此外,HiDiff在分割小物体和推广到新数据集方面表现出色。源代码可在 https://github.com/takimailto/HiDiff 上获得。
Medical image segmentation has been significantly advanced with the rapid development of deep learning (DL) techniques. Existing DL-based segmentation models are typically discriminative; i.e., they aim to learn a mapping from the input image to segmentation masks. However, these discriminative methods neglect the underlying data distribution and intrinsic class characteristics, suffering from unstable feature space. In this work, we propose to complement discriminative segmentation methods with the knowledge of underlying data distribution from generative models. To that end, we propose a novel hybrid diffusion framework for medical image segmentation, termed HiDiff, which can synergize the strengths of existing discriminative segmentation models and new generative diffusion models. HiDiff comprises two key components: discriminative segmentor and diffusion refiner. First, we utilize any conventional trained segmentation models as discriminative segmentor, which can provide a segmentation mask prior for diffusion refiner. Second, we propose a novel binary Bernoulli diffusion model (BBDM) as the diffusion refiner, which can effectively, efficiently, and interactively refine the segmentation mask by modeling the underlying data distribution. Third, we train the segmentor and BBDM in an alternate-collaborative manner to mutually boost each other. Extensive experimental results on abdomen organ, brain tumor, polyps, and retinal vessels segmentation datasets, covering four widely-used modalities, demonstrate the superior performance of HiDiff over existing medical segmentation algorithms, including the state-of-the-art transformer- and diffusion-based ones. In addition, HiDiff excels at segmenting small objects and generalizing to new datasets. Source codes are made available at https://github.com/takimailto/HiDiff.
许愿