forked from brainlife/app-tract-profile
/
main.py
42 lines (36 loc) · 1.38 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
import AFQ.utils.streamlines as aus
import AFQ.dti as dti
import nibabel as nib
import codecs
import json
import os
import AFQ.segmentation as seg
def main():
with open('config.json') as config_json:
config = json.load(config_json)
data_file = str(config['data_file'])
data_bval = str(config['data_bval'])
data_bvec = str(config['data_bvec'])
tracks = str(config['tracks'])
print("Calculating DTI...")
if not os.path.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'}
FA_img = nib.load(dti_params['FA'])
FA_data = FA_img.get_data()
print("Extracting tract profiles...")
path = os.getcwd() + '/profile/'
if not os.path.exists(path):
os.makedirs(path)
for t in os.listdir(tracks):
if t.endswith('.tck'):
tg = nib.streamlines.load(tracks+'/'+t)
streamlines = list(tg.streamlines)
profile = seg.calculate_tract_profile(FA_data, streamlines)
profile = profile.tolist()
t = os.path.splitext(os.path.basename(t))[0] #remove the .tck from string
p = path+'/'+t+'.json'
json.dump(profile, codecs.open(p, 'w', encoding='utf-8'), separators=(',', ':'), sort_keys=True, indent=4)
main()