CNVAE: Anatomically Correct Segmentation via Conditional Generative Modeling

Abstract

Supervised deep learning segmentation models, although unparalleled in performance, are not immune to anatomical mistakes. We develop a segmentation-by-generation paradigm extending the Nouveau VAE (NVAE) generative framework for anatomically correct cardiac segmentation. In particular, we propose a novel conditional NVAE model (CNVAE) for conditional shape generation, and explore its capabilities for image segmentation in two scenarios: (1) direct image-to-mask CNVAE for one-step segmentation; (2) mask-to-mask CNVAE for refinement of segmentation masks produced by third-party segmentation models. Both models are lightweight and return consistently anatomically valid segmentation masks (> 98% success rate, an improvement of 15% above state-of-the-art segmentation models) on the cardiac segmentation/editing tasks of interest, without loss of accuracy in segmented contours.