Enhancing Medical Image Generation with Anatomical Precision: A Multi-Headed VAE-Based Diffusion Model
Keywords: medical image processing, image segmentation, variational auto-encoder, diffusion modelling, controlled image generation
TL;DR: A controlled score-based image synthesis model for medical images, incorporating a multi-headed VAE to preserve style and position while enhancing segmentation capabilities.
Abstract: Score-based image generation models, also known as diffusion models, can generate highly realistic and diverse natural images. However, a common challenge emerges when applying diffusion models to medical image generation and segmentation. While these models excel at producing realistic local textures, they struggle to accurately capture global anatomical priors, such as organ shape and location. Furthermore, the model lacks the capability for controlled recalibration to transform an anatomically unrealistic image into a realistic one. Here we present a new diffusion model where the generated images exhibit both realistic style and anatomically accurate position. Specifically, this is done by guiding the reverse diffusion process with our specially designed multi-headed VAE, which produces the image's disentangled style and position embeddings. We use the position embedding to define a grid deformation function that deforms a simple position prior to a predicted segmentation mask. Then, we apply the same grid deformation on the style embedding for image generation. This alleviates the style embedding from the burden of learning position features, thereby promoting disentangling. Our proposed approach showcases promising performance in controlled image generation across a range of medical image tasks, such as skin lesions and fetal head. Furthermore, our model delivers state-of-the-art segmentation performance.
Primary Area: generative models
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 7146
Loading