Package: dmri-pcconv
Spatio-Angular Convolutions for Super-resolution in Diffusion MRIThis project performs angular super-resolution of dMRI data through a parametric continuous convolutional neural network (PCCNN). This codebase is associated with the following paper. Please cite the paper if you use this model:Spatio-Angular Convolutions for Super-resolution in Diffusion MRI[NeurIPS 2023]Table of contentsInstallationTrainingPredictionInstallationdmri-pcconvcan be installed via pip:pipinstalldmri-pcconvRequirementsdmri-pcconvusesPyTorchas the deep learning framework.Listed below are the requirements for this package, these will automatically be installed when installing via pip.torchlightningnpy-patchereinopsnibabelTrainingFollow the instructions below on how to train a PCCNN model for dMRI angular super-resolution.Data PreprocessingThis training pipeline requiresdMRIdata to be saved in.npyformat. Additionally, the angular dimension must be the first dimension within each 4D array. This is because this module usesnpy-patcherto extract training patches at runtime. Below is an example of how to convertNIfTIfiles into.npyusingnibabel.importnumpyasnpfromdmri_pcconv.core.ioimportload_niftidata,_=load_nifti('/path/to/data.nii.gz')# Load dMRI data into memorydata=data.transpose(3,0,1,2)# Move the angular dimension from last to firstnp.save('/path/to/data.npy',data,allow_pickle=False)# Save in npy format. Ensure this is on an SSD.N.B.Patches are read lazily from disk, therefore it ishighlyrecommended to store the training data on an SSD type device, as an HDD will significantly bottleneck the training process when data loading.Additionally,xmaxvalues are required prior to training, due to the lazy runtime of data extraction mentioned above. Below is an example of how to extract and savexmaxvalues for a given subject.fromdmri_pcconv.core.ioimportload_bval,load_niftifromdmri_pcconv.core.normalisationimportTrainingNormaliserbvals=load_bval('path/to/bvals')dmri,_=load_nifti('/path/to/dmri.nii.gz')mask,_=load_nifti('/path/to/brain_mask.nii.gz')xmax_dict=TrainingNormaliser.calculate_xmax(dmri,bvals,mask)TrainingNormaliser.save_xmax('/path/to/xmax.json',xmax_dict)TrainingBelow is an example of how to train thePCCNNmodel, it uses thelightningmodulePCCNNLightningModeland data modulePCCNNDataModule. ThePCCNN-Bv,PCCNN-Sp, andPCCNN-Bv-Spvariants all have their own corresponding model and data module classes.importlightning.pytorchasplfromdmri_pcconv.core.qspaceimportQSpaceInfofromdmri_pcconv.core.modelimportPCCNNLightningModelfromdmri_pcconv.core.trainingimportSubject,PCCNNDataModule# Collect dataset filepathssubj1=Subject('/path/to/first/dmri.npy','/path/to/first/bvecs','/path/to/first/bvals','/path/to/first/brain_mask.nii.gz','/path/to/first/xmax.json')subj2=Subject('/path/to/second/dmri.npy','/path/to/second/bvecs','/path/to/second/bvals','/path/to/second/brain_mask.nii.gz','/path/to/second/xmax.json')subj3=Subject('/path/to/third/dmri.npy','/path/to/third/bvecs','/path/to/third/bvals','/path/to/third/brain_mask.nii.gz','/path/to/third/xmax.json')# Assign Q-space training parametersqinfo=QSpaceInfo(q_in_min=6,# Minimum number of q-space samples each training example will holdq_in_max=20# Maximum number. Training will sample in between this range.q_out=10# Number of output samples per training example.shells=(1000,2000,3000)# Shells used in training and prediction.seed=12345# Optionally provide a random seed for sampling)# Create DataModule instance. This is a thin wrapper around `pl.LightningDataModule`.data_module=PCCNNDataModule(train_subjects=(subj1,subj2),val_subjects=(subj3),qinfo=qinfo,batch_size=16,# Batch size of each devicenum_workers=8,# Number of CPU workers that load the dataseed=12345,# Optionally provide a random seed for sampling)# Load PCCNN lightning modelmodel=PCCNNLightningModel()# Create `pl.Trainer` instance. `PCCNNDataModule` is usable in DDP distributed training strategy.trainer=pl.Trainer(devices=1,accelerator='gpu',epochs=100)# Start trainingtrainer.fit(model,data_module)PredictionHere we outline how to perform prediction after training.importtorchfromdmri_pcconv.core.weightsimportget_weightsfromdmri_pcconv.core.modelimportPCCNNBvLightningModelfromdmri_pcconv.core.predictionimportPCCNNBvPredictionProcessor# Load your pretrained weights## From the original paperweights=torch.load(get_weights('pccnn-bv'))model=PCCNNBvLightningModel()model.load_state_dict(weights)## Or from a pytorch_lightning checkpointmodel=PCCNNBvLightningModel.load_from_checkpoint('/path/to/my/checkpoint.ckpt')# Run predictionpredict=PCCNNBvPredictionProcessor(batch_size=4,num_workers=8,accelerator='gpu')predict.run_subject(model=model,dmri_in='/path/to/context_dmri.nii.gz',bvec_in='/path/to/context_bvecs',bval_in='/path/to/context_bvals',bvec_out='/path/to/target_bvecs',bval_out='/path/to/target_bvals',mask='/path/to/brain_mask.nii.gz',out_fpath='/path/to/predicted_dmri.nii.gz',)N.B.Weights provided by theget_weightsfunction are saved within~/.dmri_pcconvby default. SetDMRI_PCCONV_DIRenvironment variable to override the save directory.