annotate create_streamlines.py @ 5:c55029873167 draft

Uploaded
author greg
date Wed, 29 Nov 2017 09:50:42 -0500
parents 3b5e09c727f5
children 606a712b6322
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
rev   line source
0
1c5508f627ec Uploaded
greg
parents:
diff changeset
1 #!/usr/bin/env python
1c5508f627ec Uploaded
greg
parents:
diff changeset
2 import argparse
1c5508f627ec Uploaded
greg
parents:
diff changeset
3 import shutil
1c5508f627ec Uploaded
greg
parents:
diff changeset
4
1c5508f627ec Uploaded
greg
parents:
diff changeset
5 from dipy.data import fetch_stanford_t1, read_stanford_labels, read_stanford_t1
1c5508f627ec Uploaded
greg
parents:
diff changeset
6 from dipy.reconst import peaks, shm
1c5508f627ec Uploaded
greg
parents:
diff changeset
7 from dipy.tracking import utils
1c5508f627ec Uploaded
greg
parents:
diff changeset
8 from dipy.tracking.eudx import EuDX
1c5508f627ec Uploaded
greg
parents:
diff changeset
9 from dipy.viz import fvtk
1c5508f627ec Uploaded
greg
parents:
diff changeset
10 from dipy.viz.colormap import line_colors
1c5508f627ec Uploaded
greg
parents:
diff changeset
11
1c5508f627ec Uploaded
greg
parents:
diff changeset
12 import matplotlib.pyplot as plt
1c5508f627ec Uploaded
greg
parents:
diff changeset
13
1c5508f627ec Uploaded
greg
parents:
diff changeset
14 import nibabel as nib
1c5508f627ec Uploaded
greg
parents:
diff changeset
15
1c5508f627ec Uploaded
greg
parents:
diff changeset
16 import numpy as np
1c5508f627ec Uploaded
greg
parents:
diff changeset
17
1c5508f627ec Uploaded
greg
parents:
diff changeset
18 parser = argparse.ArgumentParser()
1c5508f627ec Uploaded
greg
parents:
diff changeset
19 parser.add_argument('--drmi_dataset', dest='drmi_dataset', help='Input dataset')
1c5508f627ec Uploaded
greg
parents:
diff changeset
20 parser.add_argument('--output_corpuscallosum_axial', dest='output_corpuscallosum_axial', help='Output corpuscallosum axial dataset')
1c5508f627ec Uploaded
greg
parents:
diff changeset
21 parser.add_argument('--output_corpuscallosum_sagittal', dest='output_corpuscallosum_sagittal', help='Output corpuscallosum sagittal dataset')
1c5508f627ec Uploaded
greg
parents:
diff changeset
22 parser.add_argument('--output_connectivity', dest='output_connectivity', help='Output connectivity dataset')
1c5508f627ec Uploaded
greg
parents:
diff changeset
23 parser.add_argument('--output_superiorfrontal_nifti', dest='output_superiorfrontal_nifti', help='Output superiorfrontal nifti1 dataset')
1c5508f627ec Uploaded
greg
parents:
diff changeset
24 parser.add_argument('--output_trackvis_header', dest='output_trackvis_header', help='Output superiorfrontal track visualization header dataset')
1c5508f627ec Uploaded
greg
parents:
diff changeset
25
1c5508f627ec Uploaded
greg
parents:
diff changeset
26 args = parser.parse_args()
1c5508f627ec Uploaded
greg
parents:
diff changeset
27
1c5508f627ec Uploaded
greg
parents:
diff changeset
28 hardi_img, gtab, labels_img = read_stanford_labels()
1c5508f627ec Uploaded
greg
parents:
diff changeset
29 data = hardi_img.get_data()
1c5508f627ec Uploaded
greg
parents:
diff changeset
30 labels = labels_img.get_data()
1c5508f627ec Uploaded
greg
parents:
diff changeset
31
3
51263bfe7b2c Uploaded
greg
parents: 0
diff changeset
32 # For possible later use: if args.drmi_dataset == 'stanford_t1':
51263bfe7b2c Uploaded
greg
parents: 0
diff changeset
33 fetch_stanford_t1()
51263bfe7b2c Uploaded
greg
parents: 0
diff changeset
34 t1 = read_stanford_t1()
0
1c5508f627ec Uploaded
greg
parents:
diff changeset
35 t1_data = t1.get_data()
1c5508f627ec Uploaded
greg
parents:
diff changeset
36 white_matter = (labels == 1) | (labels == 2)
1c5508f627ec Uploaded
greg
parents:
diff changeset
37 csamodel = shm.CsaOdfModel(gtab, 6)
1c5508f627ec Uploaded
greg
parents:
diff changeset
38 csapeaks = peaks.peaks_from_model(model=csamodel, data=data, sphere=peaks.default_sphere, relative_peak_threshold=.8, min_separation_angle=45, mask=white_matter)
1c5508f627ec Uploaded
greg
parents:
diff changeset
39 seeds = utils.seeds_from_mask(white_matter, density=2)
1c5508f627ec Uploaded
greg
parents:
diff changeset
40 streamline_generator = EuDX(csapeaks.peak_values, csapeaks.peak_indices, odf_vertices=peaks.default_sphere.vertices, a_low=.05, step_sz=.5, seeds=seeds)
1c5508f627ec Uploaded
greg
parents:
diff changeset
41 affine = streamline_generator.affine
1c5508f627ec Uploaded
greg
parents:
diff changeset
42 streamlines = list(streamline_generator)
1c5508f627ec Uploaded
greg
parents:
diff changeset
43 cc_slice = labels == 2
1c5508f627ec Uploaded
greg
parents:
diff changeset
44 cc_streamlines = utils.target(streamlines, cc_slice, affine=affine)
1c5508f627ec Uploaded
greg
parents:
diff changeset
45 cc_streamlines = list(cc_streamlines)
1c5508f627ec Uploaded
greg
parents:
diff changeset
46 other_streamlines = utils.target(streamlines, cc_slice, affine=affine, include=False)
1c5508f627ec Uploaded
greg
parents:
diff changeset
47 other_streamlines = list(other_streamlines)
1c5508f627ec Uploaded
greg
parents:
diff changeset
48 assert len(other_streamlines) + len(cc_streamlines) == len(streamlines)
1c5508f627ec Uploaded
greg
parents:
diff changeset
49 # Make display objects
1c5508f627ec Uploaded
greg
parents:
diff changeset
50 color = line_colors(cc_streamlines)
1c5508f627ec Uploaded
greg
parents:
diff changeset
51 cc_streamlines_actor = fvtk.line(cc_streamlines, line_colors(cc_streamlines))
1c5508f627ec Uploaded
greg
parents:
diff changeset
52 cc_ROI_actor = fvtk.contour(cc_slice, levels=[1], colors=[(1., 1., 0.)], opacities=[1.])
1c5508f627ec Uploaded
greg
parents:
diff changeset
53 vol_actor = fvtk.slicer(t1_data)
1c5508f627ec Uploaded
greg
parents:
diff changeset
54 vol_actor.display(40, None, None)
1c5508f627ec Uploaded
greg
parents:
diff changeset
55 vol_actor2 = vol_actor.copy()
1c5508f627ec Uploaded
greg
parents:
diff changeset
56 vol_actor2.display(None, None, 35)
1c5508f627ec Uploaded
greg
parents:
diff changeset
57 # Add display objects to canvas
1c5508f627ec Uploaded
greg
parents:
diff changeset
58 r = fvtk.ren()
1c5508f627ec Uploaded
greg
parents:
diff changeset
59 fvtk.add(r, vol_actor)
1c5508f627ec Uploaded
greg
parents:
diff changeset
60 fvtk.add(r, vol_actor2)
1c5508f627ec Uploaded
greg
parents:
diff changeset
61 fvtk.add(r, cc_streamlines_actor)
1c5508f627ec Uploaded
greg
parents:
diff changeset
62 fvtk.add(r, cc_ROI_actor)
1c5508f627ec Uploaded
greg
parents:
diff changeset
63 # Save figures
4
3b5e09c727f5 Uploaded
greg
parents: 3
diff changeset
64 fvtk.record(r, n_frames=1, out_path="corpuscallosum_axial.png", size=(800, 800))
3b5e09c727f5 Uploaded
greg
parents: 3
diff changeset
65 shutil.move("corpuscallosum_axial.png", args.output_corpuscallosum_axial)
0
1c5508f627ec Uploaded
greg
parents:
diff changeset
66 fvtk.camera(r, [-1, 0, 0], [0, 0, 0], viewup=[0, 0, 1])
4
3b5e09c727f5 Uploaded
greg
parents: 3
diff changeset
67 fvtk.record(r, n_frames=1, out_path="corpuscallosum_sagittal.png", size=(800, 800))
3b5e09c727f5 Uploaded
greg
parents: 3
diff changeset
68 shutil.move("corpuscallosum_sagittal.png", args.output_corpuscallosum_sagittal)
0
1c5508f627ec Uploaded
greg
parents:
diff changeset
69 M, grouping = utils.connectivity_matrix(cc_streamlines, labels, affine=affine, return_mapping=True, mapping_as_streamlines=True)
1c5508f627ec Uploaded
greg
parents:
diff changeset
70 M[:3, :] = 0
1c5508f627ec Uploaded
greg
parents:
diff changeset
71 M[:, :3] = 0
1c5508f627ec Uploaded
greg
parents:
diff changeset
72 plt.imshow(np.log1p(M), interpolation='nearest')
4
3b5e09c727f5 Uploaded
greg
parents: 3
diff changeset
73 plt.savefig("connectivity.png")
3b5e09c727f5 Uploaded
greg
parents: 3
diff changeset
74 shutil.move("connectivity.png", args.output_connectivity)
0
1c5508f627ec Uploaded
greg
parents:
diff changeset
75 lr_superiorfrontal_track = grouping[11, 54]
1c5508f627ec Uploaded
greg
parents:
diff changeset
76 shape = labels.shape
1c5508f627ec Uploaded
greg
parents:
diff changeset
77 dm = utils.density_map(lr_superiorfrontal_track, shape, affine=affine)
1c5508f627ec Uploaded
greg
parents:
diff changeset
78 # Save density map
1c5508f627ec Uploaded
greg
parents:
diff changeset
79 dm_img = nib.Nifti1Image(dm.astype("int16"), hardi_img.affine)
1c5508f627ec Uploaded
greg
parents:
diff changeset
80 dm_img.to_filename("lr-superiorfrontal-dm.nii")
1c5508f627ec Uploaded
greg
parents:
diff changeset
81 shutil.move('lr-superiorfrontal-dm.nii', args.output_superiorfrontal_nifti)
1c5508f627ec Uploaded
greg
parents:
diff changeset
82
1c5508f627ec Uploaded
greg
parents:
diff changeset
83 # Make a trackvis header so we can save streamlines
1c5508f627ec Uploaded
greg
parents:
diff changeset
84 voxel_size = labels_img.header.get_zooms()
1c5508f627ec Uploaded
greg
parents:
diff changeset
85 trackvis_header = nib.trackvis.empty_header()
1c5508f627ec Uploaded
greg
parents:
diff changeset
86 trackvis_header['voxel_size'] = voxel_size
1c5508f627ec Uploaded
greg
parents:
diff changeset
87 trackvis_header['dim'] = shape
1c5508f627ec Uploaded
greg
parents:
diff changeset
88 trackvis_header['voxel_order'] = "RAS"
1c5508f627ec Uploaded
greg
parents:
diff changeset
89
1c5508f627ec Uploaded
greg
parents:
diff changeset
90 # Move streamlines to "trackvis space"
1c5508f627ec Uploaded
greg
parents:
diff changeset
91 trackvis_point_space = utils.affine_for_trackvis(voxel_size)
1c5508f627ec Uploaded
greg
parents:
diff changeset
92 lr_sf_trk = utils.move_streamlines(lr_superiorfrontal_track, trackvis_point_space, input_space=affine)
1c5508f627ec Uploaded
greg
parents:
diff changeset
93 lr_sf_trk = list(lr_sf_trk)
1c5508f627ec Uploaded
greg
parents:
diff changeset
94
1c5508f627ec Uploaded
greg
parents:
diff changeset
95 # Save streamlines
1c5508f627ec Uploaded
greg
parents:
diff changeset
96 for_save = [(sl, None, None) for sl in lr_sf_trk]
1c5508f627ec Uploaded
greg
parents:
diff changeset
97 nib.trackvis.write("lr-superiorfrontal.trk", for_save, trackvis_header)
1c5508f627ec Uploaded
greg
parents:
diff changeset
98 shutil.move('lr-superiorfrontal.trk', args.output_trackvis_header)
1c5508f627ec Uploaded
greg
parents:
diff changeset
99 dm_trackvis = utils.density_map(lr_sf_trk, shape, affine=trackvis_point_space)
1c5508f627ec Uploaded
greg
parents:
diff changeset
100 assert np.all(dm == dm_trackvis)