ENH: Subcortical alignment workflow

nibabies/config.py

@@ -684,6 +684,12 @@ def init_spaces(checkpoint=True):
     if checkpoint and not spaces.is_cached():
         spaces.checkpoint()
 
+    if workflow.age_months is not None:
+        from .utils.misc import cohort_by_months
+
+        if "MNIInfant" not in spaces.get_spaces(nonstandard=False, dim=(3,)):
+            cohort = cohort_by_months("MNIInfant", workflow.age_months)
+            spaces.add(Reference("MNIInfant", {"cohort": cohort}))
     # # Add the default standard space if not already present (required by several sub-workflows)
     # if "MNI152NLin2009cAsym" not in spaces.get_spaces(nonstandard=False, dim=(3,)):
     #     spaces.add(Reference("MNI152NLin2009cAsym", {}))
@@ -695,10 +701,9 @@ def init_spaces(checkpoint=True):
     #     # Make sure there's a normalization to FSL for AROMA to use.
     #     spaces.add(Reference("MNI152NLin6Asym", {"res": "2"}))
 
-    cifti_output = workflow.cifti_output
     if workflow.cifti_output:
         # CIFTI grayordinates to corresponding FSL-MNI resolutions.
-        vol_res = "2" if cifti_output == "91k" else "1"
+        vol_res = "2" if workflow.cifti_output == "91k" else "1"
         spaces.add(Reference("fsaverage", {"den": "164k"}))
         spaces.add(Reference("MNI152NLin6Asym", {"res": vol_res}))
 

nibabies/interfaces/conftest.py → nibabies/conftest.py

@@ -1,8 +1,9 @@
 """py.test configuration"""
 from pathlib import Path
-import pytest
 from tempfile import TemporaryDirectory
+from pkg_resources import resource_filename
 
+import pytest
 
 FILES = (
     'functional.nii',
@@ -32,3 +33,4 @@ def data_dir():
 @pytest.fixture(autouse=True)
 def set_namespace(doctest_namespace, data_dir):
     doctest_namespace["data_dir"] = data_dir
+    doctest_namespace["test_data"] = Path(resource_filename("nibabies", "tests/data"))

nibabies/interfaces/nibabel.py

@@ -5,6 +5,7 @@
     BaseInterfaceInputSpec,
     File,
     SimpleInterface,
+    InputMultiObject,
 )
 
 
@@ -34,6 +35,24 @@ def _run_interface(self, runtime):
         return runtime
 
 
+class MergeROIsInputSpec(BaseInterfaceInputSpec):
+    in_files = InputMultiObject(File(exists=True), desc="ROI files to be merged")
+
+
+class MergeROIsOutputSpec(TraitedSpec):
+    out_file = File(exists=True, desc="NIfTI containing all ROIs")
+
+
+class MergeROIs(SimpleInterface):
+    """Combine multiple region of interest files (3D or 4D) into a single file"""
+    input_spec = MergeROIsInputSpec
+    output_spec = MergeROIsOutputSpec
+
+    def _run_interface(self, runtime):
+        self._results["out_file"] = _merge_rois(self.inputs.in_files, newpath=runtime.cwd)
+        return runtime
+
+
 def _dilate(in_file, radius=3, iterations=1, newpath=None):
     """Dilate (binary) input mask."""
     from pathlib import Path
@@ -55,3 +74,41 @@ def _dilate(in_file, radius=3, iterations=1, newpath=None):
     out_file = fname_presuffix(in_file, suffix="_dil", newpath=newpath or Path.cwd())
     mask.__class__(newdata.astype("uint8"), mask.affine, hdr).to_filename(out_file)
     return out_file
+
+
+def _merge_rois(in_files, newpath=None):
+    """
+    Aggregate individual 4D ROI files together into a single subcortical NIfTI.
+    All ROI images are sanity checked with regards to:
+    1) Shape
+    2) Affine
+    3) Overlap
+
+    If any of these checks fail, an ``AssertionError`` will be raised.
+    """
+    from pathlib import Path
+    import nibabel as nb
+    import numpy as np
+
+    img = nb.load(in_files[0])
+    data = np.array(img.dataobj)
+    affine = img.affine
+    header = img.header
+
+    nonzero = np.any(data, axis=3)
+    for roi in in_files[1:]:
+        img = nb.load(roi)
+        assert img.shape == data.shape, "Mismatch in image shape"
+        assert np.allclose(img.affine, affine), "Mismatch in affine"
+        roi_data = np.asanyarray(img.dataobj)
+        roi_nonzero = np.any(roi_data, axis=3)
+        assert not np.any(roi_nonzero & nonzero), "Overlapping ROIs"
+        nonzero |= roi_nonzero
+        data += roi_data
+        del roi_data
+
+    if newpath is None:
+        newpath = Path()
+    out_file = str((Path(newpath) / "combined.nii.gz").absolute())
+    img.__class__(data, affine, header).to_filename(out_file)
+    return out_file

nibabies/interfaces/tests/test_nibabel.py

@@ -0,0 +1,73 @@
+import uuid
+
+import nibabel as nb
+import numpy as np
+import pytest
+
+from ..nibabel import MergeROIs
+
+
+@pytest.fixture
+def create_roi(tmp_path):
+    files = []
+
+    def _create_roi(affine, img_data, roi_index):
+        img_data[tuple(roi_index)] = 1
+        nii = nb.Nifti1Image(img_data, affine)
+        filename = tmp_path / f"{str(uuid.uuid4())}.nii.gz"
+        files.append(filename)
+        nii.to_filename(filename)
+        return filename
+
+    yield _create_roi
+
+    for f in files:
+        f.unlink()
+
+
+# create a slightly off affine
+bad_affine = np.eye(4)
+bad_affine[0, -1] = -1
+
+
+@pytest.mark.parametrize(
+    "affine, data, roi_index, error, err_message",
+    [
+        (np.eye(4), np.zeros((2, 2, 2, 2), dtype=int), [1, 0], None, None),
+        (
+            np.eye(4),
+            np.zeros((2, 2, 3, 2), dtype=int),
+            [1, 0],
+            True,
+            "Mismatch in image shape",
+        ),
+        (
+            bad_affine,
+            np.zeros((2, 2, 2, 2), dtype=int),
+            [1, 0],
+            True,
+            "Mismatch in affine",
+        ),
+        (
+            np.eye(4),
+            np.zeros((2, 2, 2, 2), dtype=int),
+            [0, 0, 0],
+            True,
+            "Overlapping ROIs",
+        ),
+    ],
+)
+def test_merge_rois(tmpdir, create_roi, affine, data, roi_index, error, err_message):
+    tmpdir.chdir()
+    roi0 = create_roi(np.eye(4), np.zeros((2, 2, 2, 2), dtype=int), [0, 0])
+    roi1 = create_roi(np.eye(4), np.zeros((2, 2, 2, 2), dtype=int), [0, 1])
+    test_roi = create_roi(affine, data, roi_index)
+
+    merge = MergeROIs(in_files=[roi0, roi1, test_roi])
+    if error is None:
+        merge.run()
+        return
+    # otherwise check expected exceptions
+    with pytest.raises(AssertionError) as err:
+        merge.run()
+    assert err_message in str(err.value)