Label-efficient Multi-organ Segmentation Method with Diffusion Model
Yongzhi Huang Jinxin Zhu Haseeb Hassan Liyilei Su Jingyu Li Binding Huang
Accurate segmentation of multiple organs in Computed Tomography (CT) images
plays a vital role in computer-aided diagnosis systems. Various
supervised-learning approaches have been proposed recently. However, these
methods heavily depend on a large amount of high-quality labeled data, which is
expensive to obtain in practice. In this study, we present a label-efficient
learning approach using a pre-trained diffusion model for multi-organ
segmentation tasks in CT images. First, a denoising diffusion model was trained
using unlabeled CT data, generating additional two-dimensional (2D) CT images.
Then the pre-trained denoising diffusion network was transferred to the
downstream multi-organ segmentation task, effectively creating a
semi-supervised learning model that requires only a small amount of labeled
data. Furthermore, linear classification and fine-tuning decoder strategies
were employed to enhance the network's segmentation performance. Our generative
model at 256x256 resolution achieves impressive performance in terms of
Fr\'echet inception distance, spatial Fr\'echet inception distance, and
F1-score, with values of 11.32, 46.93, and 73.1\%, respectively. These results
affirm the diffusion model's ability to generate diverse and realistic 2D CT
images. Additionally, our method achieves competitive multi-organ segmentation
performance compared to state-of-the-art methods on the FLARE 2022 dataset,
particularly in limited labeled data scenarios. Remarkably, even with only 1\%
and 10\% labeled data, our method achieves Dice similarity coefficients (DSCs)
of 71.56\% and 78.51\% after fine-tuning, respectively. The method achieves a
DSC score of 51.81\% using just four labeled CT scans. These results
demonstrate the efficacy of our approach in overcoming the limitations of
supervised learning heavily reliant on large-scale labeled data.