enh: deep revision of EstimateReferenceImage

Author: oesteban

niworkflows/interfaces/registration.py

@@ -8,7 +8,7 @@
 import numpy as np
 from nilearn import image as nli
 from nilearn.image import index_img
-from nipype.utils.filemanip import fname_presuffix
+from nipype.utils.filemanip import fname_presuffix, ensure_list
 from nipype.interfaces.base import (
     traits,
     isdefined,
@@ -445,59 +445,81 @@ class _EstimateReferenceImageOutputSpec(TraitedSpec):
 
 class EstimateReferenceImage(SimpleInterface):
     """
-    Given an 4D EPI file estimate an optimal reference image that could be later
-    used for motion estimation and coregistration purposes. If detected uses
-    T1 saturated volumes (non-steady state) otherwise a median of
+    Generate a reference 3D map from BOLD and SBRef EPI images for BOLD datasets.
+
+    Given a 4D BOLD file or one or more 3/4D SBRefs, estimate a reference
+    image for subsequent motion estimation and coregistration steps.
+    For the case of BOLD datasets, it estimates a number of T1w saturated volumes
+    (non-steady state at the beginning of the scan) and calculates the median
+    across them.
+    Otherwise (SBRefs or detected zero non-steady state frames), a median of
     of a subset of motion corrected volumes is used.
+    If the input reference (BOLD or SBRef) is 3D already, it just returns a
+    copy of the image with the NIfTI header extensions removed.
+
+    LIMITATION: If one wants to extract the reference from several SBRefs
+    with several echoes each, the first echo should be selected elsewhere
+    and run this interface in ``multiecho = False`` mode.
     """
 
     input_spec = _EstimateReferenceImageInputSpec
     output_spec = _EstimateReferenceImageOutputSpec
 
     def _run_interface(self, runtime):
-        # Select first EPI file
-        ref_name = self.inputs.in_file[0]
-        ref_nii = nb.load(ref_name)
-        n_volumes_to_discard = _get_vols_to_discard(ref_nii)
-
-        self._results["n_volumes_to_discard"] = n_volumes_to_discard
-
-        if self.inputs.multiecho and (len(self.inputs.in_file) < 2):
-            raise ValueError("Argument 'multiecho' is True but "
-                             "'in_file' has only one element")
+        is_sbref = isdefined(self.inputs.sbref_file)
+        ref_input = ensure_list(
+            self.inputs.sbref_file if is_sbref else self.inputs.in_file
+        )
 
-        out_ref_fname = os.path.join(runtime.cwd, "ref_bold.nii.gz")
-        if isdefined(self.inputs.sbref_file):
-            if self.inputs.multiecho and (len(self.inputs.sbref_file) < 2):
+        if self.inputs.multiecho:
+            if len(ref_input) < 2:
+                input_name = "sbref_file" if is_sbref else "in_file"
                 raise ValueError("Argument 'multiecho' is True but "
-                                 "'sbref_file' has only one element")
-
-            # Select first SBRef file
-            ref_name = self.inputs.sbref_file[0]
-            ref_nii = nb.squeeze_image(nb.load(ref_name))
+                                 f"'{input_name}' has only one element.")
+            else:
+                # Select only the first echo (see LIMITATION above for SBRefs)
+                ref_input = ref_input[:1]
+        elif not is_sbref and len(ref_input) > 1:
+            raise ValueError("Input 'in_file' cannot point to more than one file "
+                             "for single-echo BOLD datasets.")
+
+        # Build the nibabel spatial image we will work with
+        ref_im = []
+        for im_i in ref_input:
+            nib_i = nb.squeeze_image(nb.load(im_i))
+            if nib_i.dataobj.ndim == 3:
+                ref_im.append(nib_i)
+            elif nib_i.dataobj.ndim == 4:
+                ref_im += nb.four_to_three(nib_i)
+        ref_im = nb.squeeze_image(nb.concat_images(ref_im))
+
+        # Volumes to discard only makes sense with BOLD inputs.
+        if not is_sbref:
+            n_volumes_to_discard = _get_vols_to_discard(ref_im)
+            out_ref_fname = os.path.join(runtime.cwd, "ref_bold.nii.gz")
+        else:
+            n_volumes_to_discard = 0
             out_ref_fname = os.path.join(runtime.cwd, "ref_sbref.nii.gz")
 
-            # If reference is only 1 volume, return it directly
-            if len(ref_nii.shape) == 3:
-                ref_nii.header.extensions.clear()
-                ref_nii.to_filename(out_ref_fname)
-                self._results["ref_image"] = out_ref_fname
-                return runtime
-            else:
-                # Reset this variable as it no longer applies
-                # and value for the output is stored in self._results
-                n_volumes_to_discard = 0
+        # Set interface outputs
+        self._results["n_volumes_to_discard"] = n_volumes_to_discard
+        self._results["ref_image"] = out_ref_fname
 
         # Slicing may induce inconsistencies with shape-dependent values in extensions.
         # For now, remove all. If this turns out to be a mistake, we can select extensions
         # that don't break pipeline stages.
-        ref_nii.header.extensions.clear()
+        ref_im.header.extensions.clear()
+
+        # If reference is only 1 volume, return it directly
+        if ref_im.dataobj.ndim == 3:
+            ref_im.to_filename(out_ref_fname)
+            return runtime
 
         if n_volumes_to_discard == 0:
-            if ref_nii.shape[-1] > 40:
+            if ref_im.shape[-1] > 40:
                 ref_name = os.path.join(runtime.cwd, "slice.nii.gz")
                 nb.Nifti1Image(
-                    ref_nii.dataobj[:, :, :, 20:40], ref_nii.affine, ref_nii.header
+                    ref_im.dataobj[:, :, :, 20:40], ref_im.affine, ref_im.header
                 ).to_filename(ref_name)
 
             if self.inputs.mc_method == "AFNI":
@@ -516,14 +538,12 @@ def _run_interface(self, runtime):
             median_image_data = np.median(mc_slice_nii.get_fdata(), axis=3)
         else:
             median_image_data = np.median(
-                ref_nii.dataobj[:, :, :, :n_volumes_to_discard], axis=3
+                ref_im.dataobj[:, :, :, :n_volumes_to_discard], axis=3
             )
 
-        nb.Nifti1Image(median_image_data, ref_nii.affine, ref_nii.header).to_filename(
+        nb.Nifti1Image(median_image_data, ref_im.affine, ref_im.header).to_filename(
             out_ref_fname
         )
-
-        self._results["ref_image"] = out_ref_fname
         return runtime