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