Merge pull request #2034 from oesteban/enh/N4-header

[ENH] Add copy_header option to N4BiasFieldCorrection

Author: satra

nipype/interfaces/ants/segmentation.py

@@ -259,6 +259,9 @@ class N4BiasFieldCorrectionInputSpec(ANTSCommandInputSpec):
                                   ' to file.'), xor=['bias_image'])
     bias_image = File(desc='Filename for the estimated bias.',
                       hash_files=False)
+    copy_header = traits.Bool(False, mandatory=True, usedefault=True,
+                              desc='copy headers of the original image into the '
+                                   'output (corrected) file')
 
 
 class N4BiasFieldCorrectionOutputSpec(TraitedSpec):
@@ -384,6 +387,27 @@ def _list_outputs(self):
                 self._gen_filename('bias_image'))
         return outputs
 
+    def _run_interface(self, runtime, correct_return_codes=(0,)):
+        runtime = super(N4BiasFieldCorrection, self)._run_interface(
+            runtime, correct_return_codes)
+
+        if self.inputs.copy_header and runtime.returncode in correct_return_codes:
+            self._copy_header(self._gen_filename('output_image'))
+            if self.inputs.save_bias or isdefined(self.inputs.bias_image):
+                self._copy_header(self._gen_filename('bias_image'))
+
+        return runtime
+
+    def _copy_header(self, fname):
+        """Copy header from input image to an output image"""
+        import nibabel as nb
+        in_img = nb.load(self.inputs.input_image)
+        out_img = nb.load(fname, mmap=False)
+        new_img = out_img.__class__(out_img.get_data(), in_img.affine,
+                                    in_img.header)
+        new_img.set_data_dtype(out_img.get_data_dtype())
+        new_img.to_filename(fname)
+
 
 class CorticalThicknessInputSpec(ANTSCommandInputSpec):
     dimension = traits.Enum(3, 2, argstr='-d %d', usedefault=True,

nipype/interfaces/ants/tests/test_auto_N4BiasFieldCorrection.py

@@ -14,6 +14,9 @@ def test_N4BiasFieldCorrection_inputs():
     ),
     convergence_threshold=dict(requires=['n_iterations'],
     ),
+    copy_header=dict(mandatory=True,
+    usedefault=True,
+    ),
     dimension=dict(argstr='-d %d',
     usedefault=True,
     ),