mask.py

# Download atlas with this link https://neurovault.org/images/23262/

import numpy as np
import nibabel
from nilearn import image

labels_data = nibabel.load('aparcaseg.nii.gz').get_fdata()
unique_labels = np.unique(labels_data)

# Append each image to concat in 4th dimension
imgs = []

# Just running on 10 parcels to make less intensive
for each_label in unique_labels[:10]:
    this_label_mask = (labels_data == each_label).astype(np.int)
    this_label_mask_img = image.new_img_like('aparcaseg.nii.gz',
                                             this_label_mask)
    imgs.append(this_label_mask_img)

print(len(imgs))

atlas = image.concat_imgs(imgs)