RefineCatDiff: Toward High-Quality Medical Image Segmentation via a Categorical Diffusion Refinement Framework

High-quality medical image segmentation is crucial for clinical applications. However, existing convolutional neural network-based or vision transformer-based segmentation models often struggle to produce accurate masks when dealing with complex images. In contrast, diffusion models are particularly effective at capturing fine-grained features in images, making them well-suited for these challenges. Therefore, RefineCatDiff, a refinement framework that leverages the strengths of diffusion models to achieve high-quality medical image segmentation, is proposed. In this framework, an initial-stage model predicts a coarse mask, which is subsequently refined by a diffusion model to generate a more accurate and detailed mask. Specifically, a categorical distribution-based discrete diffusion model is developed for refinement, which better aligns with the characteristics of segmentation tasks, thereby enhancing its effectiveness. Moreover, the coarse mask is incorporated as prior knowledge into the diffusion process to further improve efficiency. Additionally, the Prior Fusion Guidance Module, the Encoder Coupled Feature Fusion Module, and the Semantic Attention Conditional Module are integrated into the denoising network to enhance guidance of the diffusion process. Extensive experiments on the BTCV, BraTS-2020, and ISIC-2018 datasets demonstrate that RefineCatDiff outperforms several state-of-the-art models in segmentation performance and is highly compatible with various initial-stage models.