Merge pull request #10 from effigies/enh/ants-utils-overhaul

RF: Move header copying to a mixin

Author: oesteban

nipype/interfaces/ants/base.py

@@ -122,17 +122,3 @@ def set_default_num_threads(cls, num_threads):
     @property
     def version(self):
         return Info.version()
-
-
-class FixHeaderANTSCommand(ANTSCommand):
-    """Fix header if the copy_header input is on."""
-
-    def aggregate_outputs(self, runtime=None, needed_outputs=None):
-        """Overload the aggregation with header replacement, if required."""
-        outputs = super(FixHeaderANTSCommand, self).aggregate_outputs(
-            runtime, needed_outputs)
-        if self.inputs.copy_header:  # Fix headers
-            _copy_header(
-                self.inputs.op1, outputs["output_image"], keep_dtype=True
-            )
-        return outputs

nipype/interfaces/ants/segmentation.py

@@ -2,9 +2,9 @@
 import os
 from glob import glob
 from ...external.due import BibTeX
-from ...utils.imagemanip import copy_header as _copy_header
 from ...utils.filemanip import split_filename, copyfile, which, fname_presuffix
 from ..base import TraitedSpec, File, traits, InputMultiPath, OutputMultiPath, isdefined
+from ..mixins import CopyHeaderInterface
 from .base import ANTSCommand, ANTSCommandInputSpec
 
 
@@ -420,7 +420,7 @@ class N4BiasFieldCorrectionOutputSpec(TraitedSpec):
     bias_image = File(exists=True, desc="Estimated bias")
 
 
-class N4BiasFieldCorrection(ANTSCommand):
+class N4BiasFieldCorrection(ANTSCommand, CopyHeaderInterface):
     """
     Bias field correction.
 
@@ -491,6 +491,10 @@ class N4BiasFieldCorrection(ANTSCommand):
     _cmd = "N4BiasFieldCorrection"
     input_spec = N4BiasFieldCorrectionInputSpec
     output_spec = N4BiasFieldCorrectionOutputSpec
+    _copy_header_map = {
+        "output_image": ("input_image", False),
+        "bias_image": ("input_image", True),
+    }
 
     def __init__(self, *args, **kwargs):
         """Instantiate the N4BiasFieldCorrection interface."""
@@ -533,20 +537,6 @@ def _parse_inputs(self, skip=None):
             self._out_bias_file = bias_image
         return super(N4BiasFieldCorrection, self)._parse_inputs(skip=skip)
 
-    def _list_outputs(self):
-        outputs = super(N4BiasFieldCorrection, self)._list_outputs()
-
-        # Fix headers
-        if self.inputs.copy_header:
-            _copy_header(self.inputs.input_image, outputs["output_image"],
-                         keep_dtype=False)
-
-        if self._out_bias_file:
-            outputs["bias_image"] = os.path.abspath(self._out_bias_file)
-            if self.inputs.copy_header:
-                _copy_header(self.inputs.input_image, outputs["bias_image"])
-        return outputs
-
 
 class CorticalThicknessInputSpec(ANTSCommandInputSpec):
     dimension = traits.Enum(
@@ -1501,7 +1491,7 @@ def _format_arg(self, opt, spec, val):
             for option in (
                 self.inputs.out_intensity_fusion_name_format,
                 self.inputs.out_label_post_prob_name_format,
-                self.inputs.out_atlas_voting_weight_name_format
+                self.inputs.out_atlas_voting_weight_name_format,
             ):
                 if isdefined(option):
                     args.append(option)

nipype/interfaces/ants/tests/test_auto_AntsJointFusion.py

@@ -1,7 +1,8 @@
 """ANTs' utilities."""
 import os
 from ..base import traits, isdefined, TraitedSpec, File, Str, InputMultiObject
-from .base import ANTSCommandInputSpec, ANTSCommand, FixHeaderANTSCommand
+from ..mixins import CopyHeaderInterface
+from .base import ANTSCommandInputSpec, ANTSCommand
 
 
 class ImageMathInputSpec(ANTSCommandInputSpec):
@@ -67,7 +68,7 @@ class ImageMathOuputSpec(TraitedSpec):
     output_image = File(exists=True, desc="output image file")
 
 
-class ImageMath(FixHeaderANTSCommand):
+class ImageMath(ANTSCommand, CopyHeaderInterface):
     """
     Operations over images.
 
@@ -96,6 +97,7 @@ class ImageMath(FixHeaderANTSCommand):
     _cmd = "ImageMath"
     input_spec = ImageMathInputSpec
     output_spec = ImageMathOuputSpec
+    _copy_header_map = {"output_image": "op1"}
 
 
 class ResampleImageBySpacingInputSpec(ANTSCommandInputSpec):
@@ -146,7 +148,7 @@ class ResampleImageBySpacingOutputSpec(TraitedSpec):
     output_image = File(exists=True, desc="resampled file")
 
 
-class ResampleImageBySpacing(FixHeaderANTSCommand):
+class ResampleImageBySpacing(ANTSCommand, CopyHeaderInterface):
     """
     Resample an image with a given spacing.
 
@@ -182,6 +184,7 @@ class ResampleImageBySpacing(FixHeaderANTSCommand):
     _cmd = "ResampleImageBySpacing"
     input_spec = ResampleImageBySpacingInputSpec
     output_spec = ResampleImageBySpacingOutputSpec
+    _copy_header_map = {"output_image": "input_image"}
 
     def _format_arg(self, name, trait_spec, value):
         if name == "out_spacing":
@@ -248,7 +251,7 @@ class ThresholdImageOutputSpec(TraitedSpec):
     output_image = File(exists=True, desc="resampled file")
 
 
-class ThresholdImage(FixHeaderANTSCommand):
+class ThresholdImage(ANTSCommand, CopyHeaderInterface):
     """
     Apply thresholds on images.
 
@@ -277,6 +280,7 @@ class ThresholdImage(FixHeaderANTSCommand):
     _cmd = "ThresholdImage"
     input_spec = ThresholdImageInputSpec
     output_spec = ThresholdImageOutputSpec
+    _copy_header_map = {"output_image": "input_image"}
 
 
 class AIInputSpec(ANTSCommandInputSpec):

nipype/interfaces/ants/tests/test_auto_FixHeaderANTSCommand.py

@@ -3,3 +3,4 @@
     ReportCapableInputSpec,
     ReportCapableOutputSpec,
 )
+from .fixheader import CopyHeaderInputSpec, CopyHeaderInterface

nipype/interfaces/ants/utils.py

@@ -0,0 +1,134 @@
+from ..base import BaseInterface, BaseInterfaceInputSpec, traits
+from ...utils.imagemanip import copy_header as _copy_header
+
+
+class CopyHeaderInputSpec(BaseInterfaceInputSpec):
+    copy_header = traits.Bool(
+        desc="Copy headers of the input image into the output image"
+    )
+
+
+class CopyHeaderInterface(BaseInterface):
+    """ Copy headers if the copy_header input is ``True``
+
+    This interface mixin adds a post-run hook that allows for copying
+    an input header to an output file.
+    The subclass should specify a ``_copy_header_map`` that maps the **output**
+    image to the **input** image whose header should be copied.
+
+    This feature is intended for tools that are intended to adjust voxel data without
+    modifying the header, but for some reason do not reliably preserve the header.
+
+    Here we show an example interface that takes advantage of the mixin by simply
+    setting the data block:
+
+    >>> import os
+    >>> import numpy as np
+    >>> import nibabel as nb
+    >>> from nipype.interfaces.base import SimpleInterface, TraitedSpec, File
+    >>> from nipype.interfaces.mixins import CopyHeaderInputSpec, CopyHeaderInterface
+
+    >>> class ZerofileInputSpec(CopyHeaderInputSpec):
+    ...     in_file = File(mandatory=True, exists=True)
+
+    >>> class ZerofileOutputSpec(TraitedSpec):
+    ...     out_file = File()
+
+    >>> class ZerofileInterface(SimpleInterface, CopyHeaderInterface):
+    ...     input_spec = ZerofileInputSpec
+    ...     output_spec = ZerofileOutputSpec
+    ...     _copy_header_map = {'out_file': 'in_file'}
+    ...
+    ...     def _run_interface(self, runtime):
+    ...         img = nb.load(self.inputs.in_file)
+    ...         # Just set the data. Let the CopyHeaderInterface mixin fix the affine and header.
+    ...         nb.Nifti1Image(np.zeros(img.shape, dtype=np.uint8), None).to_filename('out.nii')
+    ...         self._results = {'out_file':  os.path.abspath('out.nii')}
+    ...         return runtime
+
+    Consider a file of all ones and a non-trivial affine:
+
+    >>> in_file = 'test.nii'
+    >>> nb.Nifti1Image(np.ones((5,5,5), dtype=np.int16),
+    ...                affine=np.diag((4, 3, 2, 1))).to_filename(in_file)
+
+    The default behavior would produce a file with similar data:
+
+    >>> res = ZerofileInterface(in_file=in_file).run()
+    >>> out_img = nb.load(res.outputs.out_file)
+    >>> out_img.shape
+    (5, 5, 5)
+    >>> np.all(out_img.get_fdata() == 0)
+    True
+
+    An updated data type:
+
+    >>> out_img.get_data_dtype()
+    dtype('uint8')
+
+    But a different affine:
+
+    >>> np.array_equal(out_img.affine, np.diag((4, 3, 2, 1)))
+    False
+
+    With ``copy_header=True``, then the affine is also equal:
+
+    >>> res = ZerofileInterface(in_file=in_file, copy_header=True).run()
+    >>> out_img = nb.load(res.outputs.out_file)
+    >>> np.array_equal(out_img.affine, np.diag((4, 3, 2, 1)))
+    True
+
+    The data properties remain as expected:
+
+    >>> out_img.shape
+    (5, 5, 5)
+    >>> out_img.get_data_dtype()
+    dtype('uint8')
+    >>> np.all(out_img.get_fdata() == 0)
+    True
+
+    By default, the data type of the output file is permitted to vary from the
+    inputs. That is, the data type is preserved.
+    If the data type of the original file is preferred, the ``_copy_header_map``
+    can indicate the output data type should **not** be preserved by providing a
+    tuple of the input and ``False``.
+
+    >>> ZerofileInterface._copy_header_map['out_file'] = ('in_file', False)
+
+    >>> res = ZerofileInterface(in_file=in_file, copy_header=True).run()
+    >>> out_img = nb.load(res.outputs.out_file)
+    >>> out_img.get_data_dtype()
+    dtype('<i2')
+
+    Again, the affine is updated.
+
+    >>> np.array_equal(out_img.affine, np.diag((4, 3, 2, 1)))
+    True
+    >>> out_img.shape
+    (5, 5, 5)
+    >>> np.all(out_img.get_fdata() == 0)
+    True
+
+    Providing a tuple where the second value is ``True`` is also permissible to
+    achieve the default behavior.
+
+    """
+
+    _copy_header_map = None
+
+    def _post_run_hook(self, runtime):
+        """Copy headers for outputs, if required."""
+        runtime = super()._post_run_hook(runtime)
+
+        if self._copy_header_map is None or not self.inputs.copy_header:
+            return runtime
+
+        inputs = self.inputs.get_traitsfree()
+        outputs = self.aggregate_outputs(runtime=runtime).get_traitsfree()
+        for out, inp in self._copy_header_map.items():
+            keep_dtype = True
+            if isinstance(inp, tuple):
+                inp, keep_dtype = inp
+            _copy_header(inputs[inp], outputs[out], keep_dtype=keep_dtype)
+
+        return runtime

nipype/interfaces/mixins/__init__.py

@@ -0,0 +1,11 @@
+# AUTO-GENERATED by tools/checkspecs.py - DO NOT EDIT
+from ..fixheader import CopyHeaderInterface
+
+
+def test_CopyHeaderInterface_inputs():
+    input_map = dict()
+    inputs = CopyHeaderInterface.input_spec()
+
+    for key, metadata in list(input_map.items()):
+        for metakey, value in list(metadata.items()):
+            assert getattr(inputs.traits()[key], metakey) == value