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

TL;DR

This paper introduces a novel Hamiltonian Variational Autoencoder-based architecture with discriminative regularization, enhancing tumor segmentation accuracy in data-scarce medical imaging scenarios by generating realistic images and masks.

Contribution

The proposed hybrid HVAE architecture with discriminative regularization improves joint distribution estimation, leading to better segmentation with limited data, especially in 3D medical imaging.

Findings

Effective on BRATS MRI dataset

Successful on HECKTOR PET dataset

Reduces artifacts in generated images

Abstract

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.