Merge pull request #197 from mgxd/enh/pepolar-alignment

ENH: Uniformize the grid&affine across EPI "blips" before TOPUP

Author: oesteban

sdcflows/interfaces/epi.py

@@ -18,6 +18,7 @@ class _GetReadoutTimeInputSpec(BaseInterfaceInputSpec):
 class _GetReadoutTimeOutputSpec(TraitedSpec):
     readout_time = traits.Float
     pe_direction = traits.Enum("i", "i-", "j", "j-", "k", "k-")
+    pe_dir_fsl = traits.Enum("x", "x-", "y", "y-", "z", "z-")
 
 
 class GetReadoutTime(SimpleInterface):
@@ -34,4 +35,10 @@ def _run_interface(self, runtime):
             self.inputs.in_file if isdefined(self.inputs.in_file) else None,
         )
         self._results["pe_direction"] = self.inputs.metadata["PhaseEncodingDirection"]
+        self._results["pe_dir_fsl"] = (
+            self.inputs.metadata["PhaseEncodingDirection"]
+            .replace("i", "x")
+            .replace("j", "y")
+            .replace("k", "z")
+        )
         return runtime

sdcflows/interfaces/utils.py

@@ -66,6 +66,67 @@ def _run_interface(self, runtime):
         return runtime
 
 
+class _UniformGridInputSpec(BaseInterfaceInputSpec):
+    in_data = InputMultiObject(
+        File(exists=True),
+        mandatory=True,
+        desc="list of input data",
+    )
+    reference = traits.Int(0, usedefault=True, desc="reference index")
+
+
+class _UniformGridOutputSpec(TraitedSpec):
+    out_data = OutputMultiObject(File(exists=True))
+    reference = File(exists=True)
+
+
+class UniformGrid(SimpleInterface):
+    """Ensure all images in input have the same spatial parameters."""
+
+    input_spec = _UniformGridInputSpec
+    output_spec = _UniformGridOutputSpec
+
+    def _run_interface(self, runtime):
+        import nibabel as nb
+        import numpy as np
+        from nitransforms.linear import Affine
+        from nipype.utils.filemanip import fname_presuffix
+
+        retval = [None] * len(self.inputs.in_data)
+        self._results["reference"] = self.inputs.in_data[self.inputs.reference]
+        retval[self.inputs.reference] = self._results["reference"]
+
+        refnii = nb.load(self._results["reference"])
+        refshape = refnii.shape[:3]
+        refaff = refnii.affine
+
+        resampler = Affine(reference=refnii)
+        for i, fname in enumerate(self.inputs.in_data):
+            if retval[i] is not None:
+                continue
+
+            nii = nb.load(fname)
+            retval[i] = fname_presuffix(
+                fname, suffix=f"_regrid{i:03d}", newpath=runtime.cwd
+            )
+
+            if np.allclose(nii.shape[:3], refshape) and np.allclose(nii.affine, refaff):
+                if np.all(nii.affine == refaff):
+                    retval[i] = fname
+                else:
+                    # Set reference's affine if difference is small
+                    nii.__class__(nii.dataobj, refaff, nii.header).to_filename(
+                        retval[i]
+                    )
+                continue
+
+            resampler.apply(nii).to_filename(retval[i])
+
+        self._results["out_data"] = retval
+
+        return runtime
+
+
 class _ConvertWarpInputSpec(BaseInterfaceInputSpec):
     in_file = File(exists=True, mandatory=True, desc="output of 3dQwarp")
 

sdcflows/workflows/fit/pepolar.py

@@ -24,7 +24,11 @@
 
 
 def init_topup_wf(
-    omp_nthreads=1, sloppy=False, debug=False, name="pepolar_estimate_wf"
+    grid_reference=0,
+    omp_nthreads=1,
+    sloppy=False,
+    debug=False,
+    name="pepolar_estimate_wf",
 ):
     """
     Create the PEPOLAR field estimation workflow based on FSL's ``topup``.
@@ -39,6 +43,8 @@ def init_topup_wf(
 
     Parameters
     ----------
+    grid_reference : :obj:`int`
+        Index of the volume (after flattening) that will be taken for gridding reference.
     sloppy : :obj:`bool`
         Whether a fast configuration of topup (less accurate) should be applied.
     debug : :obj:`bool`
@@ -76,7 +82,7 @@ def init_topup_wf(
 
     from ...utils.misc import front as _front
     from ...interfaces.epi import GetReadoutTime
-    from ...interfaces.utils import Flatten
+    from ...interfaces.utils import Flatten, UniformGrid
     from ...interfaces.bspline import TOPUPCoeffReorient
     from ..ancillary import init_brainextraction_wf
 
@@ -104,6 +110,7 @@ def init_topup_wf(
     outputnode.inputs.method = "PEB/PEPOLAR (phase-encoding based / PE-POLARity)"
 
     flatten = pe.Node(Flatten(), name="flatten")
+    regrid = pe.Node(UniformGrid(reference=grid_reference), name="regrid")
     concat_blips = pe.Node(MergeSeries(), name="concat_blips")
     readout_time = pe.MapNode(
         GetReadoutTime(),
@@ -128,23 +135,18 @@ def init_topup_wf(
 
     brainextraction_wf = init_brainextraction_wf()
 
-    def _getfirstpe(in_meta):
-        if isinstance(in_meta, list):
-            in_meta = in_meta[0]
-        return in_meta["PhaseEncodingDirection"]
-
     # fmt: off
     workflow.connect([
         (inputnode, flatten, [("in_data", "in_data"),
                               ("metadata", "in_meta")]),
         (flatten, readout_time, [("out_data", "in_file"),
                                  ("out_meta", "metadata")]),
-        (flatten, concat_blips, [("out_data", "in_files")]),
-        (flatten, topup, [(("out_meta", _pe2fsl), "encoding_direction")]),
-        (readout_time, topup, [("readout_time", "readout_times")]),
-        (concat_blips, topup, [("out_file", "in_file")]),
-        (flatten, fix_coeff, [(("out_data", _front), "fmap_ref"),
-                              (("out_meta", _getfirstpe), "pe_dir")]),
+        (flatten, regrid, [("out_data", "in_data")]),
+        (regrid, concat_blips, [("out_data", "in_files")]),
+        (readout_time, topup, [("readout_time", "readout_times"),
+                               ("pe_dir_fsl", "encoding_direction")]),
+        (regrid, fix_coeff, [("reference", "fmap_ref")]),
+        (readout_time, fix_coeff, [(("pe_direction", _front), "pe_dir")]),
         (topup, fix_coeff, [("out_fieldcoef", "in_coeff")]),
         (topup, outputnode, [("out_jacs", "jacobians"),
                              ("out_mats", "xfms")]),
@@ -158,30 +160,35 @@ def _getfirstpe(in_meta):
     if not debug:
         # fmt: off
         workflow.connect([
+            (concat_blips, topup, [("out_file", "in_file")]),
             (topup, merge_corrected, [("out_corrected", "in_files")]),
             (topup, outputnode, [("out_field", "fmap"),
                                  ("out_warps", "out_warps")]),
         ])
         # fmt: on
         return workflow
 
+    from nipype.interfaces.afni.preprocess import Volreg
+    from niworkflows.interfaces.nibabel import SplitSeries
     from ...interfaces.bspline import ApplyCoeffsField
 
+    realign = pe.Node(
+        Volreg(args=f"-base {grid_reference}", outputtype="NIFTI_GZ"),
+        name="realign_blips",
+    )
+    split_blips = pe.Node(SplitSeries(), name="split_blips")
     unwarp = pe.Node(ApplyCoeffsField(), name="unwarp")
     unwarp.interface._always_run = True
 
-    def _getpe(inlist):
-        if isinstance(inlist, dict):
-            inlist = [inlist]
-
-        return [m["PhaseEncodingDirection"] for m in inlist]
-
     # fmt:off
     workflow.connect([
+        (concat_blips, realign, [("out_file", "in_file")]),
+        (realign, topup, [("out_file", "in_file")]),
         (fix_coeff, unwarp, [("out_coeff", "in_coeff")]),
-        (flatten, unwarp, [("out_data", "in_target"),
-                           (("out_meta", _getpe), "pe_dir")]),
-        (readout_time, unwarp, [("readout_time", "ro_time")]),
+        (realign, split_blips, [("out_file", "in_file")]),
+        (split_blips, unwarp, [("out_files", "in_target")]),
+        (readout_time, unwarp, [("readout_time", "ro_time"),
+                                ("pe_direction", "pe_dir")]),
         (unwarp, outputnode, [("out_warp", "out_warps"),
                               ("out_field", "fmap")]),
         (unwarp, merge_corrected, [("out_corrected", "in_files")]),
@@ -343,25 +350,6 @@ def init_3dQwarp_wf(omp_nthreads=1, debug=False, name="pepolar_estimate_wf"):
     return workflow
 
 
-def _pe2fsl(metadata):
-    """
-    Convert ijk notation to xyz.
-
-    Example
-    -------
-    >>> _pe2fsl([{"PhaseEncodingDirection": "j-"}, {"PhaseEncodingDirection": "i"}])
-    ['y-', 'x']
-
-    """
-    return [
-        m["PhaseEncodingDirection"]
-        .replace("i", "x")
-        .replace("j", "y")
-        .replace("k", "z")
-        for m in metadata
-    ]
-
-
 def _sorted_pe(inlist):
     """
     Generate suitable inputs to ``3dQwarp``.