Discriminative Hamiltonian Variational Autoencoder for Accurate Tumor Segmentation in Data-Scarce Regimes

• Main Authors: Kebaili, Aghiles, Lapuyade-Lahorgue, Jérôme, Vera, Pierre, Ruan, Su
• Format: Journal Article
• Language: English
• Published: 17-06-2024
• Subjects: Computer Science - Computer Vision and Pattern Recognition
• DOI: 10.48550/arxiv.2406.11659

Deep learning has gained significant attention in medical image segmentation.
However, the limited availability of annotated training data presents a
challenge to achieving accurate results. In efforts to overcome this challenge,
data augmentation techniques have been proposed. However, the majority of these
approaches primarily focus on image generation. For segmentation tasks,
providing both images and their corresponding target masks is crucial, and the
generation of diverse and realistic samples remains a complex task, especially
when working with limited training datasets. To this end, we propose a new
end-to-end hybrid architecture based on Hamiltonian Variational Autoencoders
(HVAE) and a discriminative regularization to improve the quality of generated
images. Our method provides an accuracte estimation of the joint distribution
of the images and masks, resulting in the generation of realistic medical
images with reduced artifacts and off-distribution instances. As generating 3D
volumes requires substantial time and memory, our architecture operates on a
slice-by-slice basis to segment 3D volumes, capitilizing on the richly
augmented dataset. Experiments conducted on two public datasets, BRATS (MRI
modality) and HECKTOR (PET modality), demonstrate the efficacy of our proposed
method on different medical imaging modalities with limited data.