def main():
parser = _build_arg_parser()
args = parser.parse_args()
assert_inputs_exist(parser, [args.moving_tractogram,
args.static_tractogram])
if args.only_rigid:
matrix_filename = os.path.splitext(args.out_name)[0] + '_rigid.txt'
else:
matrix_filename = os.path.splitext(args.out_name)[0] + '_affine.txt'
assert_outputs_exist(parser, args, matrix_filename, args.out_name)
sft_moving = load_tractogram_with_reference(parser, args,
args.moving_tractogram,
arg_name='moving_tractogram')
sft_static = load_tractogram_with_reference(parser, args,
args.static_tractogram,
arg_name='static_tractogram')
if args.only_rigid:
transformation_type = 'rigid'
else:
transformation_type = 'affine'
ret = whole_brain_slr(sft_moving.streamlines,
sft_static.streamlines,
x0=transformation_type,
maxiter=150,
verbose=args.verbose)
_, transfo, _, _ = ret
np.savetxt(matrix_filename, transfo)
def register_tractogram(moving_filename, static_filename, only_rigid,
amount_to_load, matrix_filename, verbose):
amount_to_load = max(250000, amount_to_load)
moving_tractogram = nib.streamlines.load(moving_filename, lazy_load=True)
moving_streamlines = next(
ichunk(moving_tractogram.streamlines, amount_to_load))
static_tractogram = nib.streamlines.load(static_filename, lazy_load=True)
static_streamlines = next(
ichunk(static_tractogram.streamlines, amount_to_load))
if only_rigid:
transformation_type = 'rigid'
else:
transformation_type = 'affine'
ret = whole_brain_slr(static_streamlines,
moving_streamlines,
x0=transformation_type,
maxiter=150,
select_random=1000000,
verbose=verbose)
_, transfo, _, _ = ret
np.savetxt(matrix_filename, transfo)
def rough_reg(sub_fixed, temp_moving, save_base_abspath=None):
# template moves to the subject space
moved, transform, qb_cents1, qb_cents2 = whole_brain_slr(sub_fixed,
temp_moving,
verbose=True,
progressive=True)
if save_base_abspath is not None:
pklf = open(save_base_abspath + '.pkl', "wb")
pickle.dump(transform, pklf)
return moved, transform, qb_cents1, qb_cents2
def test_whole_brain_slr():
streams, hdr = nib.trackvis.read(get_fnames('fornix'))
fornix = [s[0] for s in streams]
f = Streamlines(fornix)
f1 = f.copy()
f2 = f.copy()
# check translation
f2._data += np.array([50, 0, 0])
moved, transform, qb_centroids1, qb_centroids2 = whole_brain_slr(
f1, f2, x0='affine', verbose=True, rm_small_clusters=2,
greater_than=0, less_than=np.inf,
qbx_thr=[5, 2, 1], progressive=False)
# we can check the quality of registration by comparing the matrices
# MAM streamline distances before and after SLR
D12 = bundles_distances_mam(f1, f2)
D1M = bundles_distances_mam(f1, moved)
d12_minsum = np.sum(np.min(D12, axis=0))
d1m_minsum = np.sum(np.min(D1M, axis=0))
print("distances= ", d12_minsum, " ", d1m_minsum)
assert_equal(d1m_minsum < d12_minsum, True)
assert_array_almost_equal(transform[:3, 3], [-50, -0, -0], 2)
# check rotation
mat = compose_matrix44([0, 0, 0, 15, 0, 0])
f3 = f.copy()
f3 = transform_streamlines(f3, mat)
moved, transform, qb_centroids1, qb_centroids2 = slr_with_qbx(
f1, f3, verbose=False, rm_small_clusters=1, greater_than=20,
less_than=np.inf, qbx_thr=[2],
progressive=True)
# we can also check the quality by looking at the decomposed transform
assert_array_almost_equal(decompose_matrix44(transform)[3], -15, 2)
moved, transform, qb_centroids1, qb_centroids2 = slr_with_qbx(
f1, f3, verbose=False, rm_small_clusters=1, select_random=400,
greater_than=20, less_than=np.inf, qbx_thr=[2],
progressive=True)
# we can also check the quality by looking at the decomposed transform
assert_array_almost_equal(decompose_matrix44(transform)[3], -15, 2)
def recobundles(streamlines, bundle_dict):
"""
Segment streamlines using the RecoBundles algorithm [Garyfallidis2017]
Parameters
----------
streamlines : list or Streamlines object.
A whole-brain tractogram to be segmented.
bundle_dict: dictionary
Of the form:
{'whole_brain': Streamlines,
'CST_L': {'sl': Streamlines, 'centroid': array},
'CST_R': {'sl': Streamlines, 'centroid': array},
...}
Returns
-------
fiber_groups : dict
Keys are names of the bundles, values are Streamline objects.
The streamlines in each object have all been oriented to have the
same orientation (using `dts.orient_by_streamline`).
"""
fiber_groups = {}
# We start with whole-brain SLR:
atlas = bundle_dict['whole_brain']
moved, transform, qb_centroids1, qb_centroids2 = whole_brain_slr(
atlas, streamlines, x0='affine', verbose=False, progressive=True)
# We generate our instance of RB with the moved streamlines:
rb = RecoBundles(moved, verbose=False)
# Next we'll iterate over bundles, registering each one:
bundle_list = list(bundle_dict.keys())
bundle_list.remove('whole_brain')
for bundle in bundle_list:
model_sl = bundle_dict[bundle]['sl']
_, rec_labels = rb.recognize(model_bundle=model_sl,
model_clust_thr=5.,
reduction_thr=10,
reduction_distance='mam',
slr=True,
slr_metric='asymmetric',
pruning_distance='mam')
# Use the streamlines in the original space:
recognized_sl = streamlines[rec_labels]
standard_sl = bundle_dict[bundle]['centroid']
oriented_sl = dts.orient_by_streamline(recognized_sl, standard_sl)
fiber_groups[bundle] = oriented_sl
return fiber_groups
def bundle_extract(atlas_track_path, atlas_bundle_path, target_track_path):
time0 = time.time()
atlas_file = atlas_track_path
target_file = target_track_path
print('loading data begin! time:', time.time() - time0)
sft_atlas = load_trk(atlas_file, "same", bbox_valid_check=False)
atlas = sft_atlas.streamlines
atlas_header = create_tractogram_header(atlas_file,
*sft_atlas.space_attributes)
sft_target = load_trk(target_file, "same", bbox_valid_check=False)
target = sft_target.streamlines
target_header = create_tractogram_header(target_file,
*sft_target.space_attributes)
moved, transform, qb_centroids1, qb_centroids2 = whole_brain_slr(
atlas,
target,
x0='affine',
verbose=True,
progressive=True,
rng=np.random.RandomState(1984))
bundle_track = StatefulTractogram(moved, target_header, Space.RASMM)
save_trk(bundle_track, 'moved.trk', bbox_valid_check=False)
np.save("slr_transform.npy", transform)
model_bundle_file = atlas_bundle_path
model_bundle = load_trk(model_bundle_file, "same", bbox_valid_check=False)
model_bundle = model_bundle.streamlines
print('comparing begin! time:', time.time() - time0)
rb = RecoBundles(moved, verbose=True, rng=np.random.RandomState(2001))
recognized_bundle, bundle_labels = rb.recognize(model_bundle=model_bundle,
model_clust_thr=0,
reduction_thr=20,
reduction_distance='mam',
slr=True,
slr_metric='asymmetric',
pruning_distance='mam')
bundle_track = StatefulTractogram(target[bundle_labels], target_header,
Space.RASMM)
return bundle_track
def test_whole_brain_slr():
streams, hdr = nib.trackvis.read(get_data('fornix'))
fornix = [s[0] for s in streams]
f = Streamlines(fornix)
f1 = f.copy()
f2 = f.copy()
# check translation
f2._data += np.array([50, 0, 0])
moved, transform, qb_centroids1, qb_centroids2 = whole_brain_slr(
f1, f2, verbose=True, rm_small_clusters=2, greater_than=0,
less_than=np.inf, qb_thr=5, progressive=False)
# we can check the quality of registration by comparing the matrices
# MAM streamline distances before and after SLR
D12 = bundles_distances_mam(f1, f2)
D1M = bundles_distances_mam(f1, moved)
d12_minsum = np.sum(np.min(D12, axis=0))
d1m_minsum = np.sum(np.min(D1M, axis=0))
assert_equal(d1m_minsum < d12_minsum, True)
assert_array_almost_equal(transform[:3, 3], [-50, -0, -0], 3)
# check rotation
mat = compose_matrix44([0, 0, 0, 15, 0, 0])
f3 = f.copy()
f3 = transform_streamlines(f3, mat)
moved, transform, qb_centroids1, qb_centroids2 = slr_with_qb(
f1, f3, verbose=False, rm_small_clusters=1, greater_than=20,
less_than=np.inf, qb_thr=2, progressive=True)
# we can also check the quality by looking at the decomposed transform
assert_array_almost_equal(decompose_matrix44(transform)[3], -15, 2)
moved, transform, qb_centroids1, qb_centroids2 = slr_with_qb(
f1, f3, verbose=False, rm_small_clusters=1, select_random=400,
greater_than=20,
less_than=np.inf, qb_thr=2, progressive=True)
# we can also check the quality by looking at the decomposed transform
assert_array_almost_equal(decompose_matrix44(transform)[3], -15, 2)
def slr_registration(moving_data,
static_data,
moving_affine=None,
static_affine=None,
moving_shape=None,
static_shape=None,
**kwargs):
"""Register a source image (moving) to a target image (static).
Parameters
----------
moving : ndarray
The source tractography data to be registered
moving_affine : ndarray
The affine associated with the moving (source) data.
moving_shape : ndarray
The shape of the space associated with the static (target) data.
static : ndarray
The target tractography data for registration
static_affine : ndarray
The affine associated with the static (target) data.
static_shape : ndarray
The shape of the space associated with the static (target) data.
**kwargs:
kwargs are passed into whole_brain_slr
Returns
-------
AffineMap
"""
from AFQ.definitions.mapping import ConformedAffineMapping
_, transform, _, _ = whole_brain_slr(static_data,
moving_data,
x0='affine',
verbose=False,
**kwargs)
return ConformedAffineMapping(
transform,
codomain_grid_shape=reduce_shape(static_shape),
codomain_grid2world=static_affine,
domain_grid_shape=reduce_shape(moving_shape),
domain_grid2world=moving_affine)
def move_streamlines(self, tg=None, reg_algo='slr'):
"""Streamline-based registration of a whole-brain tractogram to
the MNI whole-brain atlas.
registration_algo : str
"slr" or "syn"
"""
tg = self._read_tg(tg=tg)
if reg_algo is None:
if self.mapping is None:
reg_algo = 'slr'
else:
reg_algo = 'syn'
if reg_algo == "slr":
self.logger.info("Registering tractogram with SLR")
atlas = self.bundle_dict['whole_brain']
self.moved_sl, _, _, _ = whole_brain_slr(
atlas,
self.tg.streamlines,
x0='affine',
verbose=False,
progressive=self.progressive,
greater_than=self.greater_than,
rm_small_clusters=self.rm_small_clusters,
rng=self.rng)
elif reg_algo == "syn":
self.logger.info("Registering tractogram based on syn")
self.tg.to_rasmm()
delta = dts.values_from_volume(self.mapping.forward,
self.tg.streamlines, np.eye(4))
self.moved_sl = dts.Streamlines(
[d + s for d, s in zip(delta, self.tg.streamlines)])
self.tg.to_vox()
if self.save_intermediates is not None:
moved_sft = StatefulTractogram(self.moved_sl, self.reg_template,
Space.RASMM)
save_tractogram(moved_sft,
op.join(self.save_intermediates, 'sls_in_mni.trk'),
bbox_valid_check=False)
def transform_bundles(wb_tracts_name):
atlas_file, all_bundles_files = get_bundle_atlas_hcp842()
sft_atlas = load_trk(atlas_file, "same", bbox_valid_check=False)
atlas = sft_atlas.streamlines
sft_target = load_trk(wb_tracts_name, "same", bbox_valid_check=False)
target = sft_target.streamlines
#show_atlas_target_graph(atlas, target,out_path=folder_name+r'\try_atlas_target',interactive=True)
atlas = set_number_of_points(atlas, 20)
target = set_number_of_points(target, 20)
moved, transform, qb_centroids1, qb_centroids2 = whole_brain_slr(
atlas,
target,
x0='affine',
verbose=True,
progressive=True,
rng=np.random.RandomState(1984))
#np.save("slf_L_transform.npy", transform)
#show_atlas_target_graph(atlas, moved,out_path=r'',interactive=True)
return moved, target
def run_recobundles(input_folder, atlas_file, ground_truth_folder,
output_folder):
print('Running Recobundles in ' + input_folder)
# make a folder to save output
try:
Path(output_folder).mkdir(parents=True, exist_ok=True)
except OSError:
print('Could not create output dir. Aborting...')
return
# Uncomment for first exemplary use
# target_file, target_folder = fetch_target_tractogram_hcp()
# atlas_file, atlas_folder = fetch_bundle_atlas_hcp842()
# target_file = get_target_tractogram_hcp()
target_file = input_folder + 'whole_brain.trk'
# use this line to select tracts if necessary
sel_tracts = tracts
# sel_bundle_paths = [data_path + 'tractseg/599469/tracts/AF_left.trk']
# print(sel_bundle_paths)
sft_atlas = load_trk(atlas_file, 'same', bbox_valid_check=False)
atlas = sft_atlas.streamlines
sft_target = load_trk(target_file, 'same', bbox_valid_check=False)
target = sft_target.streamlines
target_header = create_tractogram_header(target_file,
*sft_atlas.space_attributes)
target, transform, qb_centroids1, qb_centroids2 = whole_brain_slr(
target, atlas, x0='affine', progressive=True)
print(transform)
sft_rec = StatefulTractogram(
nib.streamlines.array_sequence.concatenate([target, atlas], 0),
target_header, Space.RASMM)
save_trk(sft_rec, output_folder + 'test.trk', bbox_valid_check=False)
def register_tractogram(moving_tractogram, static_tractogram, only_rigid,
amount_to_load, matrix_filename, verbose):
amount_to_load = max(250000, amount_to_load)
moving_streamlines = next(
ichunk(moving_tractogram.streamlines, amount_to_load))
static_streamlines = next(
ichunk(static_tractogram.streamlines, amount_to_load))
if only_rigid:
transformation_type = 'rigid'
else:
transformation_type = 'affine'
ret = whole_brain_slr(ArraySequence(static_streamlines),
ArraySequence(moving_streamlines),
x0=transformation_type,
maxiter=150,
verbose=verbose)
_, transfo, _, _ = ret
np.savetxt(matrix_filename, transfo)
window.record(ren, out_path='tractograms_initial.png', size=(600, 600))
if interactive:
window.show(ren)
"""
.. figure:: tractograms_initial.png
:align: center
Atlas and target before registration.
"""
"""
We will register target tractogram to model atlas' space using streamlinear
registeration (SLR) [Garyfallidis15]_
"""
moved, transform, qb_centroids1, qb_centroids2 = whole_brain_slr(
atlas, target, x0='affine', verbose=True, progressive=True)
"""
We save the transform generated in this registration, so that we can use
it in the bundle profiles example
"""
np.save("slr_transform.npy", transform)
"""
let's visualize atlas tractogram and target tractogram after registration
"""
interactive = False
ren = window.Renderer()
ren.SetBackground(1, 1, 1)
ren.add(actor.line(atlas, colors=(1, 0, 1)))
def test_whole_brain_slr():
fname = get_fnames('fornix')
fornix = load_tractogram(fname, 'same', bbox_valid_check=False).streamlines
f = Streamlines(fornix)
f1 = f.copy()
f2 = f.copy()
# check translation
f2._data += np.array([50, 0, 0])
moved, transform, qb_centroids1, qb_centroids2 = whole_brain_slr(
f1,
f2,
x0='affine',
verbose=True,
rm_small_clusters=2,
greater_than=0,
less_than=np.inf,
qbx_thr=[5, 2, 1],
progressive=False)
# we can check the quality of registration by comparing the matrices
# MAM streamline distances before and after SLR
D12 = bundles_distances_mam(f1, f2)
D1M = bundles_distances_mam(f1, moved)
d12_minsum = np.sum(np.min(D12, axis=0))
d1m_minsum = np.sum(np.min(D1M, axis=0))
print("distances= ", d12_minsum, " ", d1m_minsum)
assert_equal(d1m_minsum < d12_minsum, True)
assert_array_almost_equal(transform[:3, 3], [-50, -0, -0], 2)
# check rotation
mat = compose_matrix44([0, 0, 0, 15, 0, 0])
f3 = f.copy()
f3 = transform_streamlines(f3, mat)
moved, transform, qb_centroids1, qb_centroids2 = slr_with_qbx(
f1,
f3,
verbose=False,
rm_small_clusters=1,
greater_than=20,
less_than=np.inf,
qbx_thr=[2],
progressive=True)
# we can also check the quality by looking at the decomposed transform
assert_array_almost_equal(decompose_matrix44(transform)[3], -15, 2)
moved, transform, qb_centroids1, qb_centroids2 = slr_with_qbx(
f1,
f3,
verbose=False,
rm_small_clusters=1,
select_random=400,
greater_than=20,
less_than=np.inf,
qbx_thr=[2],
progressive=True)
# we can also check the quality by looking at the decomposed transform
assert_array_almost_equal(decompose_matrix44(transform)[3], -15, 2)
def rough_reg(sub_fixed, temp_moving):
# template moves to the subject space
# qb_thr=5 errored
moved, transform, qb_centroids1, qb_centroids2 = whole_brain_slr(
sub_fixed, temp_moving, verbose=True, progressive=True)
return moved, transform, qb_centroids1, qb_centroids2
def segment_reco(self, tg=None):
"""
Segment streamlines using the RecoBundles algorithm [Garyfallidis2017]
Parameters
----------
tg : StatefulTractogram class instance
A whole-brain tractogram to be segmented.
Returns
-------
fiber_groups : dict
Keys are names of the bundles, values are Streamline objects.
The streamlines in each object have all been oriented to have the
same orientation (using `dts.orient_by_streamline`).
"""
if tg is None:
tg = self.tg
else:
self.tg = tg
fiber_groups = {}
self.logger.info("Registering Whole-brain with SLR")
# We start with whole-brain SLR:
atlas = self.bundle_dict['whole_brain']
moved, transform, qb_centroids1, qb_centroids2 = whole_brain_slr(
atlas,
self.tg.streamlines,
x0='affine',
verbose=False,
progressive=self.progressive,
greater_than=self.greater_than,
rm_small_clusters=self.rm_small_clusters,
rng=self.rng)
# We generate our instance of RB with the moved streamlines:
self.logger.info("Extracting Bundles")
rb = RecoBundles(moved, verbose=False, rng=self.rng)
# Next we'll iterate over bundles, registering each one:
bundle_list = list(self.bundle_dict.keys())
bundle_list.remove('whole_brain')
self.logger.info("Assigning Streamlines to Bundles")
for bundle in bundle_list:
model_sl = self.bundle_dict[bundle]['sl']
_, rec_labels = rb.recognize(model_bundle=model_sl,
model_clust_thr=self.model_clust_thr,
reduction_thr=self.reduction_thr,
reduction_distance='mdf',
slr=True,
slr_metric='asymmetric',
pruning_distance='mdf')
# Use the streamlines in the original space:
recognized_sl = tg.streamlines[rec_labels]
if self.refine:
_, rec_labels = rb.refine(model_sl,
recognized_sl,
self.model_clust_thr,
reduction_thr=self.reduction_thr,
pruning_thr=self.pruning_thr)
recognized_sl = tg.streamlines[rec_labels]
standard_sl = self.bundle_dict[bundle]['centroid']
oriented_sl = dts.orient_by_streamline(recognized_sl, standard_sl)
if self.return_idx:
fiber_groups[bundle] = {}
fiber_groups[bundle]['idx'] = rec_labels
fiber_groups[bundle]['sl'] = StatefulTractogram(
oriented_sl, self.img, Space.RASMM)
else:
fiber_groups[bundle] = StatefulTractogram(
oriented_sl, self.img, Space.RASMM)
self.fiber_groups = fiber_groups
return fiber_groups
"""
.. figure:: tractograms_initial.png
:align: center
Atlas and target before registration.
"""
"""
We will register target tractogram to model atlas' space using streamlinear
registeration (SLR) [Garyfallidis15]_
"""
moved, transform, qb_centroids1, qb_centroids2 = whole_brain_slr(
atlas,
target,
x0='affine',
verbose=True,
progressive=True,
rng=np.random.RandomState(1984))
"""
We save the transform generated in this registration, so that we can use
it in the bundle profiles example
"""
np.save("slr_transform.npy", transform)
"""
let's visualize atlas tractogram and target tractogram after registration
"""
interactive = False
