Parcellate volume data into N random parcel using floodfill to preserve connectivity.

vol_rand_parc.py

import argparse

import numpy as np

import nibabel as nib

from scipy.spatial.distance import pdist, squareform


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

Uses N non-recursive simultaneous 6-dir 'floodfill' to match voxel to seed points.

Seed point are randomely generated by a heuristic trying to maximize their 
 Euclidean distance (ignoring volume 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)


    # Finds n_parcel initial voxel somewhat distant

    # Voxel to parcellate are non-zeros
    pos = np.where(data)

    n_pos = len(pos[0])

    possible_seeding = []
    seeding_value = []

    # number of try for init heuristic
    N_seeding = 10000

    for idx in range(N_seeding):
        # seeding
        seed_pos = []
        while len(seed_pos) < n_parcel:
            candidate = np.random.randint(0, n_pos)
            if candidate not in seed_pos:
                seed_pos.append(candidate)


        seed_coord = np.concatenate((pos[0][seed_pos][:, None], pos[1][seed_pos][:, None], pos[2][seed_pos][:, None]), axis = 1)

        # judge seeding quality with distance between the 2 closest
        seed_dist = squareform(pdist(seed_coord, metric = 'euclidean', p=2)) + 10000*np.eye(n_parcel)
        seeding_value.append(seed_dist.min())
        possible_seeding.append(seed_pos)

    best_idx = np.argmax(np.array(seeding_value))

    seeding_pos = possible_seeding[best_idx]

    if verbose:
        print('Seeding Done.')



    # non-recursive flood fill

    # first version spatial ordering matter for filling priority
    parcelation = np.zeros_like(data).astype(np.float32)
    # init
    parcel_idx = 0
    for seed in seeding_pos:
        parcel_idx += 1
        parcelation[pos[0][seed], pos[1][seed], pos[2][seed]] = parcel_idx

    vox_pos_not_prop = range(int(n_pos))

    it_count = 0
    keep_going = True
    while keep_going:
        unfilled_spot_prev = n_pos - (parcelation!=0).sum()

        it_count += 1
        if verbose:
            if not it_count%1:
                print('iter {}, {} left'.format(it_count, unfilled_spot_prev))

        vox_idx_that_prop = []
        # one pass of filling
        for not_prop_idx in np.random.permutation(range(len(vox_pos_not_prop))):
            vox_idx = vox_pos_not_prop[not_prop_idx]

            if parcelation[pos[0][vox_idx], pos[1][vox_idx], pos[2][vox_idx]] != 0:
                parcel_idx = parcelation[pos[0][vox_idx], pos[1][vox_idx], pos[2][vox_idx]]

                # check for [+1, 0, 0]
                if (pos[0][vox_idx]+1 < data.shape[0]):
                    if data[pos[0][vox_idx]+1, pos[1][vox_idx], pos[2][vox_idx]] and (parcelation[pos[0][vox_idx]+1, pos[1][vox_idx], pos[2][vox_idx]] == 0):
                        parcelation[pos[0][vox_idx]+1, pos[1][vox_idx], pos[2][vox_idx]] = parcel_idx

                # check for [-1, 0, 0]
                if (pos[0][vox_idx] > 0):
                    if data[pos[0][vox_idx]-1, pos[1][vox_idx], pos[2][vox_idx]] and (parcelation[pos[0][vox_idx]-1, pos[1][vox_idx], pos[2][vox_idx]] == 0):
                        parcelation[pos[0][vox_idx]-1, pos[1][vox_idx], pos[2][vox_idx]] = parcel_idx

                # check for [0, +1, 0]
                if (pos[1][vox_idx]+1 < data.shape[1]):
                    if data[pos[0][vox_idx], pos[1][vox_idx]+1, pos[2][vox_idx]] and (parcelation[pos[0][vox_idx], pos[1][vox_idx]+1, pos[2][vox_idx]] == 0):
                        parcelation[pos[0][vox_idx], pos[1][vox_idx]+1, pos[2][vox_idx]] = parcel_idx

                # check for [0, -1, 0]
                if (pos[1][vox_idx] > 0):
                    if data[pos[0][vox_idx], pos[1][vox_idx]-1, pos[2][vox_idx]] and (parcelation[pos[0][vox_idx], pos[1][vox_idx]-1, pos[2][vox_idx]] == 0):
                        parcelation[pos[0][vox_idx], pos[1][vox_idx]-1, pos[2][vox_idx]] = parcel_idx

                # check for [0, 0, +1]
                if (pos[2][vox_idx]+1 < data.shape[2]):
                    if data[pos[0][vox_idx], pos[1][vox_idx], pos[2][vox_idx]+1] and (parcelation[pos[0][vox_idx], pos[1][vox_idx], pos[2][vox_idx]+1] == 0):
                        parcelation[pos[0][vox_idx], pos[1][vox_idx], pos[2][vox_idx]+1] = parcel_idx

                # check for [0, 0, -1]
                if (pos[2][vox_idx] > 0):
                    if data[pos[0][vox_idx], pos[1][vox_idx], pos[2][vox_idx]-1] and (parcelation[pos[0][vox_idx], pos[1][vox_idx], pos[2][vox_idx]-1] == 0):
                        parcelation[pos[0][vox_idx], pos[1][vox_idx], pos[2][vox_idx]-1] = parcel_idx


                vox_idx_that_prop.append(vox_idx)

        for vox_to_remove in vox_idx_that_prop:
            vox_pos_not_prop.remove(vox_to_remove)



        filled_at_iter = unfilled_spot_prev - (n_pos - (parcelation!=0).sum())
        if filled_at_iter == 0:
            keep_going = False

        if verbose:
            print('Non-parcelated voxel = {}'.format(n_pos - (parcelation!=0).sum()))

    if verbose:
        for i in range(1, n_parcel+1):
            print('Parcel {} = {}'.format(i, (parcelation==i).sum()))

    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()