A generic plug & play diffusion-based denosing module for medical image segmentation

Guangju Li, Dehu Jin, Yuanjie Zheng, Jia Cui, Wei Gai, Meng Qi

Abstract: Medical image segmentation faces challenges because of the small sample size of the dataset and the fact that images often have noise and artifacts. In recent years, diffusion models have proven very effective in image generation and have been used widely in computer vision. This paper presents a new feature map denoising module (FMD) based on the diffusion model for feature refinement, which is plug-and-play, allowing flexible integration into popular used segmentation networks for seamless end-to-end training. We evaluate the performance of the FMD module on four models, UNet, UNeXt, TransUNet, and IB-TransUNet, by conducting experiments on four datasets. The experimental data analysis shows that adding the FMD module significantly positively impacts the model performance. Furthermore, especially for small lesion areas and minor organs, adding the FMD module allows users to obtain more accurate segmentation results than the original model.