release 0.4.0

.github/workflows/ci.yaml

@@ -5,29 +5,44 @@ on:
     branches: [ main, dev ]
 
 jobs:
-  build:
+  build_env_static:
 
     runs-on: self-hosted
 
     steps:
 
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
 
     - name: test cuda
-      shell: bash -l {0}
+      shell: bash -el {0}
       run: |
-        conda activate test_env
-        python -c "import torch;assert torch.cuda.is_available()"
+        source /ritter/share/projects/ci-runner/mambaforge/bin/activate test_env_current
+        python -c "import torch; assert torch.cuda.is_available()"
 
     - name: test install
-      shell: bash -l {0}
+      shell: bash -el {0}
       run: |
-        conda activate test_env
-        pip -v install --force-reinstall .
+        source /ritter/share/projects/ci-runner/mambaforge/bin/activate test_env_current
+        export CXX="g++"
+        export TORCH_CUDA_ARCH_LIST="8.6"
+        python -m pip -v install --upgrade --no-deps --force-reinstall -e .
 
     - name:  test import
-      shell: bash -l {0}
+      shell: bash -el {0}
       run: |
-        conda activate test_env
+        source /ritter/share/projects/ci-runner/mambaforge/bin/activate test_env_current
         python -c "import brain_deform"
 
+    - name: create symlink for test data
+      shell: bash -el {0}
+      run: |
+        ln -s /ritter/share/projects/ci-runner/data/hpc_data /ritter/share/projects/ci-runner/actions-runner/_work/brain-deform/brain-deform/tests/testdata/hpc_data  
+
+    - name:  unit tests
+      shell: bash -el {0}
+      run: |
+        source /ritter/share/projects/ci-runner/mambaforge/bin/activate test_env_current
+        export LD_LIBRARY_PATH=/ritter/share/projects/ci-runner/mambaforge/envs/test_env_current/lib/$LD_LIBRARY_PATH
+        export CUDA_VISIBLE_DEVICES=$(python tests/get_gpu.py)
+        python -m pytest tests/
+

.gitignore

@@ -7,3 +7,5 @@ tests/testdata/nfbs_data
 lightning_logs/
 .vscode/settings.json
 tests/testdata.tar.gz
+tests/testdata/hpc_test_data.tar.gz
+tests/testdata/._hpc_data

README.md

@@ -3,7 +3,8 @@
 This package provides *elastic deformation* methods, based on non-linear image registration.
 Everything runs **on the GPU**, so that it can be used in near real-time for your [PyTorch Lightning](https://www.pytorchlightning.ai/) model. The setup aims to be as minimal as possible, you only need to use our custom [BrainDataModule](brain_deform/lightning.py) which is based on pytoch-lightning's [DataModule](https://pytorch-lightning.readthedocs.io/en/latest/data/datamodule.html) class.
 
-At the moment, only FSL-style [FNIRT](https://fsl.fmrib.ox.ac.uk/fsl/fslwiki/FNIRT) warp coefficient files are supported.
+At the moment, FSL-style [FNIRT](https://fsl.fmrib.ox.ac.uk/fsl/fslwiki/FNIRT) warp coefficient files and warp fields are supported.
+Also, [ANTs](http://stnava.github.io/ANTs/)-style affines and warp fields are supported.
 
 You can
 - apply linear/nonlinear transformations to MNI space
@@ -80,7 +81,8 @@ data:
     # column names in data_table file
     index_column: eid
     t1_column: t1
-    coefs_to_mni_column: coefs_to_mni
+    deformation_to_mni_column: coefs_to_mni
+    deformations_are_coefs: True
     target_column: target
 
     # provide additional image from axilliary set
@@ -105,7 +107,7 @@ data:
 
 The arguments `data_table_path` and `split_path` define which data you plan on using and how it will be split into training, validation and test sets.
 The table referred to by `data_table_path` is a csv file containing at least an index (EID), a T1 FOV image, a B-spline coefficient file, and a target label.
-The split referred to by `split_path` is a json file containing arrays of EIDs for the corresponding set (train, val, test, augmentation - targets for subject-to-subject augmentation methods, auxiliary - for the optional second "auxiliary" image used for contrastive/semisupervised learning).
+The split referred to by `split_path` is a json file containing arrays of EIDs for the corresponding set (*train*, *val*, *test*, *augmentation* - targets for subject-to-subject augmentation methods, *auxiliary* - for the optional second "auxiliary" image used for contrastive/semisupervised learning).
 
 For each batch, the module returns `(batch_main, batch_aux)`, with each batch consisting of `(registered_image, augmented_registered_image, label)`. If `auxiliary_image=False`, `batch_aux` is simply `(None, None, None)`
 
@@ -125,6 +127,36 @@ That's it! Now you can train your model.
 python main.py fit -c config.yaml
 ```
 
+## Notes for ANTs users
+
+### Training with ANTs files
+
+To use this library with ANTs, you need to tweak the configuration file a bit.
+First, we need to disable warping using coefficients by setting `deformations_are_coefs` to false.
+This allows us to use warp fields in the `deformation_to_mni_column`.
+Second, we need to provide the linear affine in `deformation_matrix_column`, since ANTs does not store the affine within the warp field file.
+Finally, we need to enable ANTs compatiblility mode by setting `ants_mode` to true.
+This flag tells the software that it should expect ANTs/ITK-style matrices and deformation fields.
+Affine files typically have the suffix `0GenericAffine.mat` and warp files the suffix `1Warp.nii.gz`.
+
+```yaml
+deformation_to_mni_column: field_to_mni
+deformation_matrix_column: affine_path
+deformations_are_coefs: false
+ants_mode: true
+```
+
+And that's it!
+
+### Alternative training
+
+You may notice that using deformation fields instead of coefficient files slows down execution speed.
+This is due to the fact that coefficient files typically require a few hundred kilobytes, while deformation fields require about thrice the space of a single MNI-space image.
+Unless you really care about the precision of the warp fields, you could consider converting your ANTs registration to FSL.
+The tooling for this is provided in the `tools` directory.
+Choose `ants_to_fsl_coefs.sh` or `ants_to_fsl_coefs_noc3d.sh`, depending on if you have C3D installed.
+The scripts allow you to convert your ANTs affine and warp field into a single FSL-style coefficient file.
+
 ## Running tests (WIP)
 
 Tests rely on the external [testdata](testdata).

brain_deform/ants.py

@@ -0,0 +1,66 @@
+import numpy as np
+import numpy.typing as npt
+from typing import Any
+import SimpleITK as sitk
+from fsl.data.image import Image
+
+
+def read_ants_affine(path: str, src: Image, ref: Image) -> npt.NDArray[Any]:
+    """Reads an ANTs affine registration matrix and converts it to be compatible with FSL."""
+    # This is an adaptation of what c3d_affine_tool does
+    # See http://www.itksnap.org/pmwiki/pmwiki.php?n=Convert3D.Convert3D
+    # specifically source utilities/AffineTransformTool.cxx
+    # and https://www.jiscmail.ac.uk/cgi-bin/webadmin?A2=fsl;c21935f9.1901
+
+    premat = sitk.ReadTransform(path)
+    mov = src
+
+    # Get voxel to physical coords
+    m_ref = ref.getAffine("voxel", "world")
+    m_mov = mov.getAffine("voxel", "world")
+
+    # Get the spacing
+    refspc = ref.nibImage.header["pixdim"][1:4]
+    movspc = mov.nibImage.header["pixdim"][1:4]
+
+    # Determine if we need to flip stuff
+    m_swpref = np.eye(4)
+    if np.linalg.det(m_ref) > 0:
+        m_swpref[0,0] = -1.0
+        m_swpref[0,3] = (ref.shape[0] - 1) * refspc[0]
+
+    m_swpmov = np.eye(4)
+    if np.linalg.det(m_mov) > 0:
+        m_swpmov[0,0] = -1.0
+        m_swpmov[0,3] = (mov.shape[0] - 1) * movspc[0]
+
+    # Convert spacing into diagonal 4x4 matrices
+    a = np.ones(4)
+    a[:3] = refspc
+    b = np.ones(4)
+    b[:3] = movspc
+    refspc = a
+    movspc = b
+    refspc = np.diag(refspc)
+    movspc = np.diag(movspc)
+
+    # Itk formula for offset
+    translation = np.asarray(premat.GetTranslation())
+    center = np.asarray(premat.GetCenter())
+    matrix = np.asarray(premat.GetMatrix()).reshape((3,3))
+    offset = translation + center - matrix @ center
+
+    # Reconstruct matrix from ITK file to 4x4 matrix
+    m_fsl = np.eye(4)
+    m_fsl[:3,:3] = matrix
+    m_fsl[:3,3] = offset
+    m_fsl[2,0] *= -1
+    m_fsl[2,1] *= -1
+    m_fsl[0,2] *= -1
+    m_fsl[1,2] *= -1
+    m_fsl[0,3] *= -1
+    m_fsl[1,3] *= -1
+
+    # Combine all the parts
+    m_out = np.linalg.inv(np.linalg.inv(m_swpmov) @ movspc @ np.linalg.inv(m_mov) @ m_fsl @ m_ref @ np.linalg.inv(refspc) @ np.linalg.inv(m_swpref))
+    return m_out