enh: added test file / reduce code duplicity

Author: oesteban

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

@@ -156,7 +156,23 @@ def from_matlab_dict(cls, mdict, index=0):
     @classmethod
     def from_h5obj(cls, fileobj, check=True):
         """Read the struct from a file object."""
-        raise NotImplementedError
+
+        _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):
@@ -295,26 +311,13 @@ def from_string(cls, string):
     @classmethod
     def from_h5obj(cls, fileobj, check=True):
         """Read the struct from a file object."""
-        h5group = fileobj["TransformGroup"]
-        typo_fallback = "Transform"
-        try:
-            h5group["1"][f"{typo_fallback}Parameters"]
-        except KeyError:
-            typo_fallback = "Tranform"
 
         _self = cls()
-        _self.xforms = []
-        for xfm in list(h5group.values())[1:]:
-            if xfm["TransformType"][0].startswith(b"AffineTransform"):
-                _params = np.asanyarray(xfm[f"{typo_fallback}Parameters"])
-                _self.xforms.append(
-                    ITKLinearTransform(
-                        parameters=from_matvec(
-                            _params[:-3].reshape(3, 3), _params[-3:]
-                        ),
-                        offset=np.asanyarray(xfm[f"{typo_fallback}FixedParameters"]),
-                    )
-                )
+        _self.xforms = ITKCompositeH5.from_h5obj(
+            fileobj,
+            check=check,
+            only_linear=True,
+        )
         return _self
 
 
@@ -347,16 +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."""
         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"]
@@ -379,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/data/affine-antsComposite.h5

@@ -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,72 @@ 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
+
+    h5xfm = itk.ITKLinearTransformArray.from_fileobj(
+        (data_path / "affine-antsComposite.h5").open()
+    )
+    assert len(h5xfm.xforms) == 1
+
+    # File loadable with single affine object
+    itk.ITKLinearTransform.from_filename(
+        data_path / "affine-antsComposite.h5"
+    )
+
+    itk.ITKLinearTransform.from_fileobj(
+        (data_path / "affine-antsComposite.h5").open()
+    )
+
+    # 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")