Wavelet-Improved Score-Based Generative Model for Medical Imaging

The score-based generative model (SGM) has demonstrated remarkable performance in addressing challenging under-determined inverse problems in medical imaging. However, acquiring high-quality training datasets for these models remains a formidable task, especially in medical image reconstructions. Prevalent noise perturbations or artifacts in low-dose Computed Tomography (CT) or under-sampled Magnetic Resonance Imaging (MRI) hinder the accurate estimation of data distribution gradients, thereby compromising the overall performance of SGMs when trained with these data. To alleviate this issue, we propose a wavelet-improved denoising technique to cooperate with the SGMs, ensuring effective and stable training. Specifically, the proposed method integrates a wavelet sub-network and the standard SGM sub-network into a unified framework, effectively alleviating inaccurate distribution of the data distribution gradient and enhancing the overall stability. The mutual feedback mechanism between the wavelet sub-network and the SGM sub-network empowers the neural network to learn accurate scores even when handling noisy samples. This combination results in a framework that exhibits superior stability during the learning process, leading to the generation of more precise and reliable reconstructed images. During the reconstruction process, we further enhance the robustness and quality of the reconstructed images by incorporating regularization constraint. Our experiments, which encompass various scenarios of low-dose and sparse-view CT, as well as MRI with varying under-sampling rates and masks, demonstrate the effectiveness of the proposed method by significantly enhanced the quality of the reconstructed images. Especially, our method with noisy training samples achieves comparable results to those obtained using clean data. Our code at https://zenodo.org/record/8266123.