diff --git a/nipype/interfaces/ants/base.py b/nipype/interfaces/ants/base.py index 4d27fa619f..df45dd6c28 100644 --- a/nipype/interfaces/ants/base.py +++ b/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 diff --git a/nipype/interfaces/ants/segmentation.py b/nipype/interfaces/ants/segmentation.py index 06b9350dbc..d3319010d1 100644 --- a/nipype/interfaces/ants/segmentation.py +++ b/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) diff --git a/nipype/interfaces/ants/utils.py b/nipype/interfaces/ants/utils.py index 58c91f629f..6c3c05e9c4 100644 --- a/nipype/interfaces/ants/utils.py +++ b/nipype/interfaces/ants/utils.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): diff --git a/nipype/interfaces/mixins/__init__.py b/nipype/interfaces/mixins/__init__.py index a64dc34ff2..e54986231f 100644 --- a/nipype/interfaces/mixins/__init__.py +++ b/nipype/interfaces/mixins/__init__.py @@ -3,3 +3,4 @@ ReportCapableInputSpec, ReportCapableOutputSpec, ) +from .fixheader import CopyHeaderInputSpec, CopyHeaderInterface diff --git a/nipype/interfaces/mixins/fixheader.py b/nipype/interfaces/mixins/fixheader.py new file mode 100644 index 0000000000..3eb15785a3 --- /dev/null +++ b/nipype/interfaces/mixins/fixheader.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('>> 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