annotate create_streamlines.py @ 12:883327c42970 draft

Uploaded
author greg
date Thu, 30 Nov 2017 10:11:00 -0500
parents 1f1fdfe9ac4d
children ab72c2d00e5a
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
12
883327c42970 Uploaded
greg
parents: 8
diff changeset
3 import os
0
1c5508f627ec Uploaded
greg
parents:
diff changeset
4 import shutil
1c5508f627ec Uploaded
greg
parents:
diff changeset
5
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()
12
883327c42970 Uploaded
greg
parents: 8
diff changeset
19 parser.add_argument('--input', dest='input', help='Input dataset')
883327c42970 Uploaded
greg
parents: 8
diff changeset
20 parser.add_argument('--input_extra_files_path', dest='input_extra_files_path', help='Input dataset extra files paths')
0
1c5508f627ec Uploaded
greg
parents:
diff changeset
21 parser.add_argument('--output_superiorfrontal_nifti', dest='output_superiorfrontal_nifti', help='Output superiorfrontal nifti1 dataset')
12
883327c42970 Uploaded
greg
parents: 8
diff changeset
22 parser.add_argument('--output_superiorfrontal_nifti_files_path', dest='output_superiorfrontal_nifti_files_path', help='Output superiorfrontal nifti1 extra files path')
0
1c5508f627ec Uploaded
greg
parents:
diff changeset
23 parser.add_argument('--output_trackvis_header', dest='output_trackvis_header', help='Output superiorfrontal track visualization header dataset')
1c5508f627ec Uploaded
greg
parents:
diff changeset
24
1c5508f627ec Uploaded
greg
parents:
diff changeset
25 args = parser.parse_args()
1c5508f627ec Uploaded
greg
parents:
diff changeset
26
12
883327c42970 Uploaded
greg
parents: 8
diff changeset
27 def move_directory_files(source_dir, destination_dir, copy=False, remove_source_dir=False):
883327c42970 Uploaded
greg
parents: 8
diff changeset
28 source_directory = os.path.abspath(source_dir)
883327c42970 Uploaded
greg
parents: 8
diff changeset
29 destination_directory = os.path.abspath(destination_dir)
883327c42970 Uploaded
greg
parents: 8
diff changeset
30 if not os.path.isdir(destination_directory):
883327c42970 Uploaded
greg
parents: 8
diff changeset
31 os.makedirs(destination_directory)
883327c42970 Uploaded
greg
parents: 8
diff changeset
32 for dir_entry in os.listdir(source_directory):
883327c42970 Uploaded
greg
parents: 8
diff changeset
33 source_entry = os.path.join(source_directory, dir_entry)
883327c42970 Uploaded
greg
parents: 8
diff changeset
34 if copy:
883327c42970 Uploaded
greg
parents: 8
diff changeset
35 shutil.copy(source_entry, destination_directory)
883327c42970 Uploaded
greg
parents: 8
diff changeset
36 else:
883327c42970 Uploaded
greg
parents: 8
diff changeset
37 shutil.move(source_entry, destination_directory)
883327c42970 Uploaded
greg
parents: 8
diff changeset
38 if remove_source_dir:
883327c42970 Uploaded
greg
parents: 8
diff changeset
39 os.rmdir(source_directory)
883327c42970 Uploaded
greg
parents: 8
diff changeset
40
883327c42970 Uploaded
greg
parents: 8
diff changeset
41 # Get input data.
883327c42970 Uploaded
greg
parents: 8
diff changeset
42 # TODO: do not hard-code 'stanford_hardi'
883327c42970 Uploaded
greg
parents: 8
diff changeset
43 input_dir = 'stanford_hardi'
883327c42970 Uploaded
greg
parents: 8
diff changeset
44 os.mkdir(input_dir)
883327c42970 Uploaded
greg
parents: 8
diff changeset
45 for f in os.listdir(args.input_extra_files_path):
883327c42970 Uploaded
greg
parents: 8
diff changeset
46 shutil.copy(os.path.join(args.input_extra_files_path, f), input_dir)
883327c42970 Uploaded
greg
parents: 8
diff changeset
47 hard_img, gtab = read_stanford_hardi()
0
1c5508f627ec Uploaded
greg
parents:
diff changeset
48 data = hardi_img.get_data()
12
883327c42970 Uploaded
greg
parents: 8
diff changeset
49 labels_file = os.path.join(input_dir, "aparc-reduced.nii.gz")
883327c42970 Uploaded
greg
parents: 8
diff changeset
50 labels_img = nib.load(labels_file)
0
1c5508f627ec Uploaded
greg
parents:
diff changeset
51 labels = labels_img.get_data()
1c5508f627ec Uploaded
greg
parents:
diff changeset
52
1c5508f627ec Uploaded
greg
parents:
diff changeset
53 white_matter = (labels == 1) | (labels == 2)
1c5508f627ec Uploaded
greg
parents:
diff changeset
54 csamodel = shm.CsaOdfModel(gtab, 6)
1c5508f627ec Uploaded
greg
parents:
diff changeset
55 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
56 seeds = utils.seeds_from_mask(white_matter, density=2)
1c5508f627ec Uploaded
greg
parents:
diff changeset
57 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
58 affine = streamline_generator.affine
1c5508f627ec Uploaded
greg
parents:
diff changeset
59 streamlines = list(streamline_generator)
1c5508f627ec Uploaded
greg
parents:
diff changeset
60 cc_slice = labels == 2
1c5508f627ec Uploaded
greg
parents:
diff changeset
61 cc_streamlines = utils.target(streamlines, cc_slice, affine=affine)
1c5508f627ec Uploaded
greg
parents:
diff changeset
62 cc_streamlines = list(cc_streamlines)
12
883327c42970 Uploaded
greg
parents: 8
diff changeset
63
0
1c5508f627ec Uploaded
greg
parents:
diff changeset
64 lr_superiorfrontal_track = grouping[11, 54]
1c5508f627ec Uploaded
greg
parents:
diff changeset
65 shape = labels.shape
1c5508f627ec Uploaded
greg
parents:
diff changeset
66 dm = utils.density_map(lr_superiorfrontal_track, shape, affine=affine)
1c5508f627ec Uploaded
greg
parents:
diff changeset
67 # Save density map
1c5508f627ec Uploaded
greg
parents:
diff changeset
68 dm_img = nib.Nifti1Image(dm.astype("int16"), hardi_img.affine)
1c5508f627ec Uploaded
greg
parents:
diff changeset
69 dm_img.to_filename("lr-superiorfrontal-dm.nii")
1c5508f627ec Uploaded
greg
parents:
diff changeset
70 shutil.move('lr-superiorfrontal-dm.nii', args.output_superiorfrontal_nifti)
12
883327c42970 Uploaded
greg
parents: 8
diff changeset
71 move_directory_files(input_dir, args.output_superiorfrontal_nifti_files_path)
883327c42970 Uploaded
greg
parents: 8
diff changeset
72
0
1c5508f627ec Uploaded
greg
parents:
diff changeset
73 # Make a trackvis header so we can save streamlines
1c5508f627ec Uploaded
greg
parents:
diff changeset
74 voxel_size = labels_img.header.get_zooms()
1c5508f627ec Uploaded
greg
parents:
diff changeset
75 trackvis_header = nib.trackvis.empty_header()
1c5508f627ec Uploaded
greg
parents:
diff changeset
76 trackvis_header['voxel_size'] = voxel_size
1c5508f627ec Uploaded
greg
parents:
diff changeset
77 trackvis_header['dim'] = shape
1c5508f627ec Uploaded
greg
parents:
diff changeset
78 trackvis_header['voxel_order'] = "RAS"
1c5508f627ec Uploaded
greg
parents:
diff changeset
79 # Move streamlines to "trackvis space"
1c5508f627ec Uploaded
greg
parents:
diff changeset
80 trackvis_point_space = utils.affine_for_trackvis(voxel_size)
1c5508f627ec Uploaded
greg
parents:
diff changeset
81 lr_sf_trk = utils.move_streamlines(lr_superiorfrontal_track, trackvis_point_space, input_space=affine)
1c5508f627ec Uploaded
greg
parents:
diff changeset
82 lr_sf_trk = list(lr_sf_trk)
1c5508f627ec Uploaded
greg
parents:
diff changeset
83 # Save streamlines
1c5508f627ec Uploaded
greg
parents:
diff changeset
84 for_save = [(sl, None, None) for sl in lr_sf_trk]
1c5508f627ec Uploaded
greg
parents:
diff changeset
85 nib.trackvis.write("lr-superiorfrontal.trk", for_save, trackvis_header)
1c5508f627ec Uploaded
greg
parents:
diff changeset
86 shutil.move('lr-superiorfrontal.trk', args.output_trackvis_header)