integrate niftyreg aladin affine registration

pyproject.toml

@@ -21,6 +21,7 @@ click = ">=8.0.1"
 torch = "^1.13.1"
 torchio = "^0.18.87"
 nibabel = "^5.0.1"
+niftyreg = "^1.5.70rc1"
 
 [tool.poetry.dev-dependencies]
 Pygments = ">=2.10.0"

src/zerodose/__main__.py

@@ -33,7 +33,7 @@ def main() -> None:
             "cuda:7",
         ]
     ),
-    default="cpu",
+    default="cuda:0",
     help="Device to use for inference.",
 )
 
@@ -69,27 +69,46 @@ def main() -> None:
     help="Print verbose output.",
 )
 
+no_registration_option = click.option(
+    "-n",
+    "--no-registration",
+    "no_registration",
+    is_flag=True,
+    default=False,
+    help="""Skip registration to MNI space.
+    Useful if the input images already are in MNI space""",
+)
+
 
 @main.command()
 @mri_option
 @mask_option
 @sbpet_output_option
 @verbose_option
 @device_option
+@no_registration_option
 def syn(
     mri_fnames: Sequence[str],
     mask_fnames: Sequence[str],
     out_fnames: Union[Sequence[str], None] = None,
-    verbose: bool = False,
+    verbose: bool = True,
     device: str = "cuda:0",
+    no_registration: bool = False,
 ) -> None:
     """Synthesize baseline PET images."""
     if out_fnames is None or len(out_fnames) == 0:
         out_fnames = [
             _create_output_fname(mri_fname, suffix="_sb") for mri_fname in mri_fnames
         ]
+
+    do_registration = not no_registration
     synthesize_baselines(
-        mri_fnames, mask_fnames, out_fnames, verbose=verbose, device=device
+        mri_fnames,
+        mask_fnames,
+        out_fnames,
+        verbose=verbose,
+        device=device,
+        do_registration=do_registration,
     )
 
 

src/zerodose/dataset.py

@@ -1,19 +1,24 @@
 """Dataset class for the ZeroDose project."""
 from typing import Any
 from typing import Dict
+from typing import List
 
 import torchio as tio
 
-from zerodose.processing import Pad
+from zerodose import utils
+from zerodose.niftyreg_wrapper import NiftyRegistration
+from zerodose.processing import PadAndCropToMNI
 from zerodose.processing import ToFloat32
 
 
 class SubjectDataset(tio.data.SubjectsDataset):
     """Dataset class for the ZeroDose project."""
 
-    def __init__(self, mri_fnames, mask_fnames, out_fnames):
+    def __init__(self, mri_fnames, mask_fnames, out_fnames, do_registration=True):
         """Initialize the dataset."""
-        transforms = self._get_augmentation_transform_val()
+        transforms = self._get_augmentation_transform_val(
+            do_registration=do_registration
+        )
         subjects = [
             self._make_subject_predict(mr_f, ma_f, ou_f)
             for mr_f, ma_f, ou_f in zip(  # noqa
@@ -41,11 +46,19 @@ def _make_subject_predict(self, mr_path, mask_path, out_fname) -> tio.Subject:
 
         return tio.Subject(subject_dict)
 
-    def _get_augmentation_transform_val(self) -> tio.Compose:
-        return tio.Compose(
+    def _get_augmentation_transform_val(self, do_registration=True) -> tio.Compose:
+        augmentations: List[tio.Transform] = []
+
+        if do_registration:
+            ref = utils.get_mni_template()
+            augmentations.append(NiftyRegistration(floating_image="mr", ref=ref))
+
+        augmentations.extend(
             [
                 tio.transforms.ZNormalization(include=["mr"], masking_method="mask"),
-                Pad(include=["mr", "mask"]),
+                PadAndCropToMNI(include=["mr", "mask"]),
                 ToFloat32(include=["mr"]),
             ]
         )
+
+        return tio.Compose(augmentations)

src/zerodose/niftyreg_wrapper.py

@@ -0,0 +1,180 @@
+"""Niftyreg wrapper for torchio."""
+import os
+import tempfile
+
+import nibabel as nib
+import niftyreg  # type: ignore
+import numpy as np
+import torch
+from torchio import Subject
+from torchio.data import io
+from torchio.transforms import SpatialTransform
+
+
+def _save_matrix_nifty(affine_mat, file_name):
+    s = ""
+    for i in range(affine_mat.shape[0]):
+        for j in range(affine_mat.shape[1]):
+            affine_mat[i, j] = affine_mat[i, j].item()
+            s += str(affine_mat[i, j]) + " "
+        s = s[:-1]
+        s += "\n"
+    with open(file_name, "w") as f:
+        f.write(s)
+
+
+def _read_matrix_nifty(file_name):
+    with open(file_name) as f:
+        s = f.read()
+    s = s.split("\n")
+    s = s[:-1]
+    s = [i.split(" ") for i in s]
+    s = [[float(j) for j in i] for i in s]
+    s = np.array(s)
+    return s
+
+
+def _register_mri_to_mni(mri_fname, ref):
+
+    temp_aff = tempfile.NamedTemporaryFile(delete=False)
+
+    out_mri_fname = os.path.dirname(mri_fname) + "/out_mri1.nii.gz"
+    mni_template = ref
+    niftyreg.main(
+        [
+            "aladin",
+            "-flo",
+            mri_fname,
+            "-ref",
+            mni_template,
+            "-res",
+            out_mri_fname,
+            "-aff",
+            temp_aff.name,
+            "-speeeeed",
+        ]
+    )
+    affine_mat = _read_matrix_nifty(temp_aff.name)
+    temp_aff.close()
+    os.remove(temp_aff.name)
+    return affine_mat
+
+
+def _nifty_reg_resample(ref_path, flo_img, affine_mat):
+
+    temp_flo = tempfile.NamedTemporaryFile(delete=False, suffix=".nii.gz")
+    temp_res = tempfile.NamedTemporaryFile(delete=False, suffix=".nii.gz")
+    temp_aff = tempfile.NamedTemporaryFile(delete=False)
+
+    io.write_image(flo_img.tensor, flo_img.affine, temp_flo.name)
+    _save_matrix_nifty(affine_mat, temp_aff.name)
+
+    niftyreg.main(
+        [
+            "resample",
+            "-ref",
+            ref_path,
+            "-flo",
+            temp_flo.name,
+            "-res",
+            temp_res.name,
+            "-aff",
+            temp_aff.name,
+        ]
+    )
+
+    # After the other process has finished writing to the files, read from them
+    res = nib.load(temp_res.name)
+    data = res.get_fdata()
+    affine = res.affine
+
+    temp_flo.close()
+    temp_res.close()
+    temp_aff.close()
+
+    os.remove(temp_flo.name)
+    os.remove(temp_res.name)
+    os.remove(temp_aff.name)
+
+    return data, affine
+
+
+class NiftyRegistration(SpatialTransform):
+    """Nifty Registration for torchio."""
+
+    def __init__(
+        self,
+        floating_image=None,
+        ref=None,
+        **kwargs,
+    ):
+        """Initialize the niftyreg registration transform."""
+        self.floating_image = floating_image
+        super().__init__(**kwargs)
+        self.ref = ref
+
+    def apply_transform(self, subject: Subject) -> Subject:
+        """Apply the registration and coregistration."""
+        temp_flo = tempfile.NamedTemporaryFile(delete=False, suffix=".nii.gz")
+
+        io.write_image(
+            subject[self.floating_image].tensor,
+            subject[self.floating_image].affine,
+            temp_flo.name,
+        )
+
+        affine = _register_mri_to_mni(temp_flo.name, self.ref)
+
+        temp_flo.close()
+        os.remove(temp_flo.name)
+        inverse_ref = subject[self.floating_image].path
+        transformer = NiftyResample(
+            ref=self.ref, affine=affine, inverse_ref=inverse_ref
+        )
+        transformed = transformer(subject)
+        return transformed  # type: ignore
+
+
+class NiftyResample(SpatialTransform):
+    """Nifty resample (coreregistration)."""
+
+    def __init__(self, affine, ref=None, is_inverse=False, inverse_ref=None, **kwargs):
+        """Initialize the nifty resample."""
+        self.affine = affine
+        self.ref = ref
+        self.is_inverse = is_inverse
+        self.inverse_ref = inverse_ref
+        super().__init__(**kwargs)
+        self.args_names = ("affine", "ref", "is_inverse", "inverse_ref")
+
+    def apply_transform(self, subject: Subject) -> Subject:
+        """Apply the transform."""
+        for image in self.get_images(subject):
+            if self.is_inverse:
+                if image.path is not None:
+                    ref = image.path
+                else:
+                    ref = self.inverse_ref
+                _apply_niftyreg_resample(image, ref, np.linalg.inv(self.affine))
+            else:
+                _apply_niftyreg_resample(image, self.ref, self.affine)
+        return subject
+
+    @staticmethod
+    def is_invertible():
+        """Whether the transform is invertible."""
+        return True
+
+    def inverse(self):
+        """Return the inverse resample."""
+        return NiftyResample(
+            affine=self.affine, is_inverse=True, inverse_ref=self.inverse_ref
+        )
+
+
+def _apply_niftyreg_resample(image, ref_path, affine_mat):
+    data, affine = _nifty_reg_resample(ref_path, image, affine_mat)
+    image.affine = affine
+    data = data.copy()
+    data = torch.as_tensor(data)
+    image.set_data(data.unsqueeze(0))

src/zerodose/paths.py

@@ -5,4 +5,9 @@
 
 
 # please refer to the readme on where to get the parameters. Save them in this folder:
-folder_with_parameter_files = os.path.join(os.path.expanduser("~"), ".zerodose_params")
+folder_with_parameter_files = os.path.join(
+    os.path.expanduser("~"), ".zerodose_data", "model_params"
+)
+folder_with_templates = os.path.join(
+    os.path.expanduser("~"), ".zerodose_data", "templates"
+)

src/zerodose/processing.py

@@ -38,36 +38,44 @@ def _crop_192_to_mni(arr: torch.Tensor) -> torch.Tensor:
     return parr
 
 
-def postprocess(arr: torch.Tensor) -> torch.Tensor:
-    """Postprocess the 192x192x192 image to MNI."""
-    return _crop_192_to_mni(arr)
-
-
-class Pad(SpatialTransform):
+class PadAndCropToMNI(SpatialTransform):
     """Pad the MNI image to 192x192x192."""
 
-    def __init__(self, **kwargs) -> None:
+    def __init__(self, is_inverse=False, **kwargs) -> None:
         """Initialize the transform."""
+        self.is_inverse = is_inverse
         super().__init__(**kwargs)
 
     def apply_transform(self, subject: Subject) -> Subject:
         """Apply the transform to the subject."""
         for image in self.get_images(subject):
-            _pad(image)
+            if self.is_inverse:
+                _pad_inv(image)
+            else:
+                _pad(image)
 
         return subject
 
     @staticmethod
     def is_invertible() -> bool:
         """Return whether the transform is invertible."""
-        return False
+        return True
+
+    def inverse(self):
+        """Returns the inverse transform."""
+        return PadAndCropToMNI(is_inverse=True)
 
 
 def _pad(image: tio.Image) -> None:
     data = processing._crop_mni_to_192(image.data)
     image.set_data(data)
 
 
+def _pad_inv(image: tio.Image) -> None:
+    data = processing._crop_192_to_mni(image.data)
+    image.set_data(data)
+
+
 class ToFloat32(SpatialTransform):
     """Convert the image to float32."""