def labels_to_image_model(labels_shape, n_channels, generation_labels, output_labels, n_neutral_labels, atlas_res, target_res, output_shape=None, output_div_by_n=None, flipping=True, aff=None, scaling_bounds=0.15, rotation_bounds=15, shearing_bounds=0.012, translation_bounds=False, nonlin_std=4., nonlin_shape_factor=.0625, randomise_res=False, buil_distance_maps=False, data_res=None, thickness=None, downsample=False, blur_range=1.15, bias_field_std=.5, bias_shape_factor=.025): """ This function builds a keras/tensorflow model to generate images from provided label maps. The images are generated by sampling a Gaussian Mixture Model (of given parameters), conditionned on the label map. The model will take as inputs: -a label map -a vector containing the means of the Gaussian Mixture Model for each label, -a vector containing the standard deviations of the Gaussian Mixture Model for each label, -if apply_affine_deformation is True: a batch*(n_dims+1)*(n_dims+1) affine matrix -if apply_non_linear_deformation is True: a small non linear field of size batch*(dim_1*...*dim_n)*n_dims that will be resampled to labels size and integrated, to obtain a diffeomorphic elastic deformation. -if apply_bias_field is True: a small bias field of size batch*(dim_1*...*dim_n)*1 that will be resampled to labels size and multiplied to the image, to add a "bias-field" noise. The model returns: -the generated image normalised between 0 and 1. -the corresponding label map, with only the labels present in output_labels (the other are reset to zero). # IMPORTANT !!! # Each time we provide a parameter with separate values for each axis (e.g. with a numpy array or a sequence), # these values refer to the RAS axes. :param labels_shape: shape of the input label maps. Can be a sequence or a 1d numpy array. :param n_channels: number of channels to be synthetised. :param generation_labels: (optional) list of all possible label values in the input label maps. Default is None, where the label values are directly gotten from the provided label maps. If not None, can be a sequence or a 1d numpy array. It should be organised as follows: background label first, then non-sided labels (e.g. CSF, brainstem, etc.), then all the structures of the same hemisphere (can be left or right), and finally all the corresponding contralateral structures (in the same order). :param output_labels: list of all the label values to keep in the output label maps, in no particular order. Should be a subset of the values contained in generation_labels. Label values that are in generation_labels but not in output_labels are reset to zero. Can be a sequence or a 1d numpy array. By default output_labels is equal to generation_labels. :param n_neutral_labels: number of non-sided generation labels. :param atlas_res: resolution of the input label maps. Can be a number (isotropic resolution), a sequence, or a 1d numpy array. :param target_res: target resolution of the generated images and corresponding label maps. Can be a number (isotropic resolution), a sequence, or a 1d numpy array. :param output_shape: (optional) desired shape of the output image, obtained by randomly cropping the generated image Can be an integer (same size in all dimensions), a sequence, a 1d numpy array, or the path to a 1d numpy array. Default is None, where no cropping is performed. :param output_div_by_n: (optional) forces the output shape to be divisible by this value. It overwrites output_shape if necessary. Can be an integer (same size in all dimensions), a sequence, or a 1d numpy array. :param flipping: (optional) whether to introduce right/left random flipping :param aff: (optional) example of an (n_dims+1)x(n_dims+1) affine matrix of one of the input label map. Used to find brain's right/left axis. Should be given if flipping is True. :param scaling_bounds: (optional) range of the random saling to apply at each mini-batch. The scaling factor for each dimension is sampled from a uniform distribution of predefined bounds. Can either be: 1) a number, in which case the scaling factor is independently sampled from the uniform distribution of bounds [1-scaling_bounds, 1+scaling_bounds] for each dimension. 2) a sequence, in which case the scaling factor is sampled from the uniform distribution of bounds (1-scaling_bounds[i], 1+scaling_bounds[i]) for the i-th dimension. 3) a numpy array of shape (2, n_dims), in which case the scaling factor is sampled from the uniform distribution of bounds (scaling_bounds[0, i], scaling_bounds[1, i]) for the i-th dimension. 4) False, in which case scaling is completely turned off. Default is scaling_bounds = 0.15 (case 1) :param rotation_bounds: (optional) same as scaling bounds but for the rotation angle, except that for cases 1 and 2, the bounds are centred on 0 rather than 1, i.e. [0+rotation_bounds[i], 0-rotation_bounds[i]]. Default is rotation_bounds = 15. :param shearing_bounds: (optional) same as scaling bounds. Default is shearing_bounds = 0.012. :param translation_bounds: (optional) same as scaling bounds. Default is translation_bounds = False, but we encourage using it when cropping is deactivated (i.e. when output_shape=None in BrainGenerator). :param nonlin_std: (optional) Maximum value for the standard deviation of the normal distribution from which we sample the first tensor for synthesising the deformation field. Set to 0 if you wish to completely turn the elastic deformation off. :param nonlin_shape_factor: (optional) if nonlin_std is strictly positive, factor between the shapes of the input label maps and the shape of the input non-linear tensor. :param randomise_res: (optional) whether to mimic images that would have been 1) acquired at low resolution, and 2) resampled to high esolution. The low resolution is uniformly resampled at each minibatch from [1mm, 9mm]. In that process, the images generated by sampling the GMM are 1) blurred at the sampled LR, 2) downsampled at LR, and 3) resampled at target_resolution. :param data_res: (optional) specific acquisition resolution to mimic, as opposed to random resolution sampled when randomis_res is True. This triggers a blurring to mimic the specified acquisition resolution, but the downsampling is optional (see param downsample). Default for data_res is None, where images are slighlty blurred. If the generated images are uni-modal, data_res can be a number (isotropic acquisition resolution), a sequence, a 1d numpy array, or the path to a 1d numy array. In the multi-modal case, it should be given as a numpy array (or a path) of size (n_mod, n_dims), where each row is the acquisition resolution of the corresponding channel. :param thickness: (optional) if data_res is provided, we can further specify the slice thickness of the low resolution images to mimic. Must be provided in the same format as data_res. Default thickness = data_res. :param downsample: (optional) whether to actually downsample the volume images to data_res after blurring. Default is False, except when thickness is provided, and thickness < data_res. :param blur_range: (optional) Randomise the standard deviation of the blurring kernels, (whether data_res is given or not). At each mini_batch, the standard deviation of the blurring kernels are multiplied by a coefficient sampled from a uniform distribution with bounds [1/blur_range, blur_range]. If None, no randomisation. Default is 1.15. :param bias_field_std: (optional) If strictly positive, this triggers the corruption of synthesised images with a bias field. It is obtained by sampling a first small tensor from a normal distribution, resizing it to full size, and rescaling it to positive values by taking the voxel-wise exponential. bias_field_std designates the std dev of the normal distribution from which we sample the first tensor. Set to 0 to deactivate biad field corruption. :param bias_shape_factor: (optional) If bias_field_std is strictly positive, this designates the ratio between the size of the input label maps and the size of the first sampled tensor for synthesising the bias field. """ # reformat resolutions labels_shape = utils.reformat_to_list(labels_shape) n_dims, _ = utils.get_dims(labels_shape) atlas_res = utils.reformat_to_n_channels_array(atlas_res, n_dims, n_channels) data_res = atlas_res if ( data_res is None) else utils.reformat_to_n_channels_array( data_res, n_dims, n_channels) thickness = data_res if ( thickness is None) else utils.reformat_to_n_channels_array( thickness, n_dims, n_channels) downsample = utils.reformat_to_list( downsample, n_channels) if downsample else (np.min(thickness - data_res, 1) < 0) atlas_res = atlas_res[0] target_res = atlas_res if ( target_res is None) else utils.reformat_to_n_channels_array( target_res, n_dims)[0] # get shapes crop_shape, output_shape = get_shapes(labels_shape, output_shape, atlas_res, target_res, output_div_by_n) # create new_label_list and corresponding LUT to make sure that labels go from 0 to N-1 new_generation_labels, lut = utils.rearrange_label_list(generation_labels) # define model inputs labels_input = KL.Input(shape=labels_shape + [1], name='labels_input') means_input = KL.Input(shape=list(new_generation_labels.shape) + [n_channels], name='means_input') stds_input = KL.Input(shape=list(new_generation_labels.shape) + [n_channels], name='std_devs_input') # convert labels to new_label_list labels = l2i_et.convert_labels(labels_input, lut) # deform labels if (scaling_bounds is not False) | (rotation_bounds is not False) | (shearing_bounds is not False) | \ (translation_bounds is not False) | (nonlin_std > 0): labels._keras_shape = tuple(labels.get_shape().as_list()) labels = layers.RandomSpatialDeformation( scaling_bounds=scaling_bounds, rotation_bounds=rotation_bounds, shearing_bounds=shearing_bounds, translation_bounds=translation_bounds, nonlin_std=nonlin_std, nonlin_shape_factor=nonlin_shape_factor, inter_method='nearest')(labels) # cropping if crop_shape != labels_shape: labels._keras_shape = tuple(labels.get_shape().as_list()) labels = layers.RandomCrop(crop_shape)(labels) # flipping if flipping: assert aff is not None, 'aff should not be None if flipping is True' labels._keras_shape = tuple(labels.get_shape().as_list()) labels = layers.RandomFlip( get_ras_axes(aff, n_dims)[0], True, new_generation_labels, n_neutral_labels)(labels) # build synthetic image labels._keras_shape = tuple(labels.get_shape().as_list()) image = layers.SampleConditionalGMM()([labels, means_input, stds_input]) # apply bias field if bias_field_std > 0: image._keras_shape = tuple(image.get_shape().as_list()) image = layers.BiasFieldCorruption(bias_field_std, bias_shape_factor, False)(image) # intensity augmentation image._keras_shape = tuple(image.get_shape().as_list()) image = layers.IntensityAugmentation(clip=300, normalise=True, gamma_std=.4, separate_channels=True)(image) # loop over channels channels = list() split = KL.Lambda(lambda x: tf.split(x, [1] * n_channels, axis=-1))( image) if (n_channels > 1) else [image] for i, channel in enumerate(split): channel._keras_shape = tuple(channel.get_shape().as_list()) if randomise_res: max_res = np.array([9.] * 3) resolution, blur_res = layers.SampleResolution( atlas_res, max_res, .05, return_thickness=True)(means_input) sigma = l2i_et.blurring_sigma_for_downsampling(atlas_res, resolution, thickness=blur_res) channel = layers.DynamicGaussianBlur( 0.75 * max_res / np.array(atlas_res), blur_range)([channel, sigma]) if buil_distance_maps: channel, dist = layers.MimicAcquisition( atlas_res, atlas_res, output_shape, True)([channel, resolution]) channels.extend([channel, dist]) else: channel = layers.MimicAcquisition(atlas_res, atlas_res, output_shape, False)([channel, resolution]) channels.append(channel) else: sigma = l2i_et.blurring_sigma_for_downsampling( atlas_res, data_res[i], thickness=thickness[i]) channel = layers.GaussianBlur(sigma, blur_range)(channel) if downsample[i]: resolution = KL.Lambda(lambda x: tf.convert_to_tensor( data_res[i], dtype='float32'))([]) channel = layers.MimicAcquisition(atlas_res, data_res[i], output_shape)( [channel, resolution]) elif output_shape != crop_shape: channel = nrn_layers.Resize(size=output_shape)(channel) channels.append(channel) # concatenate all channels back image = KL.Lambda(lambda x: tf.concat(x, -1))( channels) if len(channels) > 1 else channels[0] # resample labels at target resolution if crop_shape != output_shape: labels = l2i_et.resample_tensor(labels, output_shape, interp_method='nearest') # convert labels back to original values and reset unwanted labels to zero labels = l2i_et.convert_labels(labels, generation_labels) labels._keras_shape = tuple(labels.get_shape().as_list()) reset_values = [v for v in generation_labels if v not in output_labels] labels = layers.ResetValuesToZero(reset_values, name='labels_out')(labels) # build model (dummy layer enables to keep the labels when plugging this model to other models) image = KL.Lambda(lambda x: x[0], name='image_out')([image, labels]) brain_model = Model(inputs=[labels_input, means_input, stds_input], outputs=[image, labels]) return brain_model
def build_augmentation_model(im_shape, n_channels, label_list, image_res, target_res=None, output_shape=None, output_div_by_n=None, n_neutral_labels=1, flipping=True, flip_rl_only=False, aff=None, scaling_bounds=0.15, rotation_bounds=15, enable_90_rotations=False, shearing_bounds=0.012, translation_bounds=False, nonlin_std=3., nonlin_shape_factor=.0625, bias_field_std=.3, bias_shape_factor=0.025, same_bias_for_all_channels=False, apply_intensity_augmentation=True, noise_std=1., augment_channels_separately=True): # reformat resolutions im_shape = utils.reformat_to_list(im_shape) n_dims, _ = utils.get_dims(im_shape) image_res = utils.reformat_to_list(image_res, length=n_dims) target_res = image_res if target_res is None else utils.reformat_to_list( target_res, length=n_dims) # get shapes cropping_shape, output_shape = get_shapes(im_shape, output_shape, image_res, target_res, output_div_by_n) im_shape = im_shape + [n_channels] # create new_label_list and corresponding LUT to make sure that labels go from 0 to N-1 new_label_list, lut = utils.rearrange_label_list(label_list) # define model inputs image_input = KL.Input(shape=im_shape, name='image_input') labels_input = KL.Input(shape=im_shape[:-1] + [1], name='labels_input', dtype='int32') # convert labels to new_label_list labels = l2i_et.convert_labels(labels_input, lut) # flipping if flipping: if flip_rl_only: labels, image = layers.RandomFlip( int(edit_volumes.get_ras_axes(aff, n_dims)[0]), [True, False], new_label_list, n_neutral_labels)([labels, image_input]) else: labels, image = layers.RandomFlip( None, [True, False], new_label_list, n_neutral_labels)([labels, image_input]) else: image = image_input # transform labels to soft prob. and concatenate them to the image labels = KL.Lambda(lambda x: tf.one_hot( tf.cast(x[..., 0], dtype='int32'), depth=len(label_list), axis=-1))( labels) image = KL.concatenate([image, labels], axis=len(im_shape)) # spatial deformation if (scaling_bounds is not False) | (rotation_bounds is not False) | (shearing_bounds is not False) | \ (translation_bounds is not False) | (nonlin_std > 0) | enable_90_rotations: image._keras_shape = tuple(image.get_shape().as_list()) image = layers.RandomSpatialDeformation( scaling_bounds=scaling_bounds, rotation_bounds=rotation_bounds, shearing_bounds=shearing_bounds, translation_bounds=translation_bounds, enable_90_rotations=enable_90_rotations, nonlin_std=nonlin_std, nonlin_shape_factor=nonlin_shape_factor)(image) # cropping if cropping_shape != im_shape[:-1]: image._keras_shape = tuple(image.get_shape().as_list()) image = layers.RandomCrop(cropping_shape)(image) # resampling (image blurred separately) if cropping_shape != output_shape: sigma = l2i_et.blurring_sigma_for_downsampling(image_res, target_res) split = KL.Lambda( lambda x: tf.split(x, [n_channels, -1], axis=len(im_shape)))(image) image = split[0] image._keras_shape = tuple(image.get_shape().as_list()) image = layers.GaussianBlur(sigma=sigma)(image) image = KL.concatenate([image, split[-1]]) image = l2i_et.resample_tensor(image, output_shape) # split tensor between image and labels image, labels = KL.Lambda( lambda x: tf.split(x, [n_channels, -1], axis=len(im_shape)), name='splitting')(image) # apply bias field if bias_field_std > 0: image._keras_shape = tuple(image.get_shape().as_list()) image = layers.BiasFieldCorruption(bias_field_std, bias_shape_factor, same_bias_for_all_channels)(image) # intensity augmentation if apply_intensity_augmentation: image._keras_shape = tuple(image.get_shape().as_list()) image = layers.IntensityAugmentation( noise_std, gamma_std=0.5, separate_channels=augment_channels_separately)(image) # build model im_trans_model = Model(inputs=[image_input, labels_input], outputs=[image, labels]) return im_trans_model
def labels_to_image_model(labels_shape, n_channels, generation_labels, output_labels, n_neutral_labels, atlas_res, target_res, output_shape=None, output_div_by_n=None, padding_margin=None, flipping=True, aff=None, apply_linear_trans=True, apply_nonlin_trans=True, nonlin_std=3., nonlin_shape_factor=.0625, blur_background=True, data_res=None, thickness=None, downsample=False, blur_range=1.15, crop_channel2=None, apply_bias_field=True, bias_field_std=.3, bias_shape_factor=.025): """ This function builds a keras/tensorflow model to generate images from provided label maps. The images are generated by sampling a Gaussian Mixture Model (of given parameters), conditionned on the label map. The model will take as inputs: -a label map -a vector containing the means of the Gaussian Mixture Model for each label, -a vector containing the standard deviations of the Gaussian Mixture Model for each label, -if apply_affine_deformation is True: a batch*(n_dims+1)*(n_dims+1) affine matrix -if apply_non_linear_deformation is True: a small non linear field of size batch*(dim_1*...*dim_n)*n_dims that will be resampled to labels size and integrated, to obtain a diffeomorphic elastic deformation. -if apply_bias_field is True: a small bias field of size batch*(dim_1*...*dim_n)*1 that will be resampled to labels size and multiplied to the image, to add a "bias-field" noise. The model returns: -the generated image normalised between 0 and 1. -the corresponding label map, with only the labels present in output_labels (the other are reset to zero). :param labels_shape: shape of the input label maps. Can be a sequence or a 1d numpy array. :param n_channels: number of channels to be synthetised. :param generation_labels: (optional) list of all possible label values in the input label maps. Default is None, where the label values are directly gotten from the provided label maps. If not None, can be a sequence or a 1d numpy array. It should be organised as follows: background label first, then non-sided labels (e.g. CSF, brainstem, etc.), then all the structures of the same hemisphere (can be left or right), and finally all the corresponding contralateral structures (in the same order). :param output_labels: list of all the label values to keep in the output label maps, in no particular order. Should be a subset of the values contained in generation_labels. Label values that are in generation_labels but not in output_labels are reset to zero. Can be a sequence or a 1d numpy array. :param n_neutral_labels: number of non-sided generation labels. :param atlas_res: resolution of the input label maps. Can be a number (isotropic resolution), a sequence, or a 1d numpy array. :param target_res: target resolution of the generated images and corresponding label maps. Can be a number (isotropic resolution), a sequence, or a 1d numpy array. :param output_shape: (optional) desired shape of the output image, obtained by randomly cropping the generated image Can be an integer (same size in all dimensions), a sequence, a 1d numpy array, or the path to a 1d numpy array. :param output_div_by_n: (optional) forces the output shape to be divisible by this value. It overwrites output_shape if necessary. Can be an integer (same size in all dimensions), a sequence, or a 1d numpy array. :param padding_margin: (optional) margin by which to pad the input labels with zeros. Padding is applied prior to any other operation. Can be an integer (same padding in all dimensions), a sequence, or a 1d numpy array. Default is no padding. :param flipping: (optional) whether to introduce right/left random flipping :param aff: (optional) example of an (n_dims+1)x(n_dims+1) affine matrix of one of the input label map. Used to find brain's right/left axis. Should be given if flipping is True. :param apply_linear_trans: (optional) whether to linearly deform the input label maps prior to generation. If true, the model will take an additional input of size batch*(n_dims+1)*(n_dims+1). Default is True. :param apply_nonlin_trans: (optional) whether to non-linearly deform the input label maps prior to generation. If true, the model will take an additional input of size batch*(dim_1*...*dim_n)*n_dims. Default is True. :param nonlin_std: (optional) If apply_nonlin_trans is True, maximum value for the standard deviation of the normal distribution from which we sample the first tensor for synthesising the deformation field. :param nonlin_shape_factor: (optional) if apply_non_linear_deformation is True, factor between the shapes of the input label maps and the shape of the input non-linear tensor. :param blur_background: (optional) If True, the background is blurred with the other labels, and can be reset to zero with a probability of 0.2. If False, the background is not blurred (we apply an edge blurring correction), and can be replaced by a low-intensity background. :param data_res: ((optional) acquisition resolution to mimick. If provided, the images sampled from the GMM are blurred to mimick data that would be: 1) acquired at the given acquisition resolution, and 2) resample at target_resolution. Default is None, where images are isotropically blurred to introduce some spatial correlation between voxels. If the generated images are uni-modal, data_res can be a number (isotropic acquisition resolution), a sequence, a 1d numpy array, or the path to a 1d numy array. In the multi-modal case, it should be given as a numpy array (or a path) of size (n_mod, n_dims), where each row is the acquisition resolution of the correspionding chanel. :param thickness: (optional) if data_res is provided, we can further specify the slice thickness of the low resolution images to mimick. If the generated images are uni-modal, data_res can be a number (isotropic acquisition resolution), a sequence, a 1d numpy array, or the path to a 1d numy array. In the multi-modal case, it should be given as a numpy array (or a path) of size (n_mod, n_dims), where each row is the acquisition resolution of the correspionding chanel. :param downsample: (optional) whether to actually downsample the volume image to data_res. Default is False, except when thickness is provided, and thickness < data_res. :param blur_range: (optional) Randomise the standard deviation of the blurring kernels, (whether data_res is given or not). At each mini_batch, the standard deviation of the blurring kernels are multiplied by a coefficient sampled from a uniform distribution with bounds [1/blur_range, blur_range]. If None, no randomisation. Default is 1.15. :param crop_channel2: (optional) stats for cropping second channel along the anterior-posterior axis. Should be a vector of length 4, with bounds of uniform distribution for cropping the front and back of the image (in percentage). None is no croppping. :param apply_bias_field: (optional) whether to apply a bias field to the generated image. If true, the model will take an additional input of size batch*(dim_1*...*dim_n)*1. Default is True. :param bias_field_std: (optional) If apply_nonlin_trans is True, maximum value for the standard deviation of the normal distribution from which we sample the first tensor for synthesising the deformation field. :param bias_shape_factor: (optional) if apply_bias_field is True, factor between the shapes of the input label maps and the shape of the input bias field tensor. """ # reformat resolutions labels_shape = utils.reformat_to_list(labels_shape) n_dims, _ = utils.get_dims(labels_shape) atlas_res = utils.reformat_to_n_channels_array(atlas_res, n_dims=n_dims, n_channels=n_channels) if data_res is None: # data_res assumed to be the same as the atlas data_res = atlas_res else: data_res = utils.reformat_to_n_channels_array(data_res, n_dims=n_dims, n_channels=n_channels) atlas_res = atlas_res[0] if downsample: # same as data_res if we want to actually downsample the synthetic image downsample_res = data_res else: # set downsample_res to None if downsampling is not necessary downsample_res = None if target_res is None: target_res = atlas_res else: target_res = utils.reformat_to_n_channels_array(target_res, n_dims)[0] thickness = utils.reformat_to_n_channels_array(thickness, n_dims=n_dims, n_channels=n_channels) # get shapes crop_shape, output_shape, padding_margin = get_shapes( labels_shape, output_shape, atlas_res, target_res, padding_margin, output_div_by_n) # create new_label_list and corresponding LUT to make sure that labels go from 0 to N-1 n_generation_labels = generation_labels.shape[0] new_generation_label_list, lut = utils.rearrange_label_list( generation_labels) # define model inputs labels_input = KL.Input(shape=labels_shape + [1], name='labels_input') means_input = KL.Input(shape=list(new_generation_label_list.shape) + [n_channels], name='means_input') std_devs_input = KL.Input(shape=list(new_generation_label_list.shape) + [n_channels], name='std_devs_input') list_inputs = [labels_input, means_input, std_devs_input] if apply_linear_trans: aff_in = KL.Input(shape=(n_dims + 1, n_dims + 1), name='aff_input') list_inputs.append(aff_in) else: aff_in = None # convert labels to new_label_list labels = l2i_et.convert_labels(labels_input, lut) # pad labels if padding_margin is not None: pad = np.transpose(np.array([[0] + padding_margin + [0]] * 2)) labels = KL.Lambda(lambda x: tf.pad( x, tf.cast(tf.convert_to_tensor(pad), dtype='int32')), name='pad')(labels) labels_shape = labels.get_shape().as_list()[1:n_dims + 1] # deform labels if apply_linear_trans | apply_nonlin_trans: labels = l2i_sp.deform_tensor(labels, aff_in, apply_nonlin_trans, 'nearest', nonlin_std, nonlin_shape_factor) labels = KL.Lambda(lambda x: tf.cast(x, dtype='int32'))(labels) # cropping if crop_shape != labels_shape: labels, _ = l2i_sp.random_cropping(labels, crop_shape, n_dims) if flipping: assert aff is not None, 'aff should not be None if flipping is True' labels, _ = l2i_sp.label_map_random_flipping( labels, new_generation_label_list, n_neutral_labels, aff, n_dims) # build synthetic image image = l2i_gmm.sample_gmm_conditioned_on_labels(labels, means_input, std_devs_input, n_generation_labels, n_channels) # loop over channels if n_channels > 1: split = KL.Lambda(lambda x: tf.split(x, [1] * n_channels, axis=-1))( image) else: split = [image] mask = KL.Lambda( lambda x: tf.where(tf.greater(x, 0), tf.ones_like(x, dtype='float32'), tf.zeros_like(x, dtype='float32')))(labels) processed_channels = list() for i, channel in enumerate(split): # reset edges of second channels to zero if (crop_channel2 is not None) & ( i == 1): # randomly crop sides of second channel channel, tmp_mask = l2i_sp.restrict_tensor( channel, axes=3, boundaries=crop_channel2) else: tmp_mask = None # blur channel if thickness is not None: sigma = utils.get_std_blurring_mask_for_downsampling( data_res[i], atlas_res, thickness=thickness[i]) else: sigma = utils.get_std_blurring_mask_for_downsampling( data_res[i], atlas_res) kernels_list = l2i_et.get_gaussian_1d_kernels( sigma, blurring_range=blur_range) channel = l2i_et.blur_channel(channel, mask, kernels_list, n_dims, blur_background) if (crop_channel2 is not None) & (i == 1): channel = KL.multiply([channel, tmp_mask]) # resample channel if downsample_res is not None: channel = l2i_et.resample_tensor(channel, output_shape, 'linear', downsample_res[i], atlas_res, n_dims=n_dims) else: if thickness is not None: diff = [ thickness[i][dim_idx] - data_res[i][dim_idx] for dim_idx in range(n_dims) ] if min(diff) < 0: channel = l2i_et.resample_tensor(channel, output_shape, 'linear', data_res[i], atlas_res, n_dims=n_dims) else: channel = l2i_et.resample_tensor(channel, output_shape, 'linear', None, atlas_res, n_dims) # apply bias field if apply_bias_field: channel = l2i_ia.bias_field_augmentation(channel, bias_field_std, bias_shape_factor) # intensity augmentation channel = KL.Lambda(lambda x: K.clip(x, 0, 300))(channel) channel = l2i_ia.min_max_normalisation(channel) processed_channels.append(l2i_ia.gamma_augmentation(channel, std=0.5)) # concatenate all channels back if n_channels > 1: image = KL.concatenate(processed_channels) else: image = processed_channels[0] # resample labels at target resolution if crop_shape != output_shape: labels = KL.Lambda(lambda x: tf.cast(x, dtype='float32'))(labels) labels = l2i_et.resample_tensor(labels, output_shape, interp_method='nearest', n_dims=3) # convert labels back to original values and reset unwanted labels to zero labels = l2i_et.convert_labels(labels, generation_labels) labels_to_reset = [ lab for lab in generation_labels if lab not in output_labels ] labels = l2i_et.reset_label_values_to_zero(labels, labels_to_reset) labels = KL.Lambda(lambda x: tf.cast(x, dtype='int32'), name='labels_out')(labels) # build model (dummy layer enables to keep the labels when plugging this model to other models) image = KL.Lambda(lambda x: x[0], name='image_out')([image, labels]) brain_model = keras.Model(inputs=list_inputs, outputs=[image, labels]) return brain_model
def build_augmentation_model(im_shape, n_channels, segmentation_labels, n_neutral_labels, atlas_res, target_res, output_shape=None, output_div_by_n=None, flipping=True, aff=None, scaling_bounds=0.15, rotation_bounds=15, shearing_bounds=0.012, translation_bounds=False, nonlin_std=3., nonlin_shape_factor=.0625, data_res=None, thickness=None, downsample=False, blur_range=1.03, bias_field_std=.5, bias_shape_factor=.025): # reformat resolutions and get shapes im_shape = utils.reformat_to_list(im_shape) n_dims, _ = utils.get_dims(im_shape) if data_res is not None: data_res = utils.reformat_to_n_channels_array(data_res, n_dims, n_channels) thickness = data_res if thickness is None else utils.reformat_to_n_channels_array( thickness, n_dims, n_channels) downsample = utils.reformat_to_list( downsample, n_channels) if downsample else np.min(thickness - data_res, 1) < 0 target_res = atlas_res if ( target_res is None) else utils.reformat_to_n_channels_array( target_res, n_dims)[0] else: target_res = atlas_res # get shapes crop_shape, output_shape = get_shapes(im_shape, output_shape, atlas_res, target_res, output_div_by_n) # define model inputs image_input = KL.Input(shape=im_shape + [n_channels], name='image_input') labels_input = KL.Input(shape=im_shape + [1], name='labels_input', dtype='int32') # deform labels labels, image = layers.RandomSpatialDeformation( scaling_bounds=scaling_bounds, rotation_bounds=rotation_bounds, shearing_bounds=shearing_bounds, translation_bounds=translation_bounds, nonlin_std=nonlin_std, nonlin_shape_factor=nonlin_shape_factor, inter_method=['nearest', 'linear'])([labels_input, image_input]) # cropping if crop_shape != im_shape: labels._keras_shape = tuple(labels.get_shape().as_list()) image._keras_shape = tuple(image.get_shape().as_list()) labels, image = layers.RandomCrop(crop_shape)([labels, image]) # flipping if flipping: assert aff is not None, 'aff should not be None if flipping is True' labels._keras_shape = tuple(labels.get_shape().as_list()) image._keras_shape = tuple(image.get_shape().as_list()) labels, image = layers.RandomFlip( get_ras_axes(aff, n_dims)[0], [True, False], segmentation_labels, n_neutral_labels)([labels, image]) # apply bias field if bias_field_std > 0: image._keras_shape = tuple(image.get_shape().as_list()) image = layers.BiasFieldCorruption(bias_field_std, bias_shape_factor, False)(image) # intensity augmentation image._keras_shape = tuple(image.get_shape().as_list()) image = layers.IntensityAugmentation(6, clip=False, normalise=True, gamma_std=.4, separate_channels=True)(image) # if necessary, loop over channels to 1) blur, 2) downsample to simulated LR, and 3) upsample to target if data_res is not None: channels = list() split = KL.Lambda(lambda x: tf.split(x, [1] * n_channels, axis=-1))( image) if (n_channels > 1) else [image] for i, channel in enumerate(split): # blur channel._keras_shape = tuple(channel.get_shape().as_list()) sigma = l2i_et.blurring_sigma_for_downsampling( atlas_res, data_res[i], thickness=thickness[i]) channel = layers.GaussianBlur(sigma, blur_range)(channel) # resample if downsample[i]: resolution = KL.Lambda(lambda x: tf.convert_to_tensor( data_res[i], dtype='float32'))([]) channel = layers.MimicAcquisition(atlas_res, data_res[i], output_shape)( [channel, resolution]) elif output_shape != crop_shape: channel = nrn_layers.Resize(size=output_shape)(channel) channels.append(channel) # concatenate all channels back image = KL.Lambda(lambda x: tf.concat(x, -1))( channels) if len(channels) > 1 else channels[0] # resample labels at target resolution if crop_shape != output_shape: labels = l2i_et.resample_tensor(labels, output_shape, interp_method='nearest') # build model (dummy layer enables to keep the labels when plugging this model to other models) labels = KL.Lambda(lambda x: tf.cast(x, dtype='int32'), name='labels_out')(labels) image = KL.Lambda(lambda x: x[0], name='image_out')([image, labels]) brain_model = models.Model(inputs=[image_input, labels_input], outputs=[image, labels]) return brain_model