ENH: Validate files passed with --derivatives (#182)

* ENH: Add derivatives validation during workflow construction

* ENH: Copy T1w header to avoid precision errors

* FIX: Ensure derivatives point to filenames

* RF: Simplify physical space check

Co-authored-by: Chris Markiewicz <[email protected]>

* FIX: Reinsert `ValidateImage` for orientation safety

* FIX: Remove unused import

Co-authored-by: Chris Markiewicz <[email protected]>

Author: mgxd

nibabies/utils/bids.py

@@ -337,6 +337,7 @@ def collect_precomputed_derivatives(layout, subject_id, derivatives_filters=None
             scope='derivatives',
             subject=subject_id,
             extension=['.nii', '.nii.gz'],
+            return_type="filename",
             **query,
         )
         if not res:

nibabies/utils/validation.py

@@ -0,0 +1,52 @@
+import sys
+import nibabel as nb
+import numpy as np
+
+
+def validate_t1w_derivatives(t1w_template, *, anat_mask=None, anat_aseg=None, atol=1e-5):
+    """
+    Validate anatomical derivatives.
+    This function compares the input T1w's orientation and physical space to each derivative.
+
+    Parameters
+    ----------
+    t1w_template : str
+        T1w template
+    anat_mask : str or None
+        Precomputed anatomical brain mask
+    anat_aseg : str or None
+        Precomputed anatomical segmentations
+    atol : float
+        Absolute error tolerance between image origins
+
+    Returns
+    -------
+    validated : dict
+        A dictionary composed of derivative keys and validated filename values.
+        Derivatives that failed validation will not be included.
+    """
+
+    validated = {}
+    # T1w information
+    t1w = nb.load(t1w_template)
+    expected_ort = nb.aff2axcodes(t1w.affine)
+
+    # Ensure orientation
+    for name, deriv_fl in zip(('anat_mask', 'anat_aseg'), (anat_mask, anat_aseg)):
+        if deriv_fl is None:
+            continue
+        img = nb.load(deriv_fl)
+        if nb.aff2axcodes(img.affine) != expected_ort:
+            print(
+                f"Orientation mismatch between {name} <{deriv_fl}> and T1w <{t1w_template}>",
+                file=sys.stderr,
+            )
+            continue
+        if img.shape != t1w.shape or not np.allclose(t1w.affine, img.affine, atol=atol):
+            print(
+                f"Physical space mismatch between {name} <{deriv_fl}> and T1w <{t1w_template}>",
+                file=sys.stderr,
+            )
+            continue
+        validated[name] = deriv_fl
+    return validated

nibabies/workflows/anatomical/base.py

@@ -93,6 +93,16 @@ def init_infant_anat_wf(
     precomp_mask = existing_derivatives.get("anat_mask")
     precomp_aseg = existing_derivatives.get("anat_aseg")
 
+    # verify derivatives are relatively similar to T1w
+    if precomp_mask or precomp_aseg:
+        from ...utils.validation import validate_t1w_derivatives
+
+        validated_derivatives = validate_t1w_derivatives(  # compare derivatives to the first T1w
+            t1w[0], anat_mask=precomp_mask, anat_aseg=precomp_aseg
+        )
+        precomp_mask = validated_derivatives.get("anat_mask")
+        precomp_aseg = validated_derivatives.get("anat_aseg")
+
     wf = Workflow(name=name)
     desc = f"""\n
 Anatomical data preprocessing

nibabies/workflows/anatomical/brain_extraction.py

@@ -272,7 +272,7 @@ def init_precomputed_mask_wf(
     bspline_fitting_distance=200, omp_nthreads=None, name="precomputed_mask_wf"
 ):
     from nipype.interfaces.ants import N4BiasFieldCorrection
-    from niworkflows.interfaces.header import ValidateImage
+    from niworkflows.interfaces.header import CopyXForm, ValidateImage
     from niworkflows.interfaces.nibabel import ApplyMask, IntensityClip
 
     workflow = pe.Workflow(name=name)
@@ -283,6 +283,7 @@ def init_precomputed_mask_wf(
     )
 
     validate_mask = pe.Node(ValidateImage(), name="validate_mask")
+    fix_hdr = pe.Node(CopyXForm(), mem_gb=0.1, name="fix_hdr")
     final_n4 = pe.Node(
         N4BiasFieldCorrection(
             dimension=3,
@@ -303,12 +304,14 @@ def init_precomputed_mask_wf(
     # fmt:off
     workflow.connect([
         (inputnode, validate_mask, [("t1w_mask", "in_file")]),
+        (validate_mask, fix_hdr, [("out_file", "in_file")]),
+        (inputnode, fix_hdr, [("t1w", "hdr_file")]),
         (inputnode, final_n4, [("t1w", "input_image")]),
-        (validate_mask, final_n4, [("out_file", "weight_image")]),
-        (validate_mask, apply_mask, [("out_file", "in_mask")]),
+        (fix_hdr, final_n4, [("out_file", "weight_image")]),
+        (fix_hdr, apply_mask, [("out_file", "in_mask")]),
         (final_n4, apply_mask, [("output_image", "in_file")]),
         (final_n4, final_clip, [("output_image", "in_file")]),
-        (validate_mask, outputnode, [("out_file", "t1w_mask")]),
+        (fix_hdr, outputnode, [("out_file", "t1w_mask")]),
         (final_clip, outputnode, [("out_file", "t1w_preproc")]),
         (apply_mask, outputnode, [("out_file", "t1w_brain")]),
     ])