comparison fast_fiber_tracking.py @ 0:4e3d4331fa58 draft

Uploaded
author greg
date Tue, 07 Nov 2017 13:48:48 -0500
parents
children 63aa79b5ebd6
comparison
equal deleted inserted replaced
-1:000000000000 0:4e3d4331fa58
1 #!/usr/bin/env python
2 import argparse
3 import shutil
4
5 from dipy.data import fetch_sherbrooke_3shell
6 from dipy.data import fetch_stanford_hardi
7 from dipy.data import get_sphere
8 from dipy.data import read_sherbrooke_3shell
9 from dipy.data import read_stanford_hardi
10 from dipy.direction import peaks_from_model
11 from dipy.io.image import save_nifti
12 from dipy.reconst.csdeconv import (ConstrainedSphericalDeconvModel, auto_response)
13 from dipy.reconst.dti import TensorModel
14 from dipy.segment.mask import median_otsu
15 from dipy.tracking.local import LocalTracking, ThresholdTissueClassifier
16 from dipy.tracking.streamline import Streamlines
17 from dipy.tracking.utils import random_seeds_from_mask
18 from dipy.viz import actor, window
19
20 import numpy as np
21
22 parser = argparse.ArgumentParser()
23 parser.add_argument('--drmi_dataset', dest='drmi_dataset', help='Input dataset')
24 parser.add_argument('--output_csd_direction_field', dest='output_csd_direction_field', help='Output csd direction field dataset')
25 parser.add_argument('--output_det_streamlines', dest='output_det_streamlines', help='Output det streamlines dataset')
26 parser.add_argument('--output_fa_map', dest='output_fa_map', help='Output fa map dataset')
27
28 args = parser.parse_args()
29
30 interactive = False
31
32 # Get input data.
33 input_dir = args.drmi_dataset
34 if input_dir == 'sherbrooke_3shell':
35 fetch_sherbrooke_3shell()
36 img, gtab = read_sherbrooke_3shell()
37 elif input_dir == 'stanford_hardi':
38 fetch_stanford_hardi()
39 img, gtab = read_stanford_hardi()
40
41 data = img.get_data()
42 maskdata, mask = median_otsu(data, 3, 1, False, vol_idx=range(10, 50), dilate=2)
43
44 response, ratio = auto_response(gtab, data, roi_radius=10, fa_thr=0.7)
45 csd_model = ConstrainedSphericalDeconvModel(gtab, response)
46 sphere = get_sphere('symmetric724')
47 csd_peaks = peaks_from_model(model=csd_model, data=data, sphere=sphere, mask=mask, relative_peak_threshold=.5, min_separation_angle=25, parallel=True)
48
49 tensor_model = TensorModel(gtab, fit_method='WLS')
50 tensor_fit = tensor_model.fit(data, mask)
51 fa = tensor_fit.fa
52
53 tissue_classifier = ThresholdTissueClassifier(fa, 0.1)
54 seeds = random_seeds_from_mask(fa > 0.3, seeds_count=1)
55
56 ren = window.Renderer()
57 ren.add(actor.peak_slicer(csd_peaks.peak_dirs, csd_peaks.peak_values, colors=None))
58 window.record(ren, out_path='csd_direction_field.png', size=(900, 900))
59 shutil.move('csd_direction_field.png', args.output_csd_direction_field)
60
61 streamline_generator = LocalTracking(csd_peaks, tissue_classifier, seeds, affine=np.eye(4), step_size=0.5)
62 streamlines = Streamlines(streamline_generator)
63
64 ren.clear()
65 ren.add(actor.line(streamlines))
66 window.record(ren, out_path='det_streamlines.png', size=(900, 900))
67 shutil.move('det_streamlines.png', args.output_det_streamlines)
68
69 save_nifti('fa_map.nii', fa, img.affine)
70 shutil.move('fa_map.nii', args.output_fa_map)