Wgan for medical (MRI) data augmentation

Hi
I’m trying to implement the lesson7-wgan.ipynb notebook for MRI data augmentation
My input data: My input data: 3 channels png images (each channel - different MRI contrast)
Running the ~default parameters:


The generated images starts to get contour of brain but still far from the real images

Real images
image

Generated images (ephocs)
image
Do you have any recommendation how to change the default parameters?
Noise level, normalization, learning rate?
In addition, each one of the 3 channels is in gray scale, but the gan generated images are in RGB scale, is there a way to calibrate this?
Lastly, although I set the number of epochs to 1000, the run stops at epoch 523, without error message, and regardless to the hyperparameter (even when using resnet34 insted of resnert59). Is there a solution to this?

Thanks a lot
Moran