Example of saving a trk file with color using nibabel.

save_trk_with_color.py

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)