Plot a subset of clusters from a nifti file

plot_subset.py

from nilearn import plotting as niplt
import nibabel as nib
import seaborn as sns
import numpy as np

def plot_subset(nifti, layers, colors = None, **kwargs):
    if not isinstance(nifti, nib.Nifti1Image):    
        nifti = nib.load(nifti)
    data = nifti.get_data()

    if colors is None:
        colors = sns.color_palette('Set1', len(layers) + 1)

    for value in np.unique(np.round(data)):
        if np.in1d(value, layers)[0] == False:
            data[np.round(data) == value] = 0
    niplt.plot_roi(nifti, cmap = ListedColormap(colors),**kwargs)