VAE-GANs for Compressive Medical Image Recovery

170070013-Ameya Anjarlekar

203050010-Charan Kumar Reddy Guggulla

204070004-Sanchar Palit

203050013-Vivek Kumar

Introduction

Medical images generated from MRI scan are crucial in assisting doctors for diagnosing the patient.Quality of the reconstructed image has important role in this.

Deep learning models are entering every field and medical image recovery is no exception.

We introduce vae-gans which help for recovery of medical images and perform error analysis on them.

MRI data and Problem Statement

In MRI, undersampled versions of the Fourier transform of the image are obtained as measurements.

Let X be the vectorized form of the image, F be the Fourier Transform matrix and S be the sampling matrix (no of rows < no of columns).

Then y = SFX where y (undersampled measurements) and S are known.

We need to recover X given S and y.

Problem Statement and Motivation

We want to recover the original image from its undersampled version

Commonly done using compressive sensing algorithms such as LASSO or OMP by expressing the image in sparse basis such as Fourier or DCT.

Recently, GANs have shown to be able to represent natural images much more sparsely compared to these sparse basis.

Hence, we use GAN+VAE network for image reconstruction for biomedical images.

Architecture

• The VAE encoder is of 2 layers formed through a sequence of strided convolution operations followed by ReLU activations and batch normalization.
• The VAE decoder also has 2 layers and utilized transpose convolution operations for upsampling.
• The discriminator function of the GAN is an 4-layer CNN.

Training Loss

• The loss functions used in training were based on the mixture of pixel-wise L_2 and adversarial GAN loss.
• The constant η was the weight for a KL-divergence term designed to force the latent code (based on µx, σx for a given batch) to follow a normal distribution.
• The constant λ served as the GAN weight, controlling the influence of the discriminator.

Experimental details and Main Results

The graph shows generator loss/5000 and discriminator loss vs epochs. Experiments are performed on CIFAR-10.

Challenges and proposed solutions

• Non convergence problem while training the GANs
• Not many resources on this area on medical images

​

Proposed solutions:

• Hyperparameter tuning is crucial
• Add loss term for generator
• Changing vae latent prior
• Weaken the discriminator
• Train Generator thrice and Discriminator once in each iteration
• Use different optimizers such as RMSProp, AdaGrad, AdaDelta.
• Use dropouts and add more layers in VAE.
• Use batch normalization

​

Related work

1. [2] had first proposed the use of GANs for applications such as compressed sensing, inpainting and denoising.
2. In their paper, they use a pre-trained generator network as the image prior.
3. The latent variables (z) are optimized to minimize ||A G(z) - Y|| where Y is the undersampled image, A is the sampling matrix and G is the generator network.
4. [7] proposed Deep image prior which uses an untrained CNN network as the image prior.
5. The loss function (||A G(z)- Y||) is reduced by optimizing over the generator weights.

Suggested improvements to the paper

As mentioned in [8], if the training dataset contains full images (instead of undersampled measurements) then we could include the sampling stage as an additional layer in the generator during training.

The trained weights of the first layer can then be used to develop a better sampling strategy rather than using Gaussian or uniform random matrix for sampling.

Conclusion

Compressed sensing can be used to reduce acquisition time in imaging as well as biomedical applications such as MRI or CT.

A VAE+GAN network could be used as a much better prior compared to traditional bases such as Fourier, DCT or Wavelet.

GANs do suffer from non-convergence and instability. So the choice of layers and hyperparameters is very crucial.

​

References

[1] VAE-GANs for Probabilistic Compressive Image Recovery: Uncertainty Analysis

[2]Bora, A., Jalal, A., Price, E., and Dimakis, A. G. Compressed sensing using generative models

[3] Quan, T. M., Nguyen-Duc, T., and Jeong, W.-K. Compressed sensing MRI reconstruction with cyclic loss in generative adversarial networks. arXiv preprint arXiv:1709.00753, 2017.

[4]Mardani, M., Gong, E., Cheng, J. Y., Vasanawala, S. S., Zaharchuk, G., Xing, L., and Pauly, J. M. Deep generative adversarial neural networks for compressive sensing (GANCS) MRI. IEEE Transactions on Medical Imaging, 2018a

[5]Yang, G., Yu, S., Dong, H., Slabaugh, G., Dragotti, P. L., Ye, X., Liu, F., Arridge, S., Keegan, J., Guo, Y., et al. Dagan: Deep de-aliasing generative adversarial networks for fast compressed sensing MRI reconstruction. IEEE Transactions on Medical Imaging, 37(6):1310–1321, 2018

​

References

[6]Compressed Sensing with Deep Image Prior and Learned Regularization

[7] V. Lempitsky, A. Vedaldi and D. Ulyanov, "Deep Image Prior," 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2018, pp. 9446-9454, doi: 10.1109/CVPR.2018.00984.

[8] A Deep Learning Approach to Structured Signal Recovery