RF: ANTs h5 loading

Author: mgxd

nibabies/interfaces/resampling.py

@@ -24,6 +24,7 @@
 from scipy.sparse import hstack as sparse_hstack
 from sdcflows.transform import grid_bspline_weights
 from sdcflows.utils.tools import ensure_positive_cosines
+from transforms3d.affines import compose as compose_affine
 
 R = TypeVar('R')
 
@@ -62,16 +63,6 @@ def load_transforms(xfm_paths: list[Path], inverse: list[bool]) -> nt.base.Trans
     return chain
 
 
-FIXED_PARAMS = np.array([
-    193.0, 229.0, 193.0,  # Size
-    96.0, 132.0, -78.0,   # Origin
-    1.0, 1.0, 1.0,        # Spacing
-    -1.0, 0.0, 0.0,       # Directions
-    0.0, -1.0, 0.0,
-    0.0, 0.0, 1.0,
-])  # fmt:skip
-
-
 def load_ants_h5(filename: Path) -> nt.base.TransformBase:
     """Load ANTs H5 files as a nitransforms TransformChain"""
     # Borrowed from https://github.com/feilong/process
@@ -80,7 +71,8 @@ def load_ants_h5(filename: Path) -> nt.base.TransformBase:
     # Changes:
     #   * Tolerate a missing displacement field
     #   * Return the original affine without a round-trip
-    #   * Always return a nitransforms TransformChain
+    #   * Always return a nitransforms TransformBase
+    #   * Construct warp affine from fixed parameters
     #
     # This should be upstreamed into nitransforms
     h = h5py.File(filename)
@@ -104,22 +96,37 @@ def load_ants_h5(filename: Path) -> nt.base.TransformBase:
             msg += f'[{i}]: {h["TransformGroup"][i]["TransformType"][:][0]}\n'
         raise ValueError(msg)
 
-    fixed_params = transform2['TransformFixedParameters'][:]
-    shape = tuple(fixed_params[:3].astype(int))
+    # Warp field fixed parameters as defined in
+    # https://itk.org/Doxygen/html/classitk_1_1DisplacementFieldTransform.html
+    shape = transform2['TransformFixedParameters'][:3]
+    origin = transform2['TransformFixedParameters'][3:6]
+    spacing = transform2['TransformFixedParameters'][6:9]
+    direction = transform2['TransformFixedParameters'][9:].reshape((3, 3))
+
+    # We are not yet confident that we handle non-unit spacing
+    # or direction cosine ordering correctly.
+    # If we confirm or fix, we can remove these checks.
+    if not np.allclose(spacing, 1):
+        raise ValueError(f'Unexpected spacing: {spacing}')
+    if not np.allclose(direction, direction.T):
+        raise ValueError(f'Asymmetric direction matrix: {direction}')
+
+    # ITK uses LPS affines
+    lps_affine = compose_affine(T=origin, R=direction, Z=spacing)
+    ras_affine = np.diag([-1, -1, 1, 1]) @ lps_affine
+
     # ITK stores warps in Fortran-order, where the vector components change fastest
-    # Nitransforms expects 3 volumes, not a volume of three-vectors, so transpose
-    warp = np.reshape(
+    # Vectors are in mm LPS
+    itk_warp = np.reshape(
         transform2['TransformParameters'],
-        (3, *shape),
+        (3, *shape.astype(int)),
         order='F',
-    ).transpose(1, 2, 3, 0)
-
-    warp_affine = np.eye(4)
-    warp_affine[:3, :3] = fixed_params[9:].reshape((3, 3))
-    warp_affine[:3, 3] = fixed_params[3:6]
-    lps_to_ras = np.eye(4) * np.array([-1, -1, 1, 1])
-    warp_affine = lps_to_ras @ warp_affine
-    transforms.insert(0, nt.DenseFieldTransform(nb.Nifti1Image(warp, warp_affine)))
+    )
+
+    # Nitransforms warps are in RAS, with the vector components changing slowest
+    nt_warp = itk_warp.transpose(1, 2, 3, 0) * np.array([-1, -1, 1])
+
+    transforms.insert(0, nt.DenseFieldTransform(nb.Nifti1Image(nt_warp, ras_affine)))
     return nt.TransformChain(transforms)
 
 

pyproject.toml

@@ -35,6 +35,7 @@ dependencies = [
     "smriprep @ git+https://github.com/nipreps/smriprep.git@enh/nibabies-fit-apply",
     "tedana >= 23.0.2",
     "templateflow >= 24.2.0",
+    "transforms3d",
     "toml",
 ]
 dynamic = ["version"]