Anatomically-Controllable Medical Image Generation with Segmentation-Guided Diffusion Models

Nicholas Konz,Yuwen Chen,Haoyu Dong,Maciej A. Mazurowski
2024-06-20
Abstract:Diffusion models have enabled remarkably high-quality medical image generation, yet it is challenging to enforce anatomical constraints in generated images. To this end, we propose a diffusion model-based method that supports anatomically-controllable medical image generation, by following a multi-class anatomical segmentation mask at each sampling step. We additionally introduce a random mask ablation training algorithm to enable conditioning on a selected combination of anatomical constraints while allowing flexibility in other anatomical areas. We compare our method ("SegGuidedDiff") to existing methods on breast MRI and abdominal/neck-to-pelvis CT datasets with a wide range of anatomical objects. Results show that our method reaches a new state-of-the-art in the faithfulness of generated images to input anatomical masks on both datasets, and is on par for general anatomical realism. Finally, our model also enjoys the extra benefit of being able to adjust the anatomical similarity of generated images to real images of choice through interpolation in its latent space. SegGuidedDiff has many applications, including cross-modality translation, and the generation of paired or counterfactual data. Our code is available at <a class="link-external link-https" href="https://github.com/mazurowski-lab/segmentation-guided-diffusion" rel="external noopener nofollow">this https URL</a>.
Image and Video Processing,Computer Vision and Pattern Recognition,Machine Learning
What problem does this paper attempt to address?
The problem that this paper attempts to solve is how to precisely control anatomical structures when generating medical images. Although existing diffusion models are able to generate high - quality medical images, it is very difficult to enforce anatomical constraints during the generation process, resulting in the generated tissues or organs may be anatomically unrealistic. Specifically, standard generative models such as DDPMs (Denoising Diffusion Probability Models) may fail to create anatomically reasonable tissues (as shown in Figure 1), and such anatomical features cannot be precisely customized. To solve this problem, the authors propose a diffusion - model - based method that supports anatomically controllable medical image generation based on multi - class anatomical segmentation masks. This method achieves this goal by following a multi - class anatomical segmentation mask at each sampling step. In addition, the authors also introduce a random mask ablation training algorithm to allow conditional generation on selected combinations of anatomical constraints while maintaining flexibility in other anatomical regions. This method not only improves the fidelity of the generated image to the input anatomical mask but also reaches a level comparable to existing methods in terms of general anatomical realism. Finally, by interpolating in the latent space of the model, the anatomical similarity between the generated image and a specific real image can also be adjusted, thus expanding the application range of the model.