Skip to content

Instantly share code, notes, and snippets.

@MarcCote
Created March 14, 2014 16:06
Show Gist options
  • Save MarcCote/9550783 to your computer and use it in GitHub Desktop.
Save MarcCote/9550783 to your computer and use it in GitHub Desktop.
Example of saving a trk file with color using nibabel.
import numpy as np
from nibabel import trackvis as tv
from dipy.segment.quickbundles import QuickBundles
from dipy.data import get_data
fname = get_data('fornix')
streams, hdr = tv.read(fname)
streamlines = [i[0] for i in streams]
qb = QuickBundles(streamlines, dist_thr=10., pts=12)
centroids = qb.centroids
colormap = np.random.rand(len(centroids), 3)
colors = np.ones((len(streamlines), 3))
for i, centroid in enumerate(centroids):
inds = qb.label2tracksids(i)
colors[inds] = colormap[i]
def save_streamlines_with_color(streamlines, colors, out_file, hdr):
trk = []
for i, streamline in enumerate(streamlines):
# Color (RBG [0-255]) for each point of the streamline
scalars = np.array([colors[i]] * len(streamline)) * 255
properties = None
trk.append((streamline, scalars, properties))
tv.write(out_file, trk, hdr)
save_streamlines_with_color(streamlines, colors, 'fornix_clustered.trk', hdr)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment