-
Notifications
You must be signed in to change notification settings - Fork 0
/
mribrew_dwi_processing_x.py
288 lines (231 loc) · 11.1 KB
/
mribrew_dwi_processing_x.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
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
# %%
import os
from nipype import config, logging
from nipype.interfaces import io, fsl, mrtrix3
import nipype.interfaces.utility as niu
import nipype.pipeline.engine as pe
from mribrew.utils import colours
import mribrew.dwiproc_interface as ProcInterface
# ---------------------- Set up directory structures and constant variables ----------------------
cwd = os.getcwd()
misc_dir = os.path.join(cwd, 'misc')
data_dir = os.path.join(cwd, 'data')
raw_dir = os.path.join(data_dir, 'raw')
proc_dir = os.path.join(data_dir, 'proc', 'dwi_proc')
wf_dir = os.path.join(cwd, 'wf')
log_dir = os.path.join(wf_dir, 'log')
acqp_file = os.path.join(misc_dir, 'acqp.txt')
# List of all subjects
subject_list = next(os.walk(raw_dir))[1]
# DWI sequence file names
dwi_name = 'dir-AP_dwi' #'ep2d_diff_hardi_s2'
dwipa_name = 'dir-PA_dwi' #'ep2d_diff_hardi_s2_pa'
# Computational variables
processing_type = 'MultiProc' # or 'Linear'
cuda_processing = False
total_memory = 6 # in GB
n_cpus = 6 # number of nipype processes to run at the same time
os.environ['OMP_NUM_THREADS'] = str(n_cpus)
os.environ["NUMEXPR_NUM_THREADS"] = str(n_cpus)
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
plugin_args = {
'n_procs': n_cpus,
'memory_gb': total_memory,
'raise_insufficient': True,
'scheduler': 'mem_thread', # Prioritize jobs by memory consumption then nr of threads
}
# Set up logging
os.makedirs(log_dir, exist_ok=True)
config.update_config({'logging': {'log_directory': log_dir,'log_to_file': True}})
logging.update_logging(config)
# ---------------------- INPUT SOURCE NODES ----------------------
print(colours.CGREEN + "Creating Source Nodes." + colours.CEND)
# Set up input files
info = dict(dwi_file=[['subject_id', 'dwi', '*%s.nii.gz' % dwi_name]],
bvec_file=[['subject_id', 'dwi','*%s.bvec' % dwi_name]],
bval_file=[['subject_id', 'dwi','*%s.bval' % dwi_name]],
dwiPA_file=[['subject_id','dwi', '*%s.nii.gz' % dwipa_name]])
# Set up infosource node
infosource = pe.Node(niu.IdentityInterface(fields=['subject_id']), name='infosource')
infosource.iterables = [('subject_id', subject_list)]
infosource.inputs.acqp_file = acqp_file # add to the act similarly!!! ###############################################
# Set up datasource node
datasource = pe.Node(io.DataGrabber(infields=['subject_id'], outfields=list(info.keys())),
name='datasource')
datasource.inputs.base_directory = raw_dir
datasource.inputs.template = "%s/%s/%s"
datasource.inputs.template_args = info
datasource.inputs.sort_filelist = True
# ---------------------- OUTPUT SINK NODE ----------------------
print(colours.CGREEN + "Creating Sink Node." + colours.CEND)
# Set up sink node where all output is stored in subject folder
datasink = pe.Node(io.DataSink(parameterization=False), name='datasink')
datasink.inputs.base_directory = proc_dir
# ---------------------- PROCESSING NODES ----------------------
print(colours.CGREEN + "Creating Processing Nodes." + colours.CEND)
### DENOISE & DEGIBBS
# Using MRtrix3's DWIDenoise to reduce random noise
mrtrixDenoise = pe.Node(mrtrix3.DWIDenoise(), name='mrtrixDenoise')
# Removing Gibbs ringing artifacts using MRtrix3's MRDeGibbs function
mrtrixDegibbs = pe.Node(mrtrix3.MRDeGibbs(), name='mrtrixDegibbs')
### BRAIN MASK 1 (pre-topup/eddy)
# Extracting the brain mask from the raw data prior to any corrections using FSL's BET function
betMask1 = pe.Node(fsl.BET(), name = 'betMask1')
betMask1.inputs.mask = True
betMask1.inputs.output_type = 'NIFTI_GZ'
betMask1.inputs.no_output = True
betMask1.inputs.functional = True
### TOPUP
# Ensuring the 3D dimensionality of the input images.
checkDimension = pe.Node(ProcInterface.checkDimension(), name='checkDimension')
# Select 1st volume of dwi-ap for topup
dwiB0 = pe.Node(fsl.ExtractROI(), name='dwiB0')
dwiB0.inputs.t_min = 0
dwiB0.inputs.t_size = 1
dwiB0.inputs.output_type = 'NIFTI_GZ'
# Select 1st volume of dwi-pa for topup
dwiPAB0 = pe.Node(fsl.ExtractROI(), name='dwiPAB0')
dwiPAB0.inputs.t_min = 0
dwiPAB0.inputs.t_size = 1
dwiPAB0.inputs.output_type = 'NIFTI_GZ'
# Handle odd dimensions by cutting off a slice from all three planes
cutOddB0 = pe.Node(fsl.ExtractROI(), name='cutOddB0')
cutOddB0.inputs.x_size = -1
cutOddB0.inputs.y_size = -1
cutOddB0.inputs.z_size = -1
cutOddB0.inputs.roi_file = 'cutOddB0.nii.gz'
cutOddB0.inputs.output_type = 'NIFTI_GZ'
cutOddPA = cutOddB0.clone(name = 'cutOddPA')
cutOddPA.inputs.roi_file = 'cutOddPA.nii.gz'
# Merge b0s of dwi-ap and dwi-pa
listAPPA = pe.Node(niu.Merge(2), name='listAPPA')
mergeAPPA = pe.Node(fsl.Merge(), name='mergeAPPA')
mergeAPPA.inputs.dimension = 't'
mergeAPPA.inputs.merged_file = 'mergeAPPA.nii.gz'
mergeAPPA.inputs.output_type = 'NIFTI_GZ'
# Topup correction
topup = pe.Node(fsl.TOPUP(), name='topup')
topup.inputs.output_type = "NIFTI_GZ"
### EDDY
# Adjust b-values if both b100 and b0 exist
adjustBval = pe.Node(ProcInterface.adjustBval(), name='adjustBval')
adjustBval.inputs.valold = 100
adjustBval.inputs.valnew = 0
# Create index file for Eddy
eddyIndex = pe.Node(ProcInterface.eddyIndex(), name='eddyIndex')
# Eddy correction
eddy = pe.Node(fsl.Eddy(), name='eddy')
eddy.inputs.interp = 'spline'
eddy.inputs.use_cuda = cuda_processing
eddy.inputs.is_shelled = True
eddy.inputs.args = '--ol_nstd=5 --repol'
eddy.inputs.output_type = 'NIFTI_GZ'
### BRAIN MASK 2 (post-topup/eddy)
# Using FSL's BET function after corrections
betMask2 = betMask1.clone(name='betMask2')
# Creating a brain mask using MRtrix
mrtrixMask = pe.Node(ProcInterface.MRTRIX3BrainMask(), name='mrtrixMask')
mrtrixMask.inputs.out_name = 'mrtrix_mask.nii.gz'
# Combine different masks to create a final DWI brain mask
dwiMask = pe.Node(ProcInterface.combineDWIBrainMask(), name='dwiMask')
dwiMask.inputs.out_name = 'dwi_mask.nii.gz'
### GRADIENT CHECK
# Check gradient directions using MRtrix
mrtrixGradCheck = pe.Node(ProcInterface.MRTRIX3GradCheck(), name='mrtrixGradCheck')
# ---------------------- CREATE WORKFLOW AND CONNECT NODES ----------------------
print(colours.CGREEN + 'Connecting Nodes.\n' + colours.CEND)
workflow = pe.Workflow(name='dwiproc_wf', base_dir=f"{wf_dir}")
workflow.connect([
# ---------------------- INPUT/OUTPUT STRUCTURE (Handling input/output directories)
# Linking subject's information to data source
(infosource, datasource, [('subject_id', 'subject_id')]),
# Setting the output directory structure based on the subject ID
(infosource, datasink, [('subject_id', 'container')]),
# ---------------------- DENOISE (noise reduction)
# Apply denoising to the DWI data using mrtrixDenoise
(datasource, mrtrixDenoise, [('dwi_file', 'in_file')]),
# ---------------------- DEGIBBS (Correction for Gibbs ringing artifacts)
# Use denoised data to correct for Gibbs ringing artifacts with mrtrixDegibbs
(mrtrixDenoise, mrtrixDegibbs, [('out_file', 'in_file')]),
# ---------------------- BRAIN MASK 1 (creation of initial brain mask pre-topup/eddy)
# Extract initial brain mask from the DWI data using BET
(datasource, betMask1, [('dwi_file', 'in_file')]),
# ---------------------- TOPUP (distrortion correction)
# Extract B0 image (AP direction) for susceptibility correction
(datasource, dwiB0, [('dwi_file', 'in_file')]),
# Check the dimensions of the B0 (AP direction)
(dwiB0, checkDimension, [('roi_file', 'in_file')]),
# If dimensions are odd, cut slices to ensure compatibility with TOPUP
(dwiB0, cutOddB0, [('roi_file', 'in_file')]),
(checkDimension, cutOddB0, [('axialCutX', 'x_min'),
('axialCutY', 'y_min'),
('axialCutZ', 'z_min')]),
# Extract B0 image (PA direction)
(datasource, dwiPAB0, [('dwiPA_file', 'in_file')]),
# If dimensions are odd for the PA B0, cut slices to ensure compatibility with TOPUP
(dwiPAB0, cutOddPA, [('roi_file', 'in_file')]),
(checkDimension, cutOddPA, [('axialCutX', 'x_min'),
('axialCutY', 'y_min'),
('axialCutZ', 'z_min')]),
# Merge the B0 images from both AP and PA phase-encode directions
(cutOddB0, listAPPA, [('roi_file', 'in1')]),
(cutOddPA, listAPPA, [('roi_file', 'in2')]),
(listAPPA, mergeAPPA, [('out', 'in_files')]),
# Execute TOPUP for susceptibility-induced distortion correction
(mergeAPPA, topup, [('merged_file', 'in_file')]),
(infosource, topup, [('acqp_file', 'encoding_file')]),
# ---------------------- EDDY (motion & eddy current corrections)
# Use field coefficients and movement parameters from topup for eddy
(topup, eddy, [('out_fieldcoef', 'in_topup_fieldcoef'),
('out_movpar', 'in_topup_movpar')]),
# Provide initial brain mask for eddy
(betMask1, eddy, [('mask_file', 'in_mask')]),
# Provide b-vectors for eddy correction
(datasource, eddy, [('bvec_file', 'in_bvec')]),
# Provide denoised, degibbsed DWI data for eddy
(mrtrixDegibbs, eddy, [('out_file', 'in_file')]),
# Provide phase-encode information for eddy
(infosource, eddy, [('acqp_file', 'in_acqp')]),
# Adjust b-values before using them in eddy
(datasource, adjustBval, [('bval_file', 'in_bval')]),
# Create an index file for eddy
(datasource, eddyIndex, [('bval_file', 'in_bval')]),
# Provide adjusted b-values to eddy
(adjustBval, eddy, [('out_bval', 'in_bval')]),
# Provide index file to eddy
(eddyIndex, eddy, [('out_file', 'in_index')]),
# ---------------------- BRAIN MASK 2 (post-topup/eddy)
# Create a mask with mrtrix based on corrected DWI data
(eddy, mrtrixMask, [('out_corrected', 'in_file')]),
# Use rotated b-vectors for mrtrix mask generation
(eddy, mrtrixMask, [('out_rotated_bvecs', 'in_bvec')]),
# Provide b-values for mask generation with mrtrix
(datasource, mrtrixMask, [('bval_file', 'in_bval')]),
# Create a brain mask with BET based on eddy corrected data
(eddy, betMask2, [('out_corrected', 'in_file')]),
# Combine masks from mrtrix and BET
(mrtrixMask, dwiMask, [('out_mask', 'in_mask1')]),
(betMask2, dwiMask, [('mask_file', 'in_mask2')]),
# ---------------------- GRADIENT CHECK (ensuring consistency in b-values/vectors)
# Check gradient consistency of the eddy-corrected DWI
(eddy, mrtrixGradCheck, [('out_corrected', 'in_file')]),
# Use rotated b-vectors for gradient consistency check
(eddy, mrtrixGradCheck, [('out_rotated_bvecs', 'in_bvecs')]),
# Provide original b-values for gradient consistency check
(datasource, mrtrixGradCheck, [('bval_file', 'in_bvals')]),
# ---------------------- DATASINK (saving results)
# Save the final DWI brain mask
(dwiMask, datasink, [('out_mask', 'dwi.@dwi_mask')]),
# Save the eddy-corrected DWI
(eddy, datasink, [('out_corrected', 'dwi.@eddy_corrected')]),
# Save the checked b-values post gradient check
(mrtrixGradCheck, datasink, [('out_bvals', 'dwi.@bvals')]),
# Save the checked b-vectors post gradient check
(mrtrixGradCheck, datasink, [('out_bvecs', 'dwi.@bvecs')]),
])
# Run the script and generate a graph of the workflow
if __name__ == '__main__':
workflow.write_graph(graph2use='orig')
workflow.run(plugin=processing_type, plugin_args=plugin_args)
# %%