Mercurial > repos > greg > diffusion_signal_reconstruction

comparison diffusion_signal_reconstruction.py @ 0:fe569aad237c draft

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Uploaded
author greg
date Sun, 05 Nov 2017 09:30:58 -0500
parents
children 85df19d98cd0
comparison
equal deleted inserted replaced
-1:000000000000 0:fe569aad237c
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.reconst import dti
11 from dipy.reconst.dti import color_fa
12 from dipy.reconst.dti import fractional_anisotropy
13 from dipy.segment.mask import median_otsu
14 from dipy.viz import fvtk
15 from matplotlib import pyplot
16
17 import nibabel
18 import numpy
19
20 # http://nipy.org/dipy/examples_built/reconst_dti.html#example-reconst-dti
21 parser = argparse.ArgumentParser()
22 parser.add_argument('--drmi_dataset', dest='drmi_dataset', help='Input dataset')
23 parser.add_argument('--output_nifti1_fa', dest='output_nifti1_fa', help='Output fractional anisotropy Nifti1 dataset')
24 parser.add_argument('--output_nifti1_evecs', dest='output_nifti1_fevecs', help='Output eigen vectors Nifti1 dataset')
25 parser.add_argument('--output_nifti1_md', dest='output_nifti1_md', help='Output mean diffusivity Nifti1 dataset')
26 parser.add_argument('--output_nifti1_rgb', dest='output_nifti1_rgb', help='Output RGB-map Nifti1 dataset')
27 parser.add_argument('--output_png_ellipsoids', dest='output_png_ellipsoids', help='Output ellipsoids PNG dataset')
28 parser.add_argument('--output_png_odfs', dest='output_png_odfs', help='Output orientation distribution functions PNG dataset')
29 parser.add_argument('--output_png_middle_axial_slice', dest='output_png_middle_axial_slice', help='Output middle axial slice PNG dataset')
30
31 args = parser.parse_args()
32
33 # Get input data.
34 input_dir = args.drmi_dataset
35 if input_dir == 'sherbrooke_3shell':
36 fetch_sherbrooke_3shell()
37 img, gtab = read_sherbrooke_3shell()
38 elif input_dir == 'stanford_hardi':
39 fetch_stanford_hardi()
40 img, gtab = read_stanford_hardi()
41
42 data = img.get_data()
43 maskdata, mask = median_otsu(data, 3, 1, True, vol_idx=range(10, 50), dilate=2)
44
45 axial_middle = data.shape[2] // 2
46 pyplot.subplot(1, 2, 1).set_axis_off()
47 pyplot.imshow(data[:, :, axial_middle, 0].T, cmap='gray', origin='lower')
48 pyplot.subplot(1, 2, 2).set_axis_off()
49 pyplot.imshow(data[:, :, axial_middle, 10].T, cmap='gray', origin='lower')
50 pyplot.savefig('middle_axial.png', bbox_inches='tight')
51 shutil.move('middle_axial.png', args.output_png_middle_axial_slice)
52
53 tenmodel = dti.TensorModel(gtab)
54 tenfit = tenmodel.fit(maskdata)
55
56 fa = fractional_anisotropy(tenfit.evals)
57 fa[numpy.isnan(fa)] = 0
58 fa_img = nibabel.Nifti1Image(fa.astype(numpy.float32), img.affine)
59 nibabel.save(fa_img, 'output_fa.nii')
60 shutil.move('output_fa.nii', args.output_nifti1_fa)
61
62 evecs_img = nibabel.Nifti1Image(tenfit.evecs.astype(numpy.float32), img.affine)
63 nibabel.save(evecs_img, 'output_evecs.nii')
64 shutil.move('output_evecs.nii', args.output_nifti1_evecs)
65
66 md1 = dti.mean_diffusivity(tenfit.evals)
67 nibabel.save(nibabel.Nifti1Image(md1.astype(numpy.float32), img.affine), 'output_md.nii')
68 shutil.move('output_md.nii', args.output_nifti1_md)
69
70 fa = numpy.clip(fa, 0, 1)
71 rgb = color_fa(fa, tenfit.evecs)
72 nibabel.save(nibabel.Nifti1Image(numpy.array(255 * rgb, 'uint8'), img.affine), 'output_rgb.nii')
73 shutil.move('output_rgb.nii', args.output_nifti1_rgb)
74
75 sphere = get_sphere('symmetric724')
76 ren = fvtk.ren()
77
78 evals = tenfit.evals[13:43, 44:74, 28:29]
79 evecs = tenfit.evecs[13:43, 44:74, 28:29]
80 cfa = rgb[13:43, 44:74, 28:29]
81 cfa /= cfa.max()
82 fvtk.add(ren, fvtk.tensor(evals, evecs, cfa, sphere))
83 fvtk.record(ren, n_frames=1, out_path='tensor_ellipsoids.png', size=(600, 600))
84 shutil.move('tensor_ellipsoids.png', args.output_png_ellipsoids)
85
86 fvtk.clear(ren)
87
88 tensor_odfs = tenmodel.fit(data[20:50, 55:85, 38:39]).odf(sphere)
89 fvtk.add(ren, fvtk.sphere_funcs(tensor_odfs, sphere, colormap=None))
90 fvtk.record(ren, n_frames=1, out_path='tensor_odfs.png', size=(600, 600))
91 shutil.move('tensor_odfs.png', args.output_png_odfs)