Skip to content

khanlab/gcnn_dmri

Repository files navigation

Gauge equivariant CNNs for Diffusion MRI (gcnn_dmri)

PyPI PyPI - Python Version PyPI - License Coookiecutter - Wolt

Authors:

  • @uhussai7 (primary developer)
  • @akhanf

gcnn_dmri incorporates gauge equivariance into cnns designed to process diffusion MRI (dMRI) data. The dmri signal is realized on an antipodally identified sphere, i.e the real projective space . Inspired by Cohen et al. we model this 'half-sphere' as the top of an icosahedron. Interestingly, invoking the correct padding naturally leads us to use the full dihedral group, , to include reflections in addition to rotations of the hexagon, as shown in the image on the right. Here we show the application of such gauge equivariant layers to de-noising Human Connectome Project dMRI data limited to six gradient directions, a problem similar to the work of Tian et al.

Data preparation

  • Select random subject id's for training and testing, one approach is shown in dataHandling/subject_list_generator.py.
  • Similar to Tian et al. we use a mask that avoids CSF. For this we need a grey matter and a white matter mask, which can be made from mri_binarize with the flags --all-wm and --gm respectively.
  • Further steps are shown in niceData.py:
    • make_freesurfer_masks runs the shell script to make the mask mentioned above.
    • make_loss_mask_and_structural finalizes the mask, T1 and T2 images with the correct padding and resolution.
    • make_diffusion creates diffusion volumes with fewer gradient directions, directions are choosen in the sequence of the aquisition and then cut off at desired number.
    • dtifit_on_directions runs dtifit on the new diffusion volumes with fewer directions.
    • We obtain the following folder structure:
        ── <training/testing>
            ├── <subject_id>
            │   ├── diffusion
            │   │   └── <# of gradient directions>
            │   │       ├── diffusion
            │   │       │   ├── bvals
            │   │       │   ├── bvecs
            │   │       │   ├── data.nii.gz
            │   │       │   ├── nodif_brain_mask.nii.gz
            │   │       │   └── S0mean.nii.gz
            │   │       └── dtifit
            │   │           ├── dtifit_< >.nii.gz
            │   ├── freesurfer_mask
            │   │   ├── mask_all_wm.nii.gz
            │   │   └── mask_gm.nii.gz
            │   ├── masks
            │   │   ├── mask_all_wm.nii.gz
            │   │   ├── mask_gm.nii.gz
            │   │   └── mask.nii.gz
            │   └── structural
            │       ├── T1.nii.gz
            │       └── T2.nii.gz
      

Training

Similar to Tian et al. (and references therein) we use a residual network architecture but with the addition of gauge equivariant convolutions on the half icosahedron. The training script with the parameters used is training_script.py. Note that structural mri images (T1.nii.gz and T2.nii.gz) are also used as inputs.

Predictions

Predictions can be performed with the script predicting_script.py. This will create a diffusion volume file, data_network.nii.gz along with bvecs_network and bvals_network, upon which one may perform dtifit. Following are some results of the denoising, the left grey images are fractional anistropy and right colored images are the V1 vector:


Documentation: https://akhanf.github.io/gcnn_dmri

Source Code: https://github.com/akhanf/gcnn_dmri

PyPI: https://pypi.org/project/gcnn_dmri/


Graph-equivariant CNNs for diffusion MRI

Installation

pip install gcnn_dmri

Development

  • Clone this repository
  • Requirements:
  • Create a virtual environment and install the dependencies
poetry install
  • Activate the virtual environment
poetry shell

Testing

pytest

Documentation

The documentation is automatically generated from the content of the docs directory and from the docstrings of the public signatures of the source code. The documentation is updated and published as a Github project page automatically as part each release.

Releasing

Trigger the Draft release workflow (press Run workflow). This will update the changelog & version and create a GitHub release which is in Draft state.

Find the draft release from the GitHub releases and publish it. When a release is published, it'll trigger release workflow which creates PyPI release and deploys updated documentation.

Pre-commit

Pre-commit hooks run all the auto-formatters (e.g. black, isort), linters (e.g. mypy, flake8), and other quality checks to make sure the changeset is in good shape before a commit/push happens.

You can install the hooks with (runs for each commit):

pre-commit install

Or if you want them to run only for each push:

pre-commit install -t pre-push

Or if you want e.g. want to run all checks manually for all files:

pre-commit run --all-files

This project was generated using the wolt-python-package-cookiecutter template.