ENH: Read ITK's composite transforms with only affines

nitransforms/conftest.py

@@ -9,6 +9,7 @@
 _data = None
 _brainmask = None
 _testdir = Path(os.getenv("TEST_DATA_HOME", "~/.nitransforms/testdata")).expanduser()
+_datadir = Path(__file__).parent / "tests" / "data"
 
 
 @pytest.fixture(autouse=True)
@@ -18,7 +19,7 @@ def doctest_autoimport(doctest_namespace):
     doctest_namespace["nb"] = nb
     doctest_namespace["os"] = os
     doctest_namespace["Path"] = Path
-    doctest_namespace["regress_dir"] = Path(__file__).parent / "tests" / "data"
+    doctest_namespace["regress_dir"] = _datadir
     doctest_namespace["test_dir"] = _testdir
 
     tmpdir = tempfile.TemporaryDirectory()
@@ -35,7 +36,7 @@ def doctest_autoimport(doctest_namespace):
 @pytest.fixture
 def data_path():
     """Return the test data folder."""
-    return Path(__file__).parent / "tests" / "data"
+    return _datadir
 
 
 @pytest.fixture

nitransforms/io/itk.py

@@ -2,6 +2,7 @@
 import warnings
 import numpy as np
 from scipy.io import loadmat as _read_mat, savemat as _save_mat
+from h5py import File as H5File
 from nibabel import Nifti1Header, Nifti1Image
 from nibabel.affines import from_matvec
 from .base import (
@@ -112,6 +113,9 @@ def from_filename(cls, filename):
         if str(filename).endswith(".mat"):
             with open(str(filename), "rb") as fileobj:
                 return cls.from_binary(fileobj)
+        elif str(filename).endswith(".h5"):
+            with H5File(str(filename)) as f:
+                return cls.from_h5obj(f)
 
         with open(str(filename)) as fileobj:
             return cls.from_string(fileobj.read())
@@ -121,6 +125,10 @@ def from_fileobj(cls, fileobj, check=True):
         """Read the struct from a file object."""
         if fileobj.name.endswith(".mat"):
             return cls.from_binary(fileobj)
+        elif fileobj.name.endswith(".h5"):
+            with H5File(fileobj) as f:
+                return cls.from_h5obj(f)
+
         return cls.from_string(fileobj.read())
 
     @classmethod
@@ -145,6 +153,27 @@ def from_matlab_dict(cls, mdict, index=0):
         sa["offset"] = mdict["fixed"].flatten()
         return tf
 
+    @classmethod
+    def from_h5obj(cls, fileobj, check=True):
+        """Read the struct from a file object."""
+
+        _xfm = ITKCompositeH5.from_h5obj(
+            fileobj,
+            check=check,
+            only_linear=True,
+        )
+
+        if not _xfm:
+            raise TransformIOError(
+                "Composite transform file does not contain at least one linear transform"
+            )
+        elif len(_xfm) > 1:
+            raise TransformIOError(
+                "Composite transform file contains more than one linear transform"
+            )
+
+        return _xfm[0]
+
     @classmethod
     def from_ras(cls, ras, index=0, moving=None, reference=None):
         """Create an ITK affine from a nitransform's RAS+ matrix."""
@@ -225,6 +254,9 @@ def from_filename(cls, filename):
         if str(filename).endswith(".mat"):
             with open(str(filename), "rb") as f:
                 return cls.from_binary(f)
+        elif str(filename).endswith(".h5"):
+            with H5File(str(filename)) as f:
+                return cls.from_h5obj(f)
 
         with open(str(filename)) as f:
             string = f.read()
@@ -235,6 +267,10 @@ def from_fileobj(cls, fileobj, check=True):
         """Read the struct from a file object."""
         if fileobj.name.endswith(".mat"):
             return cls.from_binary(fileobj)
+
+        elif fileobj.name.endswith(".h5"):
+            with H5File(fileobj) as f:
+                return cls.from_h5obj(f)
         return cls.from_string(fileobj.read())
 
     @classmethod
@@ -272,6 +308,18 @@ def from_string(cls, string):
             _self.xforms.append(cls._inner_type.from_string("#%s" % xfm))
         return _self
 
+    @classmethod
+    def from_h5obj(cls, fileobj, check=True):
+        """Read the struct from a file object."""
+
+        _self = cls()
+        _self.xforms = ITKCompositeH5.from_h5obj(
+            fileobj,
+            check=check,
+            only_linear=True,
+        )
+        return _self
+
 
 class ITKDisplacementsField(DisplacementsField):
     """A data structure representing displacements fields."""
@@ -302,18 +350,16 @@ class ITKCompositeH5:
     """A data structure for ITK's HDF5 files."""
 
     @classmethod
-    def from_filename(cls, filename):
+    def from_filename(cls, filename, only_linear=False):
         """Read the struct from a file given its path."""
-        from h5py import File as H5File
-
         if not str(filename).endswith(".h5"):
             raise TransformFileError("Extension is not .h5")
 
         with H5File(str(filename)) as f:
-            return cls.from_h5obj(f)
+            return cls.from_h5obj(f, only_linear=only_linear)
 
     @classmethod
-    def from_h5obj(cls, fileobj, check=True):
+    def from_h5obj(cls, fileobj, check=True, only_linear=False):
         """Read the struct from a file object."""
         xfm_list = []
         h5group = fileobj["TransformGroup"]
@@ -336,6 +382,8 @@ def from_h5obj(cls, fileobj, check=True):
                 )
                 continue
             if xfm["TransformType"][0].startswith(b"DisplacementFieldTransform"):
+                if only_linear:
+                    continue
                 _fixed = np.asanyarray(xfm[f"{typo_fallback}FixedParameters"])
                 shape = _fixed[:3].astype("uint16").tolist()
                 offset = _fixed[3:6].astype("float")

nitransforms/tests/test_io.py

@@ -15,18 +15,19 @@
 from nibabel.affines import from_matvec
 from scipy.io import loadmat
 from nitransforms.linear import Affine
-from ..io import (
+from nitransforms.io import (
     afni,
     fsl,
     lta as fs,
     itk,
 )
-from ..io.lta import (
+from nitransforms.io.lta import (
     VolumeGeometry as VG,
     FSLinearTransform as LT,
     FSLinearTransformArray as LTA,
 )
-from ..io.base import LinearParameters, TransformIOError, TransformFileError
+from nitransforms.io.base import LinearParameters, TransformIOError, TransformFileError
+from nitransforms.conftest import _datadir, _testdir
 
 LPS = np.diag([-1, -1, 1, 1])
 ITK_MAT = LPS.dot(np.ones((4, 4)).dot(LPS))
@@ -410,33 +411,35 @@ def test_afni_Displacements():
         afni.AFNIDisplacementsField.from_image(field)
 
 
-def test_itk_h5(tmpdir, testdata_path):
+@pytest.mark.parametrize("only_linear", [True, False])
+@pytest.mark.parametrize("h5_path,nxforms", [
+    (_datadir / "affine-antsComposite.h5", 1),
+    (_testdir / "ds-005_sub-01_from-T1w_to-MNI152NLin2009cAsym_mode-image_xfm.h5", 2),
+])
+def test_itk_h5(tmpdir, only_linear, h5_path, nxforms):
     """Test displacements fields."""
     assert (
         len(
             list(
                 itk.ITKCompositeH5.from_filename(
-                    testdata_path
-                    / "ds-005_sub-01_from-T1w_to-MNI152NLin2009cAsym_mode-image_xfm.h5"
+                    h5_path,
+                    only_linear=only_linear,
                 )
             )
         )
-        == 2
+        == nxforms if not only_linear else 1
     )
 
     with pytest.raises(TransformFileError):
         list(
             itk.ITKCompositeH5.from_filename(
-                testdata_path
-                / "ds-005_sub-01_from-T1w_to-MNI152NLin2009cAsym_mode-image_xfm.x5"
+                h5_path.absolute().name.replace(".h5", ".x5"),
+                only_linear=only_linear,
             )
         )
 
     tmpdir.chdir()
-    shutil.copy(
-        testdata_path / "ds-005_sub-01_from-T1w_to-MNI152NLin2009cAsym_mode-image_xfm.h5",
-        "test.h5",
-    )
+    shutil.copy(h5_path, "test.h5")
     os.chmod("test.h5", 0o666)
 
     with H5File("test.h5", "r+") as h5file:
@@ -584,3 +587,70 @@ def _generate_reoriented(path, directions, swapaxes, parameters):
     hdr.set_qform(newaff, code=1)
     hdr.set_sform(newaff, code=1)
     return img.__class__(data, newaff, hdr), R
+
+
+def test_itk_linear_h5(tmpdir, data_path, testdata_path):
+    """Check different lower-level loading options."""
+
+    # File loadable with transform array
+    h5xfm = itk.ITKLinearTransformArray.from_filename(
+        data_path / "affine-antsComposite.h5"
+    )
+    assert len(h5xfm.xforms) == 1
+
+    with open(data_path / "affine-antsComposite.h5", "rb") as f:
+        h5xfm = itk.ITKLinearTransformArray.from_fileobj(f)
+    assert len(h5xfm.xforms) == 1
+
+    # File loadable with single affine object
+    itk.ITKLinearTransform.from_filename(
+        data_path / "affine-antsComposite.h5"
+    )
+
+    with open(data_path / "affine-antsComposite.h5", "rb") as f:
+        itk.ITKLinearTransform.from_fileobj(f)
+
+    # Exercise only_linear
+    itk.ITKCompositeH5.from_filename(
+        testdata_path / "ds-005_sub-01_from-T1w_to-MNI152NLin2009cAsym_mode-image_xfm.h5",
+        only_linear=True,
+    )
+
+    tmpdir.chdir()
+    shutil.copy(data_path / "affine-antsComposite.h5", "test.h5")
+    os.chmod("test.h5", 0o666)
+
+    with H5File("test.h5", "r+") as h5file:
+        h5group = h5file["TransformGroup"]
+        xfm = h5group.create_group("2")
+        xfm["TransformType"] = (b"AffineTransform", b"")
+        xfm["TransformParameters"] = np.zeros(12, dtype=float)
+        xfm["TransformFixedParameters"] = np.zeros(3, dtype=float)
+
+    # File loadable with transform array
+    h5xfm = itk.ITKLinearTransformArray.from_filename("test.h5")
+    assert len(h5xfm.xforms) == 2
+
+    # File loadable with generalistic object (NOTE we directly access the list)
+    h5xfm = itk.ITKCompositeH5.from_filename("test.h5")
+    assert len(h5xfm) == 2
+
+    # Error raised if the we try to use the single affine loader
+    with pytest.raises(TransformIOError):
+        itk.ITKLinearTransform.from_filename("test.h5")
+
+    shutil.copy(data_path / "affine-antsComposite.h5", "test.h5")
+    os.chmod("test.h5", 0o666)
+
+    # Generate an empty h5 file
+    with H5File("test.h5", "r+") as h5file:
+        h5group = h5file["TransformGroup"]
+        del h5group["1"]
+
+    # File loadable with generalistic object
+    h5xfm = itk.ITKCompositeH5.from_filename("test.h5")
+    assert len(h5xfm) == 0
+
+    # Error raised if the we try to use the single affine loader
+    with pytest.raises(TransformIOError):
+        itk.ITKLinearTransform.from_filename("test.h5")