prepare_data.py

import argparse
import numpy as np
import nipy
import utils

parser = argparse.ArgumentParser(description='Prepare ATLAS anatomical ground truth')
parser.add_argument('--brains_list')
args = parser.parse_args()

labels = [i for i in range(1001,1004)] + [i for i in range(1005,1036)] + [i for i in range(2001,2004)] + [i for i in range(2005, 2036)] + [10, 49, 11, 50, 12, 51, 13, 52, 17, 53, 18, 54, 26, 58, 28, 60, 2, 41, 4, 5, 43, 44, 14, 15, 24, 16, 7, 46, 8, 47, 251, 252, 253, 254, 255]
print (len(labels), labels)
filename = 'aparc+aseg.nii.gz'
output_filename = 'atlas_full_104.nii.gz'
f = open(args.brains_list, 'r')
brains = f.read().splitlines()
for i, b in enumerate(brains):
    print (b)
    img = nipy.load_image(b + filename).get_data()
    segmentation = np.zeros(img.shape)
    k = 1
    for l in labels:
        mask = np.equal(l, img)
        segmentation[mask] = k
        k += 1
    if (len(np.unique(segmentation)) == (len(labels) + 1)):
        utils.save_npy_to_nifti(segmentation, b + output_filename, b + filename)
    else:
        assert False, 'Bad {}'.format(b)