/3D-VAE

Implementation of 3D convolutional conditional variational autoencoder.

Primary LanguageJupyter NotebookMIT LicenseMIT

Conditional VAE for 3D (Brain) Imaging

A Res-Net Style CVAE for 3D Imaging data utilizing lightweight depth-wise separable convolutions. The data set I used for this repository was synthesized from the pre-trained diffusion model from Pinaya et al. (2022). However, any 3D data would work here.

At the moment, the model conditions the latent space of brain volume and ventricle volume. Here, we see an example of interpolating ventricle volume of a specific brain reconstruction from 50 ml to 150 ml:

Training Examples

python train_vae.py -mn test_run --device_index 0 --deep_model --latent_channels 128 --block_widths 1 2 4 8 --ch_multi 64 --dataset_root #path to dataset root#