import itertools
import os
import nibabel as nib
import numpy as np
import tables
from keras import Model
from scipy import ndimage
from tqdm import tqdm
from fetal.utils import get_last_model_path
from fetal_net.utils.threaded_generator import ThreadedGenerator
from fetal_net.utils.utils import get_image, list_load, pickle_load
from .augment import permute_data, generate_permutation_keys, reverse_permute_data, contrast_augment
from .training import load_old_model
from .utils.patches import get_patch_from_3d_data
def flip_it(data_, axes):
for ax in axes:
data_ = np.flip(data_, ax)
return data_
def predict_augment(data, model, overlap_factor, patch_shape, num_augments=32):
data_max = data.max()
data_min = data.min()
data = data.squeeze()
order = 2
predictions = []
for _ in range(num_augments):
# pixel-wise augmentations
val_range = data_max - data_min
contrast_min_val = data_min + 0.10 * np.random.uniform(-1, 1) * val_range
contrast_max_val = data_max + 0.10 * np.random.uniform(-1, 1) * val_range
curr_data = contrast_augment(data, contrast_min_val, contrast_max_val)
# spatial augmentations
rotate_factor = np.random.uniform(-30, 30)
to_flip = np.arange(0, 3)[np.random.choice([True, False], size=3)]
to_transpose = np.random.choice([True, False])
curr_data = flip_it(curr_data, to_flip)
if to_transpose:
curr_data = curr_data.transpose([1, 0, 2])
curr_data = ndimage.rotate(curr_data, rotate_factor, order=order, reshape=False)
curr_prediction = patch_wise_prediction(model=model, data=curr_data[np.newaxis, ...], overlap_factor=overlap_factor, patch_shape=patch_shape).squeeze()
curr_prediction = ndimage.rotate(curr_prediction, -rotate_factor)
if to_transpose:
curr_prediction = curr_prediction.transpose([1, 0, 2])
curr_prediction = flip_it(curr_prediction, to_flip)
predictions += [curr_prediction.squeeze()]
res = np.stack(predictions, axis=0)
return res
def predict_flips(data, model, overlap_factor, config):
def powerset(iterable):
"powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
s = list(iterable)
return itertools.chain.from_iterable(itertools.combinations(s, r) for r in range(0, len(s) + 1))
def predict_it(data_, axes=()):
data_ = flip_it(data_, axes)
curr_pred = \
patch_wise_prediction(model=model,
data=np.expand_dims(data_.squeeze(), 0),
overlap_factor=overlap_factor,
patch_shape=config["patch_shape"] + [config["patch_depth"]]).squeeze()
curr_pred = flip_it(curr_pred, axes)
return curr_pred
predictions = []
for axes in powerset([0, 1, 2]):
predictions += [predict_it(data, axes).squeeze()]
return predictions
def get_set_of_patch_indices_full(start, stop, step):
indices = []
for start_i, stop_i, step_i in zip(start, stop, step):
indices_i = list(range(start_i, stop_i + 1, step_i))
if stop_i % step_i > 0:
indices_i += [stop_i]
indices += [indices_i]
return np.array(list(itertools.product(*indices)))
def batch_iterator(indices, batch_size, data_0, patch_shape, truth_0, prev_truth_index, truth_patch_shape):
i = 0
while i < len(indices):
batch = []
curr_indices = []
while len(batch) < batch_size and i < len(indices):
curr_index = indices[i]
patch = get_patch_from_3d_data(data_0, patch_shape=patch_shape, patch_index=curr_index)
if truth_0 is not None:
truth_index = list(curr_index[:2]) + [curr_index[2] + prev_truth_index]
truth_patch = get_patch_from_3d_data(truth_0, patch_shape=truth_patch_shape,
patch_index=truth_index)
patch = np.concatenate([patch, truth_patch], axis=-1)
batch.append(patch)
curr_indices.append(curr_index)
i += 1
yield [batch, curr_indices]
# print('Finished! {}-{}'.format(i, len(indices)))
def patch_wise_prediction(model: Model, data, patch_shape, overlap_factor=0, batch_size=5,
permute=False, truth_data=None, prev_truth_index=None, prev_truth_size=None):
"""
:param truth_data:
:param permute:
:param overlap_factor:
:param batch_size:
:param model:
:param data:
:return:
"""
is3d = np.sum(np.array(model.output_shape[1:]) > 1) > 2
if is3d:
prediction_shape = model.output_shape[-3:]
else:
prediction_shape = model.output_shape[-3:-1] + (1,) # patch_shape[-3:-1] #[64,64]#
min_overlap = np.subtract(patch_shape, prediction_shape)
max_overlap = np.subtract(patch_shape, (1, 1, 1))
overlap = min_overlap + (overlap_factor * (max_overlap - min_overlap)).astype(np.int)
data_0 = np.pad(data[0],
[(np.ceil(_ / 2).astype(int), np.floor(_ / 2).astype(int)) for _ in
np.subtract(patch_shape, prediction_shape)],
mode='constant', constant_values=np.percentile(data[0], q=1))
pad_for_fit = [(np.ceil(_ / 2).astype(int), np.floor(_ / 2).astype(int)) for _ in
np.maximum(np.subtract(patch_shape, data_0.shape), 0)]
data_0 = np.pad(data_0,
[_ for _ in pad_for_fit],
'constant', constant_values=np.percentile(data_0, q=1))
if truth_data is not None:
truth_0 = np.pad(truth_data[0],
[(np.ceil(_ / 2).astype(int), np.floor(_ / 2).astype(int)) for _ in
np.subtract(patch_shape, prediction_shape)],
mode='constant', constant_values=0)
truth_0 = np.pad(truth_0, [_ for _ in pad_for_fit],
'constant', constant_values=0)
truth_patch_shape = list(patch_shape[:2]) + [prev_truth_size]
else:
truth_0 = None
truth_patch_shape = None
indices = get_set_of_patch_indices_full((0, 0, 0),
np.subtract(data_0.shape, patch_shape),
np.subtract(patch_shape, overlap))
b_iter = batch_iterator(indices, batch_size, data_0, patch_shape,
truth_0, prev_truth_index, truth_patch_shape)
tb_iter = iter(ThreadedGenerator(b_iter, queue_maxsize=50))
data_shape = list(data.shape[-3:] + np.sum(pad_for_fit, -1))
if is3d:
data_shape += [model.output_shape[1]]
else:
data_shape += [model.output_shape[-1]]
predicted_output = np.zeros(data_shape)
predicted_count = np.zeros(data_shape, dtype=np.int16)
with tqdm(total=len(indices)) as pbar:
for [curr_batch, batch_indices] in tb_iter:
curr_batch = np.asarray(curr_batch)
if is3d:
curr_batch = np.expand_dims(curr_batch, 1)
prediction = predict(model, curr_batch, permute=permute)
if is3d:
prediction = prediction.transpose([0, 2, 3, 4, 1])
else:
prediction = np.expand_dims(prediction, -2)
for predicted_patch, predicted_index in zip(prediction, batch_indices):
# predictions.append(predicted_patch)
x, y, z = predicted_index
x_len, y_len, z_len = predicted_patch.shape[:-1]
predicted_output[x:x + x_len, y:y + y_len, z:z + z_len, :] += predicted_patch
predicted_count[x:x + x_len, y:y + y_len, z:z + z_len] += 1
pbar.update(batch_size)
assert np.all(predicted_count > 0), 'Found zeros in count'
if np.sum(pad_for_fit) > 0:
# must be a better way :\
x_pad, y_pad, z_pad = [[None if p2[0] == 0 else p2[0],
None if p2[1] == 0 else -p2[1]] for p2 in pad_for_fit]
predicted_count = predicted_count[x_pad[0]: x_pad[1],
y_pad[0]: y_pad[1],
z_pad[0]: z_pad[1]]
predicted_output = predicted_output[x_pad[0]: x_pad[1],
y_pad[0]: y_pad[1],
z_pad[0]: z_pad[1]]
assert np.array_equal(predicted_count.shape[:-1], data[0].shape), 'prediction shape wrong'
return predicted_output / predicted_count
# return reconstruct_from_patches(predictions, patch_indices=indices, data_shape=data_shape)
def get_prediction_labels(prediction, threshold=0.5, labels=None):
n_samples = prediction.shape[0]
label_arrays = []
for sample_number in range(n_samples):
label_data = np.argmax(prediction[sample_number], axis=1)
label_data[np.max(prediction[sample_number], axis=0) < threshold] = 0
if labels:
for value in np.unique(label_data).tolist()[1:]:
label_data[label_data == value] = labels[value - 1]
label_arrays.append(np.array(label_data, dtype=np.uint8))
return label_arrays
def get_test_indices(testing_file):
return pickle_load(testing_file)
def predict_from_data_file(model, open_data_file, index):
return model.predict(open_data_file.root.data[index])
def predict_and_get_image(model, data, affine):
return nib.Nifti1Image(model.predict(data)[0, 0], affine)
def predict_from_data_file_and_get_image(model, open_data_file, index):
return predict_and_get_image(model, open_data_file.root.data[index], open_data_file.root.affine)
def predict_from_data_file_and_write_image(model, open_data_file, index, out_file):
image = predict_from_data_file_and_get_image(model, open_data_file, index)
image.to_filename(out_file)
def prediction_to_image(prediction, label_map=False, threshold=0.5, labels=None):
if prediction.shape[0] == 1:
data = prediction[0]
if label_map:
label_map_data = np.zeros(prediction[0, 0].shape, np.int8)
if labels:
label = labels[0]
else:
label = 1
label_map_data[data > threshold] = label
data = label_map_data
elif prediction.shape[1] > 1:
if label_map:
label_map_data = get_prediction_labels(prediction, threshold=threshold, labels=labels)
data = label_map_data[0]
else:
return multi_class_prediction(prediction)
else:
raise RuntimeError("Invalid prediction array shape: {0}".format(prediction.shape))
return get_image(data)
def multi_class_prediction(prediction, affine):
prediction_images = []
for i in range(prediction.shape[1]):
prediction_images.append(get_image(prediction[0, i]))
return prediction_images
def run_validation_case(data_index, output_dir, model, data_file, training_modalities, patch_shape,
overlap_factor=0, permute=False, prev_truth_index=None, prev_truth_size=None,
use_augmentations=False):
"""
Runs a test case and writes predicted images to file.
:param data_index: Index from of the list of test cases to get an image prediction from.
:param output_dir: Where to write prediction images.
:param output_label_map: If True, will write out a single image with one or more labels. Otherwise outputs
the (sigmoid) prediction values from the model.
:param threshold: If output_label_map is set to True, this threshold defines the value above which is
considered a positive result and will be assigned a label.
:param labels:
:param training_modalities:
:param data_file:
:param model:
"""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
test_data = np.asarray([data_file.root.data[data_index]])
if prev_truth_index is not None:
test_truth_data = np.asarray([data_file.root.truth[data_index]])
else:
test_truth_data = None
for i, modality in enumerate(training_modalities):
image = get_image(test_data[i])
image.to_filename(os.path.join(output_dir, "data_{0}.nii.gz".format(modality)))
test_truth = get_image(data_file.root.truth[data_index])
test_truth.to_filename(os.path.join(output_dir, "truth.nii.gz"))
if patch_shape == test_data.shape[-3:]:
prediction = predict(model, test_data, permute=permute)
else:
if use_augmentations:
prediction = predict_augment(data=test_data, model=model, overlap_factor=overlap_factor,
patch_shape=patch_shape)
else:
prediction = \
patch_wise_prediction(model=model, data=test_data, overlap_factor=overlap_factor,
patch_shape=patch_shape,
truth_data=test_truth_data, prev_truth_index=prev_truth_index,
prev_truth_size=prev_truth_size)[np.newaxis]
prediction = prediction.squeeze()
prediction_image = get_image(prediction)
if isinstance(prediction_image, list):
for i, image in enumerate(prediction_image):
image.to_filename(os.path.join(output_dir, "prediction_{0}.nii.gz".format(i + 1)))
else:
filename = os.path.join(output_dir, "prediction.nii.gz")
prediction_image.to_filename(filename)
return filename
def run_validation_cases(validation_keys_file, model_file, training_modalities, hdf5_file, patch_shape,
output_dir=".", overlap_factor=0, permute=False,
prev_truth_index=None, prev_truth_size=None, use_augmentations=False):
file_names = []
validation_indices = pickle_load(validation_keys_file)
model = load_old_model(get_last_model_path(model_file))
data_file = tables.open_file(hdf5_file, "r")
for index in validation_indices:
if 'subject_ids' in data_file.root:
case_directory = os.path.join(output_dir, data_file.root.subject_ids[index].decode('utf-8'))
else:
case_directory = os.path.join(output_dir, "validation_case_{}".format(index))
file_names.append(
run_validation_case(data_index=index, output_dir=case_directory, model=model, data_file=data_file,
training_modalities=training_modalities, overlap_factor=overlap_factor,
permute=permute, patch_shape=patch_shape, prev_truth_index=prev_truth_index,
prev_truth_size=prev_truth_size, use_augmentations=use_augmentations))
data_file.close()
return file_names
def predict(model, data, permute=False):
if permute:
predictions = list()
for batch_index in range(data.shape[0]):
predictions.append(predict_with_permutations(model, data[batch_index]))
return np.asarray(predictions)
else:
return model.predict(data)
def predict_with_permutations(model, data):
predictions = list()
for permutation_key in generate_permutation_keys():
temp_data = permute_data(data, permutation_key)[np.newaxis]
predictions.append(reverse_permute_data(model.predict(temp_data)[0], permutation_key))
return np.mean(predictions, axis=0)
