Parcellate volume data into N random parcel using kmeans.

vol_rand_parc_fast.py

import argparse

import numpy as np

import nibabel as nib

import sklearn.cluster as clu


DESCRIPTION = """
Parcellate volume data into N random parcel.

Uses sklearn kmean clustering.

Does not respect mask topology.
"""

EPILOG = """
"""


def buildArgsParser():
    p = argparse.ArgumentParser(
    formatter_class = argparse.RawDescriptionHelpFormatter,
    description = DESCRIPTION,
    epilog = EPILOG)
    p._optionals.title = "Options and Parameters"

    p.add_argument(
        'volume', action='store', metavar='volume', type=str,
        help='Filename for mask to parcellate.')
    
    p.add_argument(
        'outfile', action='store', metavar='outfile', type=str,
        help='Output filename (don\'t include extension).')

    p.add_argument(
        'nparcel', action='store', metavar='nparcel', type=int,
        help='Number of desired parcel.')

    p.add_argument(
        '-v', action='store_true', dest='v', default=False,
        help='Enable verbose')

    return p


def main():

    # Parser
    parser = buildArgsParser()
    args = parser.parse_args()

    volume = args.volume
    outfile = args.outfile
    n_parcel = args.nparcel
    verbose = args.v


    vol = nib.load(volume)
    data = vol.get_data()

    if verbose:
        print('Volume loaded.')

    # # Controled RNG seeding
    # i_parcel = 0
    # np.random.seed(i_parcel)

    pos = np.where(data)
    pts = np.zeros((np.array(pos[0]).shape[0], 3))
    pts[:,0] = np.array(pos[0])
    pts[:,1] = np.array(pos[1])
    pts[:,2] = np.array(pos[2])


    n_init = 3
    clusterer = clu.MiniBatchKMeans(n_clusters = n_parcel, n_init = n_init, init = "k-means++", compute_labels = True)
    fitted_cluster = clusterer.fit(pts)

    if verbose:
        print('Clustering done.')


    if verbose:
        ll = np.array([(fitted_cluster.labels_==i).sum() for i in range(n_parcel)])
        # print(ll.sum(), pts.shape[0])
        print('parcel_id   #voxel')
        for i in range(n_parcel):
            print(i, ll[i])


    parcelation = np.zeros_like(data).astype(np.float32)
    for ii in range(pts.shape[0]):
        parcelation[int(pts[ii,0]), int(pts[ii,1]), int(pts[ii,2])] = fitted_cluster.labels_[ii]



    output = nib.Nifti1Image(parcelation.astype(np.float32), vol.get_affine())
    nib.save(output, outfile + '.nii.gz')

    if verbose:
        print('Saved results as \n' + outfile + '.nii.gz')



if __name__ == "__main__":
    main()