forked from brainlife/app-dipy-afq
/
main.py
94 lines (76 loc) · 2.7 KB
/
main.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
import os.path as op
import numpy as np
import nibabel as nib
import dipy.data as dpd
import json
from dipy.data import fetcher
from dipy.io.gradients import read_bvals_bvecs
from dipy.core.gradients import gradient_table
import AFQ.utils.streamlines as aus
import AFQ.data as afd
import AFQ.tractography as aft
import AFQ.registration as reg
import AFQ.dti as dti
import AFQ.segmentation as seg
import os
def main():
with open('config.json') as config_json:
config = json.load(config_json)
#Paths to data
data_file = str(config['data_file'])
data_bval = str(config['data_bval'])
data_bvec = str(config['data_bvec'])
img = nib.load(data_file)
print('Loaded Data')
print('Calculating DTI')
if not op.exists('./dti_FA.nii.gz'):
dti_params = dti.fit_dti(data_file, data_bval, data_bvec,out_dir='.')
else:
dti_params = {'FA': './dti_FA.nii.gz','params': './dti_params.nii.gz'}
tg = nib.streamlines.load(str(config['tck_data'])).tractogram
#cannot remove inv, affine
streamlines = tg.apply_affine(np.linalg.inv(img.affine)).streamlines
#streamlines = tg.streamlines
print('Loaded streamlines')
# Use only a small portion of the streamlines, for expedience:
streamlines = streamlines[::100]
templates = afd.read_templates()
bundle_names = ["CST", "ILF"]
bundles = {}
for name in bundle_names:
for hemi in ['_R', '_L']:
bundles[name + hemi] = {'ROIs': [templates[name + '_roi1' + hemi],
templates[name + '_roi1' + hemi]],
'rules': [True, True]}
print('Set Bundles')
MNI_T2_img = dpd.read_mni_template()
print("Registering to template...")
bvals, bvecs = read_bvals_bvecs(data_bval, data_bvec)
if not op.exists('mapping.nii.gz'):
#bvals, bvecs = read_bvals_bvecs(data_bval, data_bvec)
gtab = gradient_table(bvals, bvecs)
mapping = reg.syn_register_dwi(data_file, gtab)
reg.write_mapping(mapping, './mapping.nii.gz')
else:
mapping = reg.read_mapping('./mapping.nii.gz', img, MNI_T2_img)
print("Segmenting fiber groups...")
fiber_groups = seg.segment(data_file,
data_bval,
data_bvec,
streamlines,
bundles,
reg_template=MNI_T2_img,
mapping=mapping,
as_generator=False,
affine=img.affine)
path = os.getcwd() + '/tract1/'
if not os.path.exists(path):
os.makedirs(path)
print('Creating tck files')
for fg in fiber_groups:
streamlines = fiber_groups[fg]
fname = fg + ".tck"
trg = nib.streamlines.Tractogram(streamlines, affine_to_rasmm=img.affine)
nib.streamlines.save(trg,path+fname)
print('Finished segment')
main()