Whole brain registration with SLR

whole_brain_slr.py

dname = 'C:\\Users\\elef\\Data\\rodrigo\\'

fname1 = dname + 'ADRC_TMP_wholeBrain.trk.gz'

fname2 = dname + 'IIT2mean.trk'

from nibabel.streamlines import load


trkfile1 = load(fname1)

streamlines1 = trkfile1.streamlines

trkfile2 = load(fname2)

streamlines2 = trkfile2.streamlines

from dipy.viz import window, actor

ren = window.Renderer()

ren.add(actor.line(streamlines1))
ren.add(actor.line(streamlines2))

window.show(ren)

from dipy.tracking.streamline import length, select_random_set_of_streamlines, set_number_of_points, transform_streamlines
from dipy.segment.clustering import QuickBundles
from dipy.align.streamlinear import StreamlineLinearRegistration
from time import time

def remove_clusters_by_size(clusters, min_size=0):
    by_size = lambda c: len(c) >= min_size
    return filter(by_size, clusters)

def whole_brain_slr(static, moving,
                    x0='affine',
                    rm_small_clusters=50,
                    maxiter=100,
                    select_random=None,
                    verbose=False,
                    greater_than=50,
                    less_than=250,
                    qb_thr=15,
                    nb_pts=20,
                    progressive=True):

    if verbose:
        print('Static streamlines size {}'.format(len(static)))
        print('Moving streamlines size {}'.format(len(moving)))

    def check_range(streamline, gt=greater_than, lt=less_than):

        if (length(streamline) > gt) & (length(streamline) < lt):
            return True
        else:
            return False

    streamlines1 = [s for s in static if check_range(s)]
    streamlines2 = [s for s in moving if check_range(s)]

    if verbose:

        print('Static streamlines after length reduction {}'
              .format(len(streamlines1)))
        print('Moving streamlines after length reduction {}'
              .format(len(streamlines2)))

    if select_random is not None:
        rstreamlines1 = select_random_set_of_streamlines(streamlines1,
                                                         select_random)
    else:
        rstreamlines1 = streamlines1

    rstreamlines1 = set_number_of_points(rstreamlines1, nb_pts)
    qb1 = QuickBundles(threshold=qb_thr)
    rstreamlines1 = [s.astype('f4') for s in rstreamlines1]
    cluster_map1 = qb1.cluster(rstreamlines1)
    clusters1 = remove_clusters_by_size(cluster_map1, rm_small_clusters)
    qb_centroids1 = [cluster.centroid for cluster in clusters1]

    if select_random is not None:
        rstreamlines2 = select_random_set_of_streamlines(streamlines2,
                                                         select_random)
    else:
        rstreamlines2 = streamlines2

    rstreamlines2 = set_number_of_points(rstreamlines2, nb_pts)
    qb2 = QuickBundles(threshold=qb_thr)
    rstreamlines2 = [s.astype('f4') for s in rstreamlines2]
    cluster_map2 = qb2.cluster(rstreamlines2)
    clusters2 = remove_clusters_by_size(cluster_map2, rm_small_clusters)
    qb_centroids2 = [cluster.centroid for cluster in clusters2]

    t = time()

    slr = StreamlineLinearRegistration(x0=x0,
                                       options={'maxiter': maxiter})
    slm = slr.optimize(qb_centroids1, qb_centroids2)

    if verbose:
        print('QB static centroids size %d' % len(qb_centroids1,))
        print('QB moving centroids size %d' % len(qb_centroids2,))

    duration = time() - t
    if verbose:
        print('SLR finished in  %0.3f seconds.' % (duration,))
        print('SLR iterations: %d ' % (slm.iterations,))

    moved = slm.transform(moving)

    return moved, slm.matrix, qb_centroids1, qb_centroids2


moved, transform, centroids1, centroids2 = whole_brain_slr(streamlines1,
                                                           streamlines2, 
                                                           'affine')

ren.clear()
ren.add(actor.line(centroids1, colors=(1, 0, 0)))
ren.add(actor.line(transform_streamlines(centroids2, transform), colors=(0, 1, 0)))

window.show(ren)

ren.clear()
ren.add(actor.line(streamlines1, colors=(1, 0, 0)))
ren.add(actor.line(moved, colors=(0, 1, 0)))

window.show(ren)