enh(afni): add warpdrive tests

Author: oesteban

nitransforms/io/afni.py

@@ -147,6 +147,21 @@ def from_image(cls, imgobj):
 
 
 def _is_oblique(affine, thres=OBLIQUITY_THRESHOLD_DEG):
+    """
+    Determine whether the dataset is oblique.
+
+    Examples
+    --------
+    >>> _is_oblique(np.eye(4))
+    False
+
+    >>> _is_oblique(nb.affines.from_matvec(
+    ...     nb.eulerangles.euler2mat(x=0.9, y=0.001, z=0.001),
+    ...     [4.0, 2.0, -1.0],
+    ... ))
+    True
+
+    """
     return (obliquity(affine).min() * 180 / pi) > thres
 
 
@@ -161,7 +176,8 @@ def _afni_warpdrive(oblique, shape, forward=True, ras=False):
     plumb : 4x4 numpy.array
         corresponding affine that is aligned to the cardinal axes.
     forward : :obj:`bool`
-        Transforms the affine of oblique into an AFNI's plumb (if ``True``)
+        Returns the forward transformation if True, i.e.,
+        the matrix to convert an oblique affine into an AFNI's plumb (if ``True``)
         or viceversa plumb -> oblique (if ``false``).
 
     Returns
@@ -177,21 +193,19 @@ def _afni_warpdrive(oblique, shape, forward=True, ras=False):
     plumb_r *= voxel_sizes(oblique)
     plumb = np.eye(4)
     plumb[:3, :3] = plumb_r
-    obliq_o = apply_affine(oblique, 0.5 * (shape - 1))
-    plumb_c = apply_affine(plumb, 0.5 * (shape - 1))
-    plumb[:3, 3] = -plumb_c + obliq_o
-    print(obliq_o, apply_affine(plumb, 0.5 * (shape - 1)))
-
+    obliq_c = oblique @ np.hstack((0.5 * shape, 1.0))
+    plumb_c = plumb @ np.hstack((0.5 * shape, 1.0))
+    plumb[:3, 3] = -plumb_c[:3] + obliq_c[:3]
     R = plumb @ np.linalg.inv(oblique)
+
     if not ras:
         # Change sign to match AFNI's warpdrive_matvec signs
-        B = np.ones((2, 2))
-        R *= np.block([[B, -1.0 * B], [-1.0 * B, B]])
+        R = LPS @ R @ LPS
 
     return R if forward else np.linalg.inv(R)
 
 
-def _afni_header(nii, field="WARPDRIVE_MATVEC_FOR_000000"):
+def _afni_header(nii, field="WARPDRIVE_MATVEC_FOR_000000", to_ras=False):
     from lxml import etree
     root = etree.fromstring(nii.header.extensions[0].get_content().decode())
     retval = np.fromstring(
@@ -200,6 +214,8 @@ def _afni_header(nii, field="WARPDRIVE_MATVEC_FOR_000000"):
         dtype="float32"
     )
     if retval.size == 12:
-        return np.vstack((retval.reshape((3, 4)), (0, 0, 0, 1)))
+        retval = np.vstack((retval.reshape((3, 4)), (0, 0, 0, 1)))
+        if to_ras:
+            retval = LPS @ retval @ LPS
 
     return retval

nitransforms/tests/test_io.py

@@ -1,11 +1,13 @@
 # emacs: -*- mode: python-mode; py-indent-offset: 4; indent-tabs-mode: nil -*-
 # vi: set ft=python sts=4 ts=4 sw=4 et:
 """I/O test cases."""
+from subprocess import check_call
+from io import StringIO
+import filecmp
+import shutil
 import numpy as np
 import pytest
-from io import StringIO
 
-import filecmp
 import nibabel as nb
 from nibabel.eulerangles import euler2mat
 from nibabel.affines import from_matvec
@@ -433,3 +435,80 @@ def test_itk_h5(testdata_path):
 )
 def test_regressions(file_type, test_file, data_path):
     file_type.from_filename(data_path / "regressions" / test_file)
+
+
+@pytest.mark.parametrize("parameters", [
+    {"x": 0.1, "y": 0.03, "z": 0.002},
+    {"x": 0.001, "y": 0.3, "z": 0.002},
+    {"x": 0.01, "y": 0.03, "z": 0.2},
+])
+@pytest.mark.parametrize("dir_x", (-1, 1))
+@pytest.mark.parametrize("dir_y", (-1, 1))
+@pytest.mark.parametrize("dir_z", (1, -1))
+@pytest.mark.parametrize("swapaxes", [
+    None, (0, 1), (1, 2), (0, 2),
+])
+def test_afni_oblique(tmpdir, parameters, swapaxes, testdata_path, dir_x, dir_y, dir_z):
+    tmpdir.chdir()
+    img = nb.load(testdata_path / "someones_anatomy.nii.gz")
+    shape = np.array(img.shape[:3])
+    hdr = img.header.copy()
+    aff = img.affine.copy()
+    data = np.asanyarray(img.dataobj, dtype="uint8")
+
+    directions = (dir_x, dir_y, dir_z)
+    if directions != (1, 1, 1):
+        last_ijk = np.hstack((shape - 1, 1.0))
+        last_xyz = aff @ last_ijk
+        aff = np.diag((dir_x, dir_y, dir_z, 1)) @ aff
+
+        for ax in range(3):
+            if (directions[ax] == -1):
+                aff[ax, 3] = last_xyz[ax]
+                data = np.flip(data, ax)
+
+        hdr.set_qform(aff, code=1)
+        hdr.set_sform(aff, code=1)
+        img.__class__(data, aff, hdr).to_filename("flips.nii.gz")
+
+    if swapaxes is not None:
+        data = np.swapaxes(data, swapaxes[0], swapaxes[1])
+        aff[reversed(swapaxes), :] = aff[(swapaxes), :]
+
+        hdr.set_qform(aff, code=1)
+        hdr.set_sform(aff, code=1)
+        img.__class__(data, aff, hdr).to_filename("swaps.nii.gz")
+
+    R = from_matvec(euler2mat(**parameters), [0.0, 0.0, 0.0])
+
+    img_center = aff @ np.hstack((shape * 0.5, 1.0))
+    R[:3, 3] += (img_center - (R @ aff @ np.hstack((shape * 0.5, 1.0))))[:3]
+    newaff = R @ aff
+    hdr.set_qform(newaff, code=1)
+    hdr.set_sform(newaff, code=1)
+    img = img.__class__(data, newaff, hdr)
+    img.to_filename("oblique.nii.gz")
+
+    # Run AFNI's 3dDeoblique
+    if not shutil.which("3dWarp"):
+        pytest.skip("Command 3dWarp not found on host")
+
+    cmd = f"3dWarp -verb -deoblique -prefix {tmpdir}/deob.nii.gz {tmpdir}/oblique.nii.gz"
+
+    # resample mask
+    assert check_call([cmd], shell=True) == 0
+    afni_warpdrive_inv = afni._afni_header(
+        nb.load("deob.nii.gz"),
+        field="WARPDRIVE_MATVEC_INV_000000",
+        to_ras=True,
+    )
+
+    deobnii = nb.load("deob.nii.gz")
+
+    # Confirm AFNI's rotation of axis is consistent with the one we introduced
+    assert np.allclose(afni_warpdrive_inv[:3, :3], R[:3, :3])
+
+    # Check nitransforms' estimation of warpdrive with header
+    nt_warpdrive_inv = afni._afni_warpdrive(newaff, img.shape, forward=False)
+    import pdb;pdb.set_trace()
+    assert np.allclose(afni_warpdrive_inv[:3, :3], nt_warpdrive_inv[:3, :3])

nitransforms/tests/test_linear.py

@@ -165,7 +165,7 @@ def test_linear_save(tmpdir, data_path, get_testdata, image_orientation, sw_tool
     assert_affines_by_filename(xfm_fname1, xfm_fname2)
 
 
-@pytest.mark.parametrize("image_orientation", ["RAS", "LAS", "LPS", ])  # 'oblique',
+@pytest.mark.parametrize("image_orientation", ["RAS", "LAS", "LPS", 'oblique', ])
 @pytest.mark.parametrize("sw_tool", ["itk", "fsl", "afni", "fs"])
 def test_apply_linear_transform(tmpdir, get_testdata, get_testmask, image_orientation, sw_tool):
     """Check implementation of exporting affines to formats."""