This is the official code for "Generating 3D TOF-MRA Volumes and Segmentation Labels using Generative Adversarial Networks with Mixed Precision" - Subramaniam et al. 2021 - under review
Aim: Generating reasonable looking 3D TOF-MRA patch-label pairs. Utilizing mixed precision for improving the quality and performance of the patch-label pairs
Datasets used: PEGASUS study (72) - 47 (train), 12 (val), 13 (test); 1000Plus study (20) - used for testing purposes only
Folders:
-
data_processing: contains modules used for extraction of patches, and applying skull-stripping
-
evaluation: contains modules related to evaluation methods: Precision and Recall for Distribution (PRD) and Fréchet Inception Distance (FID), and modules related to extracting embedded features for evaluation purposes under the MedicalNet folder under evaluation
-
Images: Images used below for documentation purposes
Files:
-
config.py: configuration parameters used for training different GAN varieties:
- GP model
- SN model
- SN-MP model
- c-SN-MP model
- DPGAN Also contains config parameters to continue training and generate images from saved models; the config parameter changes required to run DPGAN are specified as comments in the file
-
dataset.py: contains custom pytorch dataset class GANDataset as the each patch/label needs to be read from independent zipped files
-
utils.py: contains helper functions for training - inline documentation of the functions available in the code
-
train.py: training code for all the GANs
-
generate.py: generate patches using the saved models; config file can be used to change model and epoch numbers
-
evaluate.py: evaluation with FID and PRD
-
DPGAN_train.py: training code for DPGAN preliminary results in supplementary material
Input: Generator - Noise vector of size 128 Critic - output of generator (fake patch-label pair) or PEGASUS pairs (real patch-label pair).
Output: Generator - Generated patch and label pair Critic - Critic's score used to train the generator and critic
Hyperparameters:
- Noise vector of 128
- Adam optimizer for both discriminator generator for all architectures
- LR: two-time scale update rule (Heusel et al., 2018) Critic = 0.0004; Generator = 0.0002 for all models
- Beta_1 = 0 and beta_2 = 0.9
- Max number of features for critic and discriminator as ndf and ngf respectively:
- 512 for GP, SN, and SN-MP
- 1024 for c-SN-MP
Results:
Sample images of mid-axial slice of patch-label pairs, and the corresponding 3D vessel structure from A: GP model; B: SN model; C: SN-MP model; D: c-SN-MP model; E: Real
Samples of 3D volumes in axial view as GIFs:
GP:
SN:
SN-MP:
c-SN-MP:
Real: