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

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.

References in corpus (19)

• Very Deep Convolutional Networks for Large-Scale Image Recognition
• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Auto-Encoding Variational Bayes
• Denoising Diffusion Probabilistic Models
• Generative Adversarial Networks
• Attention U-Net: Learning Where to Look for the Pancreas
• Segment Anything in Medical Images
• An overview of deep learning in medical imaging focusing on MRI
• Semi-Supervised Learning with Generative Adversarial Networks
• Deep Unsupervised Clustering with Gaussian Mixture Variational Autoencoders
• Deep Learning Approaches for Data Augmentation in Medical Imaging: A Review
• Tackling the Generative Learning Trilemma with Denoising Diffusion GANs
• Data Augmentation in High Dimensional Low Sample Size Setting Using a Geometry-Based Variational Autoencoder
• Lymphoma segmentation from 3D PET-CT images using a deep evidential network
• BerDiff: Conditional Bernoulli Diffusion Model for Medical Image Segmentation
• Diffusion Models for Medical Image Analysis: A Comprehensive Survey
• A Comprehensive Survey on Data-Efficient GANs in Image Generation
• Swin MAE: Masked Autoencoders for Small Datasets
• End-to-end autoencoding architecture for the simultaneous generation of medical images and corresponding segmentation masks