Package: dmri-rcnn
Angular Super-Resolution in Diffusion MRI with a 3D Recurrent Convolutional AutoencoderThis project enhances the angular resolution of dMRI data through the use of a Recurrent CNN. This codebase is associated with the
following paper. Please cite the paper if you use this model:Angular Super-Resolution in Diffusion MRI with a 3D Recurrent Convolutional Autoencoder[MIDL 2022]Table of contentsInstallationInferenceTrainingDockerSpherical Harmonic BaselineInstallationdMRI-RCNNcan be installed by via pip:pipinstalldmri-rcnnRequirementsdMRI-RCNNusesTensorFlowas the deep learning architecture. To enableGPU usage within TensorFlow, you should ensure the appropriate prerequisites are installed.Listed below are the requirements for this package.tensorflow>=2.6.0numpyeinopsnibabeltqdmInferenceOnce installed, userun_dmri_rcnn.pyto perform inference of new dMRI volumes. Below lists the data requirements to use the script, and the command-line arguments available for inference.DataTo run this script, dMRI data is required in the following format:Context dMRI file. The dMRI data used as context within the model to infer other volumesFile format:NIfTISingle-shell: containing only one b-value.Dimensions:(i, j, k, q_in).(i, j, k)are the spatial dimensions of the dataq_innumber of samples within the q-space dimension. This can either be6,10, or30and will affect which of the trained models is used.Context b-vector file. The corresponding b-vectors for the context dMRI file.File format: text file, whitespace delimited.3rows corresponding to thex, y, zco-ordinates of q-spaceq_incolumns corresponding to the q-space directions sampled.q_inmust either be6,10, or30.Target b-vector file. The corresponding b-vectors for the inferred dMRI data.File format: text file, whitespace delimited.3rows corresponding to thex, y, zco-ordinates of q-spaceq_outcolumns corresponding to the q-space directions sampled.Brain mask file. Binary brain mask file for dMRI data.File format:NIfTIDimensions:(i, j, k). Same spatial dimensions as used in the dMRI data.The script will create the following data:Inferred dMRI file. dMRI volumes inferred from the model as defined by the target b-vectors.File format:NIfTIDimensions:(i, j, k, q_out).q_outnumber of samples within the q-space dimension. This can any number, though using higher numbers will require more GPU memory if using.Command-lineBring up the following help message viarun_dmri_rcnn.py -h:usage: `run_dmri_rcnn.py` [-h] -dmri_in DMRI_IN -bvec_in BVEC_IN -bvec_out BVEC_OUT -mask MASK -dmri_out DMRI_OUT -s {1000,2000,3000} [-m {1,3}] [-c] [-b BATCH_SIZE]

optional arguments:
  -h, --help            show this help message and exit
  -dmri_in DMRI_IN      Context dMRI NIfTI volume. Must be single-shell and contain q_in 3D volumes
  -bvec_in BVEC_IN      Context b-vector text file. Whitespace delimited with 3 rows and q_in columns
  -bvec_out BVEC_OUT    Target b-vector text file. Whitespace delimited with 3 rows and q_out columns
  -mask MASK            Brain mask NIfTI volume. Must have space spatial dimensions as dmri_in.
  -dmri_out DMRI_OUT    Inferred dMRI NIfTI volume. This will contain q_out inferred volumes.
  -s {1000,2000,3000}, --shell {1000,2000,3000}
                        Shell to perform inference with. Must be same shell as context/target dMRI and b-vectors
  -m {1,3}, --model-dim {1,3}
                        Model dimensionality, choose either 1 or 3. Default: 3.
  -c, --combined        Use combined shell model. Currently only applicable with 3D model and 10 q_in.
  -n, --norm            Perform normalisation using 99 percentile of data. Only implemented with --combined flag, and only for q_in = 10
  -b BATCH_SIZE, --batch-size BATCH_SIZE
                        Batch size to run model inference with.N.B.Weights are downloaded and stored within~/.dmri_rcnnby default. To store weights in a different directory, set environment variableDMRI_RCNN_DIR="/your/custom/directory"ExampleThe following example performsb = 1000inference with the 3D dMRI RCNN onHCP data.run_dmri_rcnn.py-dmri_incontext_dmri.nii.gz-bvec_incontext_bvecs-bvec_outtarget_bvecs-maskbrain_mask.nii.gz-dmri_outinferred_dmri.nii.gz-s1000-m3This example would take ~2 minutes to infer 80 volumes on anNVIDIA RTX 3080.To perform inference on data outside of the HCP dataset, use the flags-cand-n. This is currently only implemented for $q_{in} = 10$.TrainingBelow are details on how to train a given model, and the preprocessing steps involved.Data Pre-ProcessingA training dataset is typically too large to fit into memory all at once. To overcome this, this project uses TensorFlow's.tfrecordfile format and thetf.data.DatasetAPI. Therefore training data should be saved in this format before starting. Below is an example on how to do this using thedMRI-RCNNproject.importnumpyasnpfromdmri_rcnn.coreimportiofromdmri_rcnn.core.processingimportsave_tfrecord_data# First load a subject into memorydmri,_=io.load_nifti('/path/to/dmri/data.nii.gz')mask,_=io.load_nifti('/path/to/brain/mask.nii.gz',dtype=np.int8)bvecs=io.load_bvec('/path/to/dmri/bvecs')# bvecs & bvals should be in FSL formatbvals=io.load_bval('/path/to/dmri/bvals')# Optionally crop image datadmri,mask=io.autocrop_dmri(dmri,mask)# .tfrecord format uses a maximum filesize of 2 GiB, therefore for high# resolution dMRI data, the image may need to be split into smaller parts# to do this use the function below. It is recommended to first try to save# each subject as a whole before splitting the image into separate files.dmri_list=io.split_image_to_octants(dmri)mask_list=io.split_image_to_octants(mask)# Now save data in .tfrecord formatsave_tfrecord_data(dmri,bvecs,bvals,mask,'/path/to/saved/data.tfrecord')# Alternatively save the list of image parts if dmri is too largeforiinrange(len(dmri_list)):save_tfrecord_data(dmri_list[i],bvecs,bvals,mask_list[i],'/path/to/saved/data'+str(i)+'.tfrecord')Training a ModelOnce pre-processing is complete, you can then train a model.fromdmri_rcnn.core.weightsimportget_weightsfromdmri_rcnn.core.modelimportget_1d_autoencoder,get_3d_autoencoderfromdmri_rcnn.core.processingimportTrainingProcessor,TrainingProcessorNorm# If we want to fine-tune the model we can load the previously obtained weights.# In this example we'll load the weights for the 3D RCNN trained on the b = 1000# shell and 6 q-space samples per input.weights=get_weights(model_dim=3,shell=1000,q_in=6)# Now we can instantiate the pre-compiled 3D modelmodel=get_3d_autoencoder(weights)# Omit the weights argument to load without pre-trained weights# Instantiate the training processorprocessor=TrainingProcessor(shells=[1000],q_in=6)# If using non-HCP data, the TrainingProcessorNorm should be used instead.processor=TrainingProcessorNorm(shells=[1000],q_in=6)# Important: Here our q_in = 6, and the processor uses a default q_out = 10, therefore our dmri data must# contain at least 16 volumes.# Load dataset mappingtrain_data=processor.load_data(['/path/to/train_data0.tfrecord','/path/to/train_data1.tfrecord'])validation_data=processor.load_data(['/path/to/val_data0.tfrecord'],validation=True)# Begin trainingmodel.fit(train_data,epochs=10,validation_data=validation_data)DockerYou can also usedMRI-RCNNdirectly viaDocker. Both a CPU and GPU version of the project are available.CPUTo usedMRI-RCNNwith the CPU only, use:sudodockerrun-v/absolute/path/to/my/data/directory:/data-it-tmlyon93/dmri-rcnn-cpu:latestGPUTo usedMRI-RCNNwith the GPU, first ensure theappropriate NVIDIA prerequisiteshave been installed. Then use:sudodockerrun--gpusall-v/absolute/path/to/my/data/directory:/data-it-tmlyon93/dmri-rcnn-gpu:latestSpherical Harmonic BaselineTo run the Spherical Harmonic baseline model used in the paper, first ensuredipyis installed. You can installdipydirectly viapipor by installing this project using the following prompt.pipinstalldmri-rcnn[sh]Command-lineBring up the following help message viadmri_sh_baseline.py -h:usage: dMRI Spherical Harmonic Baseline Inference [-h] -dmri_in DMRI_IN -bvec_in BVEC_IN -bvec_out BVEC_OUT -dmri_out DMRI_OUT -s SHELL

optional arguments:
  -h, --help            show this help message and exit
  -dmri_in DMRI_IN      Context dMRI NIfTI volume. Must be single-shell and contain q_in 3D volumes
  -bvec_in BVEC_IN      Context b-vector text file. Whitespace delimited with 3 rows and q_in columns
  -bvec_out BVEC_OUT    Target b-vector text file. Whitespace delimited with 3 rows and q_out columns
  -dmri_out DMRI_OUT    Inferred dMRI NIfTI volume. This will contain q_out inferred volumes.
  -s SHELL, --shell SHELL
                        Shell to perform inference on. Must be same shell as context/target dMRI and b-vecsExampleThe following example performsb = 1000spherical harmonic inference.dmri_sh_baseline.py-dmri_incontext_dmri.nii.gz-bvec_incontext_bvecs-bvec_outtarget_bvecs-dmri_outinferred_dmri.nii.gz-s1000The use or inspect the spherical harmonic model, the code can be found withindmri_rcnn.core.processing.sph_harmonic.RoadmapFuture Additions & Improvements:Plot functionality