Merge pull request #736 from mgxd/enh/reorient-image

ENH: ReorientImage interface

Author: mgxd

niworkflows/interfaces/cifti.py

@@ -29,7 +29,6 @@
 from nibabel import cifti2 as ci
 import numpy as np
 from nilearn.image import resample_to_img
-
 from nipype.utils.filemanip import split_filename
 from nipype.interfaces.base import (
     BaseInterfaceInputSpec,
@@ -41,6 +40,8 @@
 )
 import templateflow.api as tf
 
+from niworkflows.interfaces.nibabel import reorient_image
+
 CIFTI_SURFACES = ("fsaverage5", "fsaverage6", "fsLR")
 CIFTI_VOLUMES = ("MNI152NLin2009cAsym", "MNI152NLin6Asym")
 CIFTI_STRUCT_WITH_LABELS = {  # CITFI structures with corresponding labels
@@ -356,8 +357,8 @@ def _create_cifti_image(
         bold_img = resample_to_img(bold_img, label_img)
 
     # ensure images match HCP orientation (LAS)
-    bold_img = _reorient_image(bold_img, orientation="LAS")
-    label_img = _reorient_image(label_img, orientation="LAS")
+    bold_img = reorient_image(bold_img, target_ornt="LAS")
+    label_img = reorient_image(label_img, target_ornt="LAS")
 
     bold_data = bold_img.get_fdata(dtype="float32")
     timepoints = bold_img.shape[3]
@@ -466,66 +467,3 @@ def _create_cifti_image(
     out_file = "{}.dtseries.nii".format(split_filename(bold_file)[1])
     ci.save(img, out_file)
     return Path.cwd() / out_file
-
-
-def _reorient_image(img, *, target_img=None, orientation=None):
-    """
-    Coerce an image to a target orientation.
-
-    .. note::
-        Only RAS -> LAS conversion is currently supported
-
-    Parameters
-    ----------
-    img : :obj:`SpatialImage`
-        image to be reoriented
-    target_img : :obj:`SpatialImage`, optional
-        target in desired orientation
-    orientation : :obj:`str` or :obj:`tuple`, optional
-        desired orientation, if no target image is provided
-
-    .. testsetup::
-    >>> img = nb.load(Path(test_data) / 'testSpatialNormalizationRPTMovingWarpedImage.nii.gz')
-    >>> las_img = img.as_reoriented([[0, -1], [1, 1], [2, 1]])
-
-    Examples
-    --------
-    >>> nimg = _reorient_image(img, target_img=img)
-    >>> nb.aff2axcodes(nimg.affine)
-    ('R', 'A', 'S')
-
-    >>> nimg = _reorient_image(img, target_img=las_img)
-    >>> nb.aff2axcodes(nimg.affine)
-    ('L', 'A', 'S')
-
-    >>> nimg = _reorient_image(img, orientation='LAS')
-    >>> nb.aff2axcodes(nimg.affine)
-    ('L', 'A', 'S')
-
-    >>> _reorient_image(img, orientation='LPI')
-    Traceback (most recent call last):
-      ...
-    NotImplementedError: Cannot reorient ...
-
-    >>> _reorient_image(img)
-    Traceback (most recent call last):
-      ...
-    RuntimeError: No orientation ...
-
-    """
-    orient0 = nb.aff2axcodes(img.affine)
-    if target_img is not None:
-        orient1 = nb.aff2axcodes(target_img.affine)
-    elif orientation is not None:
-        orient1 = tuple(orientation)
-    else:
-        raise RuntimeError("No orientation to reorient to!")
-
-    if orient0 == orient1:  # already in desired orientation
-        return img
-    elif orient0 == tuple("RAS") and orient1 == tuple("LAS"):  # RAS -> LAS
-        return img.as_reoriented([[0, -1], [1, 1], [2, 1]])
-    else:
-        raise NotImplementedError(
-            "Cannot reorient {0} to {1}.".format(orient0, orient1)
-        )

niworkflows/interfaces/nibabel.py

@@ -21,6 +21,8 @@
 #     https://www.nipreps.org/community/licensing/
 #
 """Nibabel-based interfaces."""
+from pathlib import Path
+
 import numpy as np
 import nibabel as nb
 from nipype import logging
@@ -494,6 +496,79 @@ def _run_interface(self, runtime):
         return runtime
 
 
+class ReorientImageInputSpec(BaseInterfaceInputSpec):
+    in_file = File(exists=True, mandatory=True, desc="Moving file")
+    target_file = File(
+        exists=True, xor=["target_orientation"], desc="Reference file to reorient to"
+    )
+    target_orientation = traits.Str(
+        xor=["target_file"], desc="Axis codes of coordinate system to reorient to"
+    )
+
+
+class ReorientImageOutputSpec(TraitedSpec):
+    out_file = File(desc="Reoriented file")
+
+
+class ReorientImage(SimpleInterface):
+    input_spec = ReorientImageInputSpec
+    output_spec = ReorientImageOutputSpec
+
+    def _run_interface(self, runtime):
+        self._results["out_file"] = reorient_file(
+            self.inputs.in_file,
+            target_file=self.inputs.target_file,
+            target_ornt=self.inputs.target_orientation,
+        )
+        return runtime
+
+
+def reorient_file(
+    in_file: str, *, target_file: str = None, target_ornt: str = None, newpath: str = None,
+) -> str:
+    """
+    Reorient an image.
+
+    New orientation targets can be either another image, or a string representation of the
+    orientation axis.
+
+    Parameters
+    ----------
+    in_file : Image to be reoriented
+    target_file : Reference image of desired orientation
+    target_ornt : Orientation denoted by the first letter of each axis (i.e., "RAS", "LPI")
+    """
+    import nibabel as nb
+
+    img = nb.load(in_file)
+    if not target_file and not target_ornt:
+        raise TypeError("No target orientation or file is specified.")
+
+    if target_file:
+        target_img = nb.load(target_file)
+        target_ornt = nb.aff2axcodes(target_img.affine)
+
+    reoriented = reorient_image(img, target_ornt)
+
+    if newpath is None:
+        newpath = Path()
+    out_file = str((Path(newpath) / "reoriented.nii.gz").absolute())
+    reoriented.to_filename(out_file)
+    return out_file
+
+
+def reorient_image(img: nb.spatialimages.SpatialImage, target_ornt: str):
+    """Reorient an image in memory."""
+    import nibabel as nb
+
+    img_axcodes = nb.aff2axcodes(img.affine)
+    in_ornt = nb.orientations.axcodes2ornt(img_axcodes)
+    out_ornt = nb.orientations.axcodes2ornt(target_ornt)
+    ornt_xfm = nb.orientations.ornt_transform(in_ornt, out_ornt)
+    r_img = img.as_reoriented(ornt_xfm)
+    return r_img
+
+
 def _gen_reference(
     fixed_image,
     moving_image,

niworkflows/interfaces/tests/test_nibabel.py

@@ -29,7 +29,9 @@
 import nibabel as nb
 import pytest
 
-from ..nibabel import Binarize, ApplyMask, SplitSeries, MergeSeries, MergeROIs, MapLabels
+from ..nibabel import (
+    Binarize, ApplyMask, SplitSeries, MergeSeries, MergeROIs, MapLabels, ReorientImage
+)
 
 
 @pytest.fixture
@@ -280,3 +282,55 @@ def test_map_labels(tmpdir, data, mapping, tojson, expected):
     Path(in_file).unlink()
     if tojson:
         Path(map_file).unlink()
+
+
+def create_save_img(ornt: str):
+    data = np.random.rand(2, 2, 2)
+    img = nb.Nifti1Image(data, affine=np.eye(4))
+    # img will alway be in RAS at the start
+    ras = nb.orientations.axcodes2ornt("RAS")
+    if ornt != 'RAS':
+        new = nb.orientations.axcodes2ornt(ornt)
+        xfm = nb.orientations.ornt_transform(ras, new)
+        img = img.as_reoriented(xfm)
+    out_file = f'{uuid.uuid4()}.nii.gz'
+    img.to_filename(out_file)
+    return out_file
+
+
+@pytest.mark.parametrize(
+    "in_ornt,out_ornt",
+    [
+        ("RAS", "RAS"),
+        ("RAS", "LAS"),
+        ("LAS", "RAS"),
+        ("RAS", "RPI"),
+        ("LPI", "RAS"),
+    ],
+)
+def test_reorient_image(tmpdir, in_ornt, out_ornt):
+    tmpdir.chdir()
+
+    in_file = create_save_img(ornt=in_ornt)
+    in_img = nb.load(in_file)
+    assert ''.join(nb.aff2axcodes(in_img.affine)) == in_ornt
+
+    # test string representation
+    res = ReorientImage(in_file=in_file, target_orientation=out_ornt).run()
+    out_file = res.outputs.out_file
+    out_img = nb.load(out_file)
+    assert ''.join(nb.aff2axcodes(out_img.affine)) == out_ornt
+    Path(out_file).unlink()
+
+    # test with target file
+    target_file = create_save_img(ornt=out_ornt)
+    target_img = nb.load(target_file)
+    assert ''.join(nb.aff2axcodes(target_img.affine)) == out_ornt
+    res = ReorientImage(in_file=in_file, target_file=target_file).run()
+    out_file = res.outputs.out_file
+    out_img = nb.load(out_file)
+    assert ''.join(nb.aff2axcodes(out_img.affine)) == out_ornt
+
+    # cleanup
+    for f in (in_file, target_file, out_file):
+        Path(f).unlink()