"""

Attributes:

pick_vx(function): Function to pick voxels

pick_patch(function): Function to pick patches

pick_tg(function): Function to pick patches

files: List of pairs (mri_file, label_file)

atlases: List of pairs (mri array, label array)

n_files (int): Number of files

n_out_features (int): Number of output classes in the datasets

"""

# See Miccai rules

ignored_labels = range(1,4)+range(5,11)+range(12,23)+range(24,30)+[33,34]+[42,43]+[53,54]+range(63,69)+[70,74]+\

range(80,100)+[110,111]+[126,127]+[130,131]+[158,159]+[188,189]

true_labels = [4, 11, 23, 30, 31, 32, 35, 36, 37, 38, 39, 40, 41, 44, 45, 46, 47, 48, 49, 50, 51, 52, 55, 56, 57,

58, 59, 60, 61, 62, 69, 71, 72, 73, 75, 76, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 112,

113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 128, 129, 132, 133, 134, 135, 136,

137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156,

157, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178,

179, 180, 181, 182, 183, 184, 185, 186, 187, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200,

201, 202, 203, 204, 205, 206, 207]

def __init__(self):

self.pick_vx = None

self.pick_features = None

self.pick_tg = None

self.files = None

self.n_files = None

self.atlases = []

self.ls_region_centroids = []

self.n_out_features = None

def init_from_config(self, config):

self.pick_vx = create_pick_voxel(config)

self.pick_features = create_pick_features(config)

self.pick_tg = create_pick_target(config)

self.files = list_miccai_files(**config.general["source_kwargs"])

self.__init_common()

def init_from(self, files, pick_vx, pick_patch, pick_tg):

self.pick_vx = pick_vx

self.pick_features = pick_patch

self.pick_tg = pick_tg

self.files = files

self.__init_common()

def __init_common(self):

self.n_files = len(self.files)

self.ls_region_centroids = [None]*self.n_files

print " preprocess the atlases ..."

for i, file_names in enumerate(self.files):

mri_file, lab_file = file_names

print " " + mri_file

# nib.nifti1.FLOAT32_EPS_3 = -1e-6

mri = nib.load(mri_file).get_data().squeeze()

mri = mri.astype(np.float32, copy=False)

lab = nib.load(lab_file).get_data().squeeze()

lab = lab.astype(np.int16, copy=False)

mri, lab = crop_brain_and_pad(mri, lab, self.pick_features.required_pad)

self.scale_atlas(mri, lab)

# plt.imshow(mri[100,:,:], cmap = cm.Greys_r)

# plt.savefig("salut1.png")

affine = nib.load(mri_file).get_affine()

self.atlases.append((mri, lab, affine))

# Compute the centroids

if self.pick_features.has_instance_of(PickCentroidDistances):

region_centroids = RegionCentroids(134)

temp = lab.nonzero()

vxs = np.asarray(temp).T

region_centroids.update_barycentres(vxs, lab[temp])

self.ls_region_centroids[i] = region_centroids

# Number of classes

self.n_out_features = 135

def scale_atlas(self, mri, label):

only_brain = mri[label.nonzero()]

scalar_mean = np.mean(only_brain)

scalar_std = np.std(only_brain)

mri -= scalar_mean

mri /= scalar_std

def generate_single_atlas(self, atlas_id, n_points, region_centroid, batch_size, verbose=False):

print(" file {} \n".format(self.files[atlas_id]))

mri, lab, _ = self.atlases[atlas_id]

vx_batches_generator = self.pick_vx.pick(n_points, lab, verbose=verbose, batch_size=batch_size)

for vx_batch in vx_batches_generator:

patch = self.pick_features.pick(vx_batch, mri, lab, region_centroid)[0]

# patch_lab = self.pick_features.pick(vx_batch, lab, lab, self.ls_region_centroids[atlas_id])[0]

# image1_1 = PIL.Image.fromarray(tile_raster_images(X=patch[0:10],

# img_shape=(29, 29), tile_shape=(1, 10),

# tile_spacing=(1, 1)))

# image1_1.save("patches_2D_mri_1.png")

# patch_lab = self.pick_features.pick(vx_batch, lab, lab, self.ls_region_centroids[atlas_id])[0]

# image1_1 = PIL.Image.fromarray(tile_raster_images(X=patch[10:20],

# img_shape=(29, 29), tile_shape=(1, 10),

# tile_spacing=(1, 1)))

# image1_1.save("patches_2D_mri_2.png")

# temp_arr = tile_raster_images(X=patch_lab[0:10],

# img_shape=(29, 29), tile_shape=(1, 10),

# tile_spacing=(1, 1))

# image2_1 = PIL.Image.fromarray(np.uint8(cm.spectral(temp_arr)*255))

# image2_1.save("patches_2D_seg_1.png")

# temp_arr = tile_raster_images(X=patch_lab[10:20],

# img_shape=(29, 29), tile_shape=(1, 10),

# tile_spacing=(1, 1))

# image2_1 = PIL.Image.fromarray(np.uint8(cm.spectral(temp_arr)*255))

# image2_1.save("patches_2D_seg_2.png")

tg = self.pick_tg.pick(vx_batch, self.n_out_features, mri, lab)

yield vx_batch, patch, tg

def generate_parallel(self, batch_size):

print "Generate data ..."

### Initialization of the containers

# Compute the number of voxels to extract from each atlas

voxels_per_atlas = distrib_balls_in_bins(batch_size, self.n_files)

### Fill in the containers

# Function that will be run in parallel

def generate_from_one_brain(atlas_id):

n_points = voxels_per_atlas[atlas_id]

# Large batch_size so it can not be reached. We want to store everything, so we don't split

vx, patch, tg = \

next(self.generate_single_atlas(atlas_id, n_points,

self.ls_region_centroids[atlas_id], batch_size=1000000))

return vx, patch, tg, atlas_id

# Generate the patches in parallel

if self.n_files == 1: # This special case is necessary to avoid a bug on the server

res_all = map(generate_from_one_brain, range(self.n_files))

else:

res_all = parmap(generate_from_one_brain, range(self.n_files))

# Initialize the containers

vx = np.zeros((batch_size, 3), dtype=int)

patch = np.zeros((batch_size, self.pick_features.n_features), dtype=theano.config.floatX)

tg = np.zeros((batch_size, self.n_out_features), dtype=theano.config.floatX)

file_id = np.zeros((batch_size, 1), dtype=int)

# Aggregate the data

idx1 = 0

for res in res_all:

idx2 = idx1 + res[0].shape[0]

vx[idx1:idx2], patch[idx1:idx2], tg[idx1:idx2], file_id[idx1:idx2] = res

idx1 = idx2

"""

List the the pairs (mri_file_name, label_file_name) of the miccai data.

"""

mode = kwargs["mode"]

path = kwargs["path"]

label_path = path + "label/"

mri_files = glob.glob(path + "mri/*.nii")

if mode == "folder":

idx_files = xrange(len(mri_files))

elif mode == "idx_files":

idx_files = kwargs["idx_files"]

else:

raise Exception("Error to list the MICCAI files, the mode does not exist.")

return [(mri_files[i], label_path + os.path.splitext(os.path.basename(mri_files[i]))[0] + "_glm.nii")

"""

Find the boundaries of the non-zero region of the image

"""

def check_limit_one_side(fun, iterations):

for i in iterations:

if np.any(fun(i)):

return i

return iterations[-1]

lim = np.zeros((3, 2), dtype=int)

dims = img.shape

f0 = lambda i: img[i, :, :]

f1 = lambda i: img[:, i, :]

f2 = lambda i: img[:, :, i]

f = (f0, f1, f2)

for j in xrange(3):

lim[j, 0] = check_limit_one_side(f[j], xrange(dims[j]))

lim[j, 1] = check_limit_one_side(f[j], reversed(xrange(dims[j])))

"""

Extract the brain from an mri image

"""

lim = check_img_limits(lab)

lim[:, 0] -= pad

lim[:, 1] += pad

dim_orig = np.array(mri.shape)

pad_inf = np.zeros((3,), dtype=int)

too_low = lim[:, 0] < 0

pad_inf[too_low] = -lim[too_low, 0]

lim[too_low, 0] = 0

pad_sup = np.zeros((3,), dtype=int)

too_high = lim[:, 1] > dim_orig

pad_sup[too_high] = lim[too_high, 1] - dim_orig[too_high]

lim[too_high, 1] = dim_orig[too_high]

lim0 = slice(lim[0, 0], lim[0, 1])

lim1 = slice(lim[1, 0], lim[1, 1])

lim2 = slice(lim[2, 0], lim[2, 1])

mri = mri[lim0, lim1, lim2]

lab = lab[lim0, lim1, lim2]

mri = np.lib.pad(mri, zip(pad_inf, pad_sup), 'constant', constant_values=0)

lab = np.lib.pad(lab, zip(pad_inf, pad_sup), 'constant', constant_values=0)

"""

Specialized dataset class for the brain parcellation data.

Attributes:

vx(array n_data x 3): Array containing the coordinates x, y, z of the voxels

file_ids(array n_data x 3): Array containing the file id of the datapoint

"""

def __init__(self):

Dataset.__init__(self)

# Initialize the additional containers

self.vx = None

self.file_ids = None

def populate_from_config(self, config):

data_generator = DataGeneratorBrain()

data_generator.init_from_config(config)

vx, inputs, outputs, file_ids = data_generator.generate_parallel(config.general["n_data"])

self.populate(inputs, outputs, vx, file_ids)

self.shuffle_data()

def populate(self, inputs, outputs, vx, file_ids):

self.inputs = inputs

self.outputs = outputs

self.vx = vx

self.file_ids = file_ids

def shuffle_data_virtual(self, perm):

self.vx = self.vx[perm]

self.file_ids = self.file_ids[perm]

def write_virtual(self, h5file):

h5file.create_dataset("voxels", data=self.vx, dtype='f')

h5file.create_dataset("file_id", data=self.file_ids, dtype='f')

def read_virtual(self, h5file):

self.vx = h5file["voxels"].value

self.file_ids = h5file["file_id"].value

def duplicate_datapoints_slice_virtual(self, ds, slice_idx):

ds.vx = self.vx[slice_idx]

ds.file_ids = self.file_ids

def __init__(self, n_regions):

self.n_regions = n_regions

self.barycentres = np.zeros((n_regions, 3))

def update_barycentres(self, vxs, regions):

self.barycentres = np.zeros((self.n_regions, 3))

for i in xrange(self.n_regions):

idxs = regions == i+1

if vxs[idxs].size == 0:

continue

self.barycentres[i] = np.mean(vxs[idxs], axis=0)

# For zero values (with no regions present), set them to the mean

self.barycentres[self.barycentres == 0] = self.barycentres[self.barycentres != 0].mean()

def compute_scaled_distances(self, vx):

distances = np.linalg.norm(self.barycentres - vx, axis=1)

file_path = config.general["file_path"]

ds = DatasetBrainParcellation()

ds.populate_from_config(config)

ds.write(file_path)

# for ignored_label in self.ignored_labels:

# lab[lab == ignored_label] = 0

# for idx, label in enumerate(self.true_labels):

# lab[lab==label] = idx+1

#

# aa = nib.Nifti1Image(mri, nib.load(mri_file).get_affine())

# nib.save(aa, mri_file)

#

# bb = nib.Nifti1Image(lab, nib.load(lab_file).get_affine())

# nib.save(bb, lab_file)