/
fixes.py
121 lines (98 loc) · 4.16 KB
/
fixes.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
import os
import nibabel as nb
from nipype.interfaces.base import traits
from nipype.utils.filemanip import fname_presuffix
from nipype.interfaces.ants.resampling import ApplyTransforms
from nipype.interfaces.ants.registration import Registration
from nipype.interfaces.ants.segmentation import (
N4BiasFieldCorrection as VanillaN4,
N4BiasFieldCorrectionOutputSpec as VanillaN4OutputSpec,
)
from .. import __version__
from .utils import _copyxform
class FixHeaderApplyTransforms(ApplyTransforms):
"""
A replacement for nipype.interfaces.ants.resampling.ApplyTransforms that
fixes the resampled image header to match the xform of the reference
image
"""
def _run_interface(self, runtime, correct_return_codes=(0,)):
# Run normally
runtime = super(FixHeaderApplyTransforms, self)._run_interface(
runtime, correct_return_codes
)
_copyxform(
self.inputs.reference_image,
os.path.abspath(self._gen_filename("output_image")),
message="%s (niworkflows v%s)" % (self.__class__.__name__, __version__),
)
return runtime
class FixHeaderRegistration(Registration):
"""
A replacement for nipype.interfaces.ants.registration.Registration that
fixes the resampled image header to match the xform of the reference
image
"""
def _run_interface(self, runtime, correct_return_codes=(0,)):
# Run normally
runtime = super(FixHeaderRegistration, self)._run_interface(
runtime, correct_return_codes
)
# Forward transform
out_file = self._get_outputfilenames(inverse=False)
if out_file is not None and out_file:
_copyxform(
self.inputs.fixed_image[0],
os.path.abspath(out_file),
message="%s (niworkflows v%s)" % (self.__class__.__name__, __version__),
)
# Inverse transform
out_file = self._get_outputfilenames(inverse=True)
if out_file is not None and out_file:
_copyxform(
self.inputs.moving_image[0],
os.path.abspath(out_file),
message="%s (niworkflows v%s)" % (self.__class__.__name__, __version__),
)
return runtime
class _FixN4BiasFieldCorrectionOutputSpec(VanillaN4OutputSpec):
negative_values = traits.Bool(
False,
usedefault=True,
desc="Indicates whether the input was corrected for "
"nonpositive values by adding a constant offset.",
)
class FixN4BiasFieldCorrection(VanillaN4):
"""Checks and fixes for nonpositive values in the input to ``N4BiasFieldCorrection``."""
output_spec = _FixN4BiasFieldCorrectionOutputSpec
def __init__(self, *args, **kwargs):
"""Add a private property to keep the path to the right input."""
self._input_image = None
self._negative_values = False
super(FixN4BiasFieldCorrection, self).__init__(*args, **kwargs)
def _format_arg(self, name, trait_spec, value):
if name == "input_image":
return trait_spec.argstr % self._input_image
return super(FixN4BiasFieldCorrection, self)._format_arg(
name, trait_spec, value
)
def _parse_inputs(self, skip=None):
self._input_image = self.inputs.input_image
# Check intensities
input_nii = nb.load(self.inputs.input_image)
datamin = input_nii.get_fdata().min()
if datamin < 0:
self._input_image = fname_presuffix(
self.inputs.input_image, suffix="_scaled", newpath=os.getcwd()
)
data = input_nii.get_fdata() - datamin
newnii = input_nii.__class__(data, input_nii.affine, input_nii.header)
newnii.to_filename(self._input_image)
self._negative_values = True
return super(FixN4BiasFieldCorrection, self)._parse_inputs(skip=skip)
def _list_outputs(self):
outputs = super(FixN4BiasFieldCorrection, self)._list_outputs()
outputs["negative_values"] = self._negative_values
return outputs