def test_read_write_trk():
sl = [
np.array([[0, 0, 0], [0, 0, 0.5], [0, 0, 1], [0, 0, 1.5]]),
np.array([[0, 0, 0], [0, 0.5, 0.5], [0, 1, 1]])
]
with nbtmp.InTemporaryDirectory() as tmpdir:
fname = op.join(tmpdir, 'sl.trk')
aus.write_trk(fname, sl)
new_sl = aus.read_trk(fname)
npt.assert_equal(list(new_sl), sl)
# What happens if this set of streamlines has some funky affine
# associated with it?
aff = np.eye(4) * np.random.rand()
aff[:3, 3] = np.array([1, 2, 3])
aff[3, 3] = 1
# We move the streamlines, and report the inverse of the affine:
aus.write_trk(fname,
move_streamlines(sl, aff),
affine=np.linalg.inv(aff))
# When we read this, we get back what we put in:
new_sl = aus.read_trk(fname)
# Compare each streamline:
for new, old in zip(new_sl, sl):
npt.assert_almost_equal(new, old, decimal=5)
def _streamlines(row,
wm_labels,
odf_model="DTI",
directions="det",
n_seeds=2,
random_seeds=False,
force_recompute=False,
wm_fa_thresh=0.2):
"""
wm_labels : list
The values within the segmentation that are considered white matter. We
will use this part of the image both to seed tracking (seeding
throughout), and for stopping.
"""
streamlines_file = _get_fname(
row, '%s_%s_streamlines.trk' % (odf_model, directions))
if not op.exists(streamlines_file) or force_recompute:
if odf_model == "DTI":
params_file = _dti(row)
elif odf_model == "CSD":
params_file = _csd(row)
dwi_img = nib.load(row['dwi_file'])
dwi_data = dwi_img.get_data()
if 'seg_file' in row.index:
# If we found a white matter segmentation in the
# expected location:
seg_img = nib.load(row['seg_file'])
seg_data_orig = seg_img.get_data()
# For different sets of labels, extract all the voxels that
# have any of these values:
wm_mask = np.sum(
np.concatenate([(seg_data_orig == l)[..., None]
for l in wm_labels], -1), -1)
# Resample to DWI data:
wm_mask = np.round(
reg.resample(wm_mask, dwi_data[..., 0], seg_img.affine,
dwi_img.affine)).astype(int)
else:
# Otherwise, we'll identify the white matter based on FA:
dti_fa = nib.load(_dti_fa(row)).get_data()
wm_mask = dti_fa > wm_fa_thresh
streamlines = aft.track(params_file,
directions=directions,
n_seeds=n_seeds,
random_seeds=random_seeds,
seed_mask=wm_mask,
stop_mask=wm_mask)
aus.write_trk(streamlines_file,
dtu.move_streamlines(streamlines,
np.linalg.inv(dwi_img.affine)),
affine=dwi_img.affine)
return streamlines_file
def _streamlines(row,
wm_labels,
odf_model="DTI",
directions="det",
force_recompute=False):
"""
wm_labels : list
The values within the segmentation that are considered white matter. We
will use this part of the image both to seed tracking (seeding
throughout), and for stopping.
"""
streamlines_file = _get_fname(
row, '%s_%s_streamlines.trk' % (odf_model, directions))
if not op.exists(streamlines_file) or force_recompute:
if odf_model == "DTI":
params_file = _dti(row)
else:
raise (NotImplementedError)
seg_img = nib.load(row['seg_file'])
dwi_img = nib.load(row['dwi_file'])
seg_data_orig = seg_img.get_data()
# For different sets of labels, extract all the voxels that have any
# of these values:
wm_mask = np.sum(
np.concatenate([(seg_data_orig == l)[..., None]
for l in wm_labels], -1), -1)
dwi_data = dwi_img.get_data()
resamp_wm = np.round(
reg.resample(wm_mask, dwi_data[..., 0], seg_img.affine,
dwi_img.affine)).astype(int)
streamlines = aft.track(params_file,
directions='det',
seeds=2,
seed_mask=resamp_wm,
stop_mask=resamp_wm)
aus.write_trk(streamlines_file, streamlines, affine=row['dwi_affine'])
return streamlines_file
def _streamlines(row, wm_labels, odf_model="DTI", directions="det",
force_recompute=False):
"""
wm_labels : list
The values within the segmentation that are considered white matter. We
will use this part of the image both to seed tracking (seeding
throughout), and for stopping.
"""
streamlines_file = _get_fname(row,
'%s_%s_streamlines.trk' % (odf_model,
directions))
if not op.exists(streamlines_file) or force_recompute:
if odf_model == "DTI":
params_file = _dti(row)
else:
raise(NotImplementedError)
seg_img = nib.load(row['seg_file'])
dwi_img = nib.load(row['dwi_file'])
seg_data_orig = seg_img.get_data()
# For different sets of labels, extract all the voxels that have any
# of these values:
wm_mask = np.sum(np.concatenate([(seg_data_orig == l)[..., None]
for l in wm_labels], -1), -1)
dwi_data = dwi_img.get_data()
resamp_wm = np.round(reg.resample(wm_mask, dwi_data[..., 0],
seg_img.affine,
dwi_img.affine)).astype(int)
streamlines = aft.track(params_file,
directions='det',
seeds=2,
seed_mask=resamp_wm,
stop_mask=resamp_wm)
aus.write_trk(streamlines_file, streamlines,
affine=row['dwi_affine'])
return streamlines_file
def _streamlines(row, odf_model="DTI", directions="det",
force_recompute=False):
streamlines_file = _get_fname(row,
'%s_%s_streamlines.trk' % (odf_model,
directions))
if not op.exists(streamlines_file) or force_recompute:
if odf_model == "DTI":
params_file = _dti(row)
else:
raise(NotImplementedError)
fa_file = _dti_fa(row)
fa = nib.load(fa_file).get_data()
wm_mask = np.zeros_like(fa)
wm_mask[fa > 0.2] = 1
streamlines = aft.track(params_file,
directions=directions,
seeds=1,
seed_mask=wm_mask,
stop_mask=fa)
aus.write_trk(streamlines_file, streamlines,
affine=row['dwi_affine'])
return streamlines_file
def test_read_write_trk():
sl = [np.array([[0, 0, 0], [0, 0, 0.5], [0, 0, 1], [0, 0, 1.5]]),
np.array([[0, 0, 0], [0, 0.5, 0.5], [0, 1, 1]])]
with nbtmp.InTemporaryDirectory() as tmpdir:
fname = op.join(tmpdir, 'sl.trk')
aus.write_trk(fname, sl)
new_sl = aus.read_trk(fname)
npt.assert_equal(list(new_sl), sl)
# What happens if this set of streamlines has some funky affine
# associated with it?
aff = np.eye(4) * np.random.rand()
aff[:3, 3] = np.array([1, 2, 3])
aff[3, 3] = 1
# We move the streamlines, and report the inverse of the affine:
aus.write_trk(fname, move_streamlines(sl, aff),
affine=np.linalg.inv(aff))
# When we read this, we get back what we put in:
new_sl = aus.read_trk(fname)
# Compare each streamline:
for new, old in zip(new_sl, sl):
npt.assert_almost_equal(new, old, decimal=4)
def _export_bundles(row,
wm_labels,
bundle_dict,
reg_template,
odf_model="DTI",
directions="det",
n_seeds=2,
random_seeds=False,
force_recompute=False):
for func, folder in zip([_clean_bundles, _bundles],
['clean_bundles', 'bundles']):
bundles_file = func(row,
wm_labels,
bundle_dict,
reg_template,
odf_model=odf_model,
directions=directions,
n_seeds=n_seeds,
random_seeds=random_seeds,
force_recompute=force_recompute)
bundles_dir = op.join(row['results_dir'], folder)
os.makedirs(bundles_dir, exist_ok=True)
trk = nib.streamlines.load(bundles_file)
tg = trk.tractogram
streamlines = tg.streamlines
for bundle in bundle_dict:
uid = bundle_dict[bundle]['uid']
idx = np.where(tg.data_per_streamline['bundle'] == uid)[0]
this_sl = (streamlines[idx])
fname = op.join(bundles_dir, '%s.trk' % bundle)
aus.write_trk(fname,
dtu.move_streamlines(
this_sl, np.linalg.inv(row['dwi_affine'])),
affine=row['dwi_affine'])
def test_streamline_registration():
sl1 = [
np.array([[0, 0, 0], [0, 0, 0.5], [0, 0, 1], [0, 0, 1.5]]),
np.array([[0, 0, 0], [0, 0.5, 0.5], [0, 1, 1]])
]
affine = np.eye(4)
affine[:3, 3] = np.random.randn(3)
sl2 = list(move_streamlines(sl1, affine))
aligned, matrix = streamline_registration(sl2, sl1)
npt.assert_almost_equal(matrix, np.linalg.inv(affine))
npt.assert_almost_equal(aligned[0], sl1[0])
npt.assert_almost_equal(aligned[1], sl1[1])
# We assume the two tracks come from the same space, but it might have
# some affine associated with it:
base_aff = np.eye(4) * np.random.rand()
base_aff[:3, 3] = np.array([1, 2, 3])
base_aff[3, 3] = 1
with nbtmp.InTemporaryDirectory() as tmpdir:
for use_aff in [None, base_aff]:
fname1 = op.join(tmpdir, 'sl1.trk')
fname2 = op.join(tmpdir, 'sl2.trk')
if use_aff is not None:
# Move the streamlines to this other space, and report it:
write_trk(fname1, move_streamlines(sl1,
np.linalg.inv(use_aff)),
use_aff)
write_trk(fname2, move_streamlines(sl2,
np.linalg.inv(use_aff)),
use_aff)
else:
write_trk(fname1, sl1)
write_trk(fname2, sl2)
aligned, matrix = streamline_registration(fname2, fname1)
npt.assert_almost_equal(aligned[0], sl1[0], decimal=5)
npt.assert_almost_equal(aligned[1], sl1[1], decimal=5)
def test_streamline_registration():
sl1 = [np.array([[0, 0, 0], [0, 0, 0.5], [0, 0, 1], [0, 0, 1.5]]),
np.array([[0, 0, 0], [0, 0.5, 0.5], [0, 1, 1]])]
affine = np.eye(4)
affine[:3, 3] = np.random.randn(3)
sl2 = list(move_streamlines(sl1, affine))
aligned, matrix = streamline_registration(sl2, sl1)
npt.assert_almost_equal(matrix, np.linalg.inv(affine))
npt.assert_almost_equal(aligned[0], sl1[0])
npt.assert_almost_equal(aligned[1], sl1[1])
# We assume the two tracks come from the same space, but it might have
# some affine associated with it:
base_aff = np.eye(4) * np.random.rand()
base_aff[:3, 3] = np.array([1, 2, 3])
base_aff[3, 3] = 1
with nbtmp.InTemporaryDirectory() as tmpdir:
for use_aff in [None, base_aff]:
fname1 = op.join(tmpdir, 'sl1.trk')
fname2 = op.join(tmpdir, 'sl2.trk')
if use_aff is not None:
# Move the streamlines to this other space, and report it:
write_trk(fname1,
move_streamlines(sl1, np.linalg.inv(use_aff)),
use_aff)
write_trk(fname2,
move_streamlines(sl2, np.linalg.inv(use_aff)),
use_aff)
else:
write_trk(fname1, sl1)
write_trk(fname2, sl2)
aligned, matrix = streamline_registration(fname2, fname1)
npt.assert_almost_equal(aligned[0], sl1[0], decimal=5)
npt.assert_almost_equal(aligned[1], sl1[1], decimal=5)
img = nib.load(hardi_fdata)
print("Calculating DTI...")
if not op.exists('./dti_FA.nii.gz'):
dti_params = dti.fit_dti(hardi_fdata,
hardi_fbval,
hardi_fbvec,
out_dir='.')
else:
dti_params = {'FA': './dti_FA.nii.gz', 'params': './dti_params.nii.gz'}
print("Tracking...")
if not op.exists('dti_streamlines.trk'):
streamlines = list(aft.track(dti_params['params']))
aus.write_trk('./dti_streamlines.trk', streamlines, affine=img.affine)
else:
tg = nib.streamlines.load('./dti_streamlines.trk').tractogram
streamlines = tg.apply_affine(np.linalg.inv(img.affine)).streamlines
# Use only a small portion of the streamlines, for expedience:
streamlines = streamlines[::100]
templates = afd.read_templates()
bundle_names = ["CST", "ILF"]
bundles = {}
for name in bundle_names:
for hemi in ['_R', '_L']:
bundles[name + hemi] = {
'ROIs': [
def test_AFQ_data2():
"""
Test with some actual data again, this time for track segmentation
"""
tmpdir = nbtmp.InTemporaryDirectory()
afd.organize_stanford_data(path=tmpdir.name)
preafq_path = op.join(tmpdir.name, 'stanford_hardi',
'derivatives', 'preafq')
myafq = api.AFQ(preafq_path=preafq_path,
sub_prefix='sub',
bundle_list=["SLF", "ARC", "CST", "FP"])
# Replace the mapping and streamlines with precomputed:
file_dict = afd.read_stanford_hardi_tractography()
mapping = file_dict['mapping.nii.gz']
streamlines = file_dict['tractography_subsampled.trk']
streamlines = dts.Streamlines(
dtu.move_streamlines([s for s in streamlines if s.shape[0] > 100],
np.linalg.inv(myafq.dwi_affine[0])))
sl_file = op.join(myafq.data_frame.results_dir[0],
'sub-01_sess-01_dwiDTI_det_streamlines.trk')
aus.write_trk(sl_file, streamlines, affine=myafq.dwi_affine[0])
mapping_file = op.join(myafq.data_frame.results_dir[0],
'sub-01_sess-01_dwi_mapping.nii.gz')
nib.save(mapping, mapping_file)
reg_prealign_file = op.join(myafq.data_frame.results_dir[0], 'sub-01_sess-01_dwi_reg_prealign.npy')
np.save(reg_prealign_file, np.eye(4))
tgram = nib.streamlines.load(myafq.bundles[0]).tractogram
bundles = aus.tgram_to_bundles(tgram, myafq.bundle_dict)
npt.assert_equal(len(bundles['CST_R']), 2)
# Test ROI exporting:
myafq.export_rois()
assert op.exists(op.join(myafq.data_frame['results_dir'][0],
'ROIs',
'CST_R_roi1_include.nii.gz'))
# Test bundles exporting:
myafq.export_bundles()
assert op.exists(op.join(myafq.data_frame['results_dir'][0],
'bundles',
'CST_R.trk'))
tract_profiles = pd.read_csv(myafq.tract_profiles[0])
assert tract_profiles.shape == (800, 5)
# Before we run the CLI, we'll remove the bundles and ROI folders, to see
# that the CLI generates them
shutil.rmtree(op.join(myafq.data_frame['results_dir'][0],
'bundles'))
shutil.rmtree(op.join(myafq.data_frame['results_dir'][0],
'ROIs'))
# Test the CLI:
print("Running the CLI:")
cmd = "pyAFQ " + preafq_path
out = os.system(cmd)
assert out == 0
# The combined tract profiles should already exist from the CLI Run:
from_file = pd.read_csv(op.join(myafq.afq_dir, 'tract_profiles.csv'))
# And should be identical to what we would get by rerunning this:
combined_profiles = myafq.combine_profiles()
assert combined_profiles.shape == (800, 7)
assert_frame_equal(combined_profiles, from_file)
# Make sure the CLI did indeed generate these:
assert op.exists(op.join(myafq.data_frame['results_dir'][0],
'ROIs',
'CST_R_roi1_include.nii.gz'))
assert op.exists(op.join(myafq.data_frame['results_dir'][0],
'bundles',
'CST_R.trk'))
def test_AFQ_data_planes():
"""
Test with some actual data again, this time for track segmentation
"""
tmpdir = nbtmp.InTemporaryDirectory()
afd.organize_stanford_data(path=tmpdir.name)
dmriprep_path = op.join(tmpdir.name, 'stanford_hardi', 'derivatives',
'dmriprep')
seg_algo = "planes"
bundle_names = ["SLF", "ARC", "CST", "FP"]
myafq = api.AFQ(dmriprep_path=dmriprep_path,
sub_prefix='sub',
seg_algo=seg_algo,
bundle_names=bundle_names,
odf_model="DTI")
# Replace the mapping and streamlines with precomputed:
file_dict = afd.read_stanford_hardi_tractography()
mapping = file_dict['mapping.nii.gz']
streamlines = file_dict['tractography_subsampled.trk']
streamlines = dts.Streamlines(
dtu.move_streamlines([s for s in streamlines if s.shape[0] > 100],
np.linalg.inv(myafq.dwi_affine[0])))
sl_file = op.join(myafq.data_frame.results_dir[0],
'sub-01_sess-01_dwiDTI_det_streamlines.trk')
aus.write_trk(sl_file, streamlines, affine=myafq.dwi_affine[0])
mapping_file = op.join(myafq.data_frame.results_dir[0],
'sub-01_sess-01_dwi_mapping.nii.gz')
nib.save(mapping, mapping_file)
reg_prealign_file = op.join(myafq.data_frame.results_dir[0],
'sub-01_sess-01_dwi_reg_prealign.npy')
np.save(reg_prealign_file, np.eye(4))
tgram = nib.streamlines.load(myafq.bundles[0]).tractogram
bundles = aus.tgram_to_bundles(tgram, myafq.bundle_dict)
npt.assert_(len(bundles['CST_L']) > 0)
# Test ROI exporting:
myafq.export_rois()
assert op.exists(
op.join(myafq.data_frame['results_dir'][0], 'ROIs',
'CST_R_roi1_include.nii.gz'))
# Test bundles exporting:
myafq.export_bundles()
assert op.exists(
op.join(myafq.data_frame['results_dir'][0], 'bundles', 'CST_R.trk'))
tract_profiles = pd.read_csv(myafq.tract_profiles[0])
assert tract_profiles.shape == (800, 5)
# Before we run the CLI, we'll remove the bundles and ROI folders, to see
# that the CLI generates them
shutil.rmtree(op.join(myafq.data_frame['results_dir'][0], 'bundles'))
shutil.rmtree(op.join(myafq.data_frame['results_dir'][0], 'ROIs'))
# Test the CLI:
print("Running the CLI:")
cmd = "pyAFQ " + dmriprep_path
out = os.system(cmd)
assert out == 0
# The combined tract profiles should already exist from the CLI Run:
from_file = pd.read_csv(op.join(myafq.afq_dir, 'tract_profiles.csv'))
# And should be identical to what we would get by rerunning this:
combined_profiles = myafq.combine_profiles()
assert combined_profiles.shape == (800, 7)
assert_frame_equal(combined_profiles, from_file)
# Make sure the CLI did indeed generate these:
assert op.exists(
op.join(myafq.data_frame['results_dir'][0], 'ROIs',
'CST_R_roi1_include.nii.gz'))
assert op.exists(
op.join(myafq.data_frame['results_dir'][0], 'bundles', 'CST_R.trk'))
# def test_AFQ_data_recobundles():
# tmpdir = nbtmp.InTemporaryDirectory()
# afd.fetch_hcp(["100206"], hcp_bucket='hcp-openaccess', profile_name="hcp",
# path=tmpdir.name)
# dmriprep_path = op.join(tmpdir.name, 'HCP', 'derivatives', 'dmriprep')
# seg_algo = "recobundles"
# bundle_names = ["F", "CST", "AF", "CC_ForcepsMajor"]
# myafq = api.AFQ(dmriprep_path=dmriprep_path,
# sub_prefix='sub',
# seg_algo=seg_algo,
# bundle_names=bundle_names,
# odf_model="DTI",
# b0_threshold=15)
# # Replace the streamlines with precomputed:
# path_to_trk = dpd.fetcher.fetch_target_tractogram_hcp()
# path_to_trk = dpd.fetcher.get_target_tractogram_hcp()
# sl_file = op.join(myafq.data_frame.results_dir[0], 'sub-100206_sess-01_dwiDTI_det_streamlines.trk')
# shutil.copy(path_to_trk, sl_file)
# myafq.data_frame["streamlines_file"] = sl_file
# print("here")
# tgram = nib.streamlines.load(myafq.bundles[0]).tractogram
# print("here")
# bundles = aus.tgram_to_bundles(tgram, myafq.bundle_dict)
# npt.assert_(len(bundles['CST_L']) > 0)
if not op.exists('dti_streamlines.trk'):
FA = nib.load(dti_params["FA"]).get_data()
wm_mask = np.zeros_like(FA)
wm_mask[FA > 0.2] = 1
step_size = 1
min_length_mm = 50
streamlines = dts.Streamlines(
aft.track(dti_params['params'],
directions="det",
seed_mask=wm_mask,
seeds=2,
stop_mask=FA,
stop_threshold=0.2,
step_size=step_size,
min_length=min_length_mm / step_size))
aus.write_trk('./dti_streamlines.trk', streamlines, affine=img.affine)
else:
tg = nib.streamlines.load('./dti_streamlines.trk').tractogram
streamlines = tg.apply_affine(np.linalg.inv(img.affine)).streamlines
print("We're looking at: %s streamlines" % len(streamlines))
templates = afd.read_templates()
templates['ARC_roi1_L'] = templates['SLF_roi1_L']
templates['ARC_roi1_R'] = templates['SLF_roi1_R']
templates['ARC_roi2_L'] = templates['SLFt_roi2_L']
templates['ARC_roi2_R'] = templates['SLFt_roi2_R']
bundle_names = ["ATR", "CGC", "CST", "HCC", "IFO", "ILF", "SLF", "ARC", "UNC"]
