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 run_validation_case_from_image_simple(output_dir,
model,
processed_image,
patch_shape,
image_gt=None,
image_pred=None,
overlap_factor=0.8,
prev_truth_index=None,
prev_truth_size=None,
pred_index=None,
pred_size=None,
is3d=False):
"""
:param output_dir: folder to save prediction results in
:param model: loaded trained model
:param processed_image: as a numpy ndarray, after any processing (scaling, normalization, augmentation) wanted
:param patch_shape: size of patch on which to run model
:param image_gt: optional, {0,1} ndarray of ground truth segmentation
:param image_pred: optional, {0-1} ndarray of predicted segmentation
:param overlap_factor: amount of overlap between consecutive patches, float 0-1
:param prev_truth_index: if truth used in prediction - the starting index in input batch (depth), else None
:param prev_truth_size:if truth used in prediction - amount of truth slices in input batch, else None
:param pred_index: if other pred used in prediction - the starting index in input batch (depth), else None
:param pred_size: if other pred used in prediction - amount of pred slices in input batch (depth), else None
:return: prediction + path to saved location
"""
prediction, _ = \
patch_wise_prediction(model=model, data=np.expand_dims(processed_image.squeeze(), 0), overlap_factor=overlap_factor,
patch_shape=patch_shape, truth_data=image_gt, prev_truth_index=prev_truth_index,
prev_truth_size=prev_truth_size, pred_data=image_pred, pred_index=pred_index,
pred_size=pred_size, is3d=is3d) # [np.newaxis]
prediction = prediction.squeeze()
prediction_image = get_image(prediction)
filename = os.path.join(output_dir, "prediction.nii.gz")
name_counter = 0
while os.path.exists(filename):
name_counter += 1
filename = os.path.join(output_dir,
"prediction_{0}.nii.gz".format(name_counter))
print("Saving to {}".format(filename))
prediction_image.to_filename(filename)
return prediction, filename
def save_nifti(data, path):
nifti = get_image(data)
nib.save(nifti, path)
def augment_data(data,
truth,
data_min,
data_max,
mask=None,
scale_deviation=None,
iso_scale_deviation=None,
rotate_deviation=None,
translate_deviation=None,
flip=None,
contrast_deviation=None,
poisson_noise=None,
gaussian_noise=None,
speckle_noise=None,
piecewise_affine=None,
elastic_transform=None,
intensity_multiplication_range=None,
gaussian_filter=None,
coarse_dropout=None,
data_range=None,
truth_range=None,
prev_truth_range=None):
n_dim = len(truth.shape)
if scale_deviation:
scale_factor = random_scale_factor(n_dim, std=scale_deviation)
else:
scale_factor = [1, 1, 1]
if iso_scale_deviation:
iso_scale_factor = np.random.uniform(1, iso_scale_deviation["max"])
if random_boolean():
iso_scale_factor = 1 / iso_scale_factor
scale_factor[0] *= iso_scale_factor
scale_factor[1] *= iso_scale_factor
else:
iso_scale_factor = None
if rotate_deviation:
rotate_factor = random_rotation_angle(n_dim, std=rotate_deviation)
rotate_factor = np.deg2rad(rotate_factor)
else:
rotate_factor = None
if flip is not None and flip:
flip_axis = random_flip_dimensions(n_dim, flip)
else:
flip_axis = None
if translate_deviation is not None:
translate_factor = random_translate_factor(
n_dim, -np.array(translate_deviation),
np.array(translate_deviation))
translate_factor[-1] = np.floor(
translate_factor[-1]) # z-translate should be int
else:
translate_factor = None
if contrast_deviation is not None:
val_range = data_max - data_min
contrast_min_val = data_min + contrast_deviation[
"min_factor"] * np.random.uniform(-1, 1) * val_range
contrast_max_val = data_max + contrast_deviation[
"max_factor"] * np.random.uniform(-1, 1) * val_range
else:
contrast_min_val, contrast_max_val = None, None
if poisson_noise is not None:
apply_poisson_noise = poisson_noise > np.random.random()
else:
apply_poisson_noise = False
if gaussian_noise is not None:
apply_gaussian_noise = gaussian_noise["prob"] > np.random.random()
else:
apply_gaussian_noise = False
if speckle_noise is not None:
apply_speckle_noise = speckle_noise["prob"] > np.random.random()
else:
apply_speckle_noise = False
if gaussian_filter is not None and gaussian_filter["prob"] > 0:
gaussian_sigma = gaussian_filter["max_sigma"] * np.random.random()
apply_gaussian = gaussian_filter["prob"] > np.random.random()
else:
apply_gaussian, gaussian_sigma = False, None
if piecewise_affine is not None:
piecewise_affine_scale = np.random.random() * piecewise_affine["scale"]
else:
piecewise_affine_scale = 0
if (elastic_transform is not None) and (elastic_transform["alpha"] > 0):
elastic_transform_scale = np.random.random(
) * elastic_transform["alpha"]
else:
elastic_transform_scale = 0
if intensity_multiplication_range is not None:
a, b = intensity_multiplication_range
intensity_multiplication = np.random.random() * (b - a) + a
else:
intensity_multiplication = 1
if coarse_dropout is not None:
coarse_dropout_rate = coarse_dropout['rate']
coarse_dropout_size = coarse_dropout['size_percent']
image, affine = data, np.eye(4)
distorted_data, distorted_affine = distort_image(
image,
affine,
flip_axis=flip_axis,
scale_factor=scale_factor,
rotate_factor=rotate_factor,
translate_factor=translate_factor)
if data_range is None:
data = resample_to_img(get_image(distorted_data, distorted_affine),
image,
interpolation="continuous",
copy=False,
clip=True).get_fdata()
else:
data = interpolate_affine_range(distorted_data,
distorted_affine,
data_range,
order=1,
mode='constant',
cval=data_min)
truth_image, truth_affine = truth, np.eye(4)
distorted_truth_data, distorted_truth_affine = distort_image(
truth_image,
truth_affine,
flip_axis=flip_axis,
scale_factor=scale_factor,
rotate_factor=rotate_factor,
translate_factor=translate_factor)
if truth_range is None:
truth_data = resample_to_img(get_image(distorted_truth_data,
distorted_truth_affine),
truth_image,
interpolation="nearest",
copy=False,
clip=True).get_data()
else:
truth_data = interpolate_affine_range(distorted_truth_data,
distorted_truth_affine,
truth_range,
order=0,
mode='constant',
cval=0)
if prev_truth_range is None:
prev_truth_data = None
else:
prev_truth_data = interpolate_affine_range(distorted_truth_data,
distorted_truth_affine,
prev_truth_range,
order=0,
mode='constant',
cval=0)
if mask is None:
mask_data = None
else:
mask_image, mask_affine = mask, np.eye(4)
distorted_mask_data, distorted_mask_affine = distort_image(
mask_image,
mask_affine,
flip_axis=flip_axis,
scale_factor=scale_factor,
rotate_factor=rotate_factor,
translate_factor=translate_factor)
if truth_range is None:
mask_data = resample_to_img(get_image(distorted_mask_data,
distorted_mask_affine),
mask_image,
interpolation="nearest",
copy=False,
clip=True).get_data()
else:
mask_data = interpolate_affine_range(distorted_mask_data,
distorted_mask_affine,
truth_range,
order=0,
mode='constant',
cval=0)
if piecewise_affine_scale > 0:
data, truth_data, prev_truth_data, mask_data = apply_piecewise_affine(
data, truth_data, prev_truth_data, mask_data,
piecewise_affine_scale)
if elastic_transform_scale > 0:
data, truth_data, prev_truth_data, mask_data = apply_elastic_transform(
data, truth_data, prev_truth_data, mask_data,
elastic_transform_scale, elastic_transform["sigma"])
if contrast_deviation is not None:
data = contrast_augment(data, contrast_min_val, contrast_max_val)
if intensity_multiplication != 1:
data = data * intensity_multiplication
if apply_gaussian:
data = apply_gaussian_filter(data, gaussian_sigma)
if apply_poisson_noise:
data = shot_noise(data)
if apply_speckle_noise:
data = add_speckle_noise(data, speckle_noise["sigma"])
if apply_gaussian_noise:
data = add_gaussian_noise(data, gaussian_noise["sigma"])
if coarse_dropout is not None:
data = apply_coarse_dropout(data,
rate=coarse_dropout_rate,
size_percent=coarse_dropout_size,
per_channel=coarse_dropout["per_channel"])
return data, truth_data, prev_truth_data, mask_data
def main(pred_dir,
config,
split='test',
overlap_factor=1,
preprocess_method=None):
padding = [16, 16, 8]
prediction2_dir = os.path.abspath(
os.path.join(config['base_dir'], 'predictions2', split))
model = load_old_model(get_last_model_path(config["model_file"]))
with open(os.path.join(opts.config_dir, 'norm_params.json'), 'r') as f:
norm_params = json.load(f)
for sample_folder in glob(os.path.join(pred_dir, split, '*')):
mask_path = os.path.join(sample_folder, 'prediction.nii.gz')
truth_path = os.path.join(sample_folder, 'truth.nii.gz')
subject_id = Path(sample_folder).name
dest_folder = os.path.join(prediction2_dir, subject_id)
Path(dest_folder).mkdir(parents=True, exist_ok=True)
truth = nib.load(truth_path)
nib.save(truth, os.path.join(dest_folder, Path(truth_path).name))
mask = nib.load(mask_path)
mask = process_pred(mask.get_data(), gaussian_std=0.5, threshold=0.5)
bbox_start, bbox_end = find_bounding_box(mask)
check_bounding_box(mask, bbox_start, bbox_end)
if padding is not None:
bbox_start = np.maximum(bbox_start - padding, 0)
bbox_end = np.minimum(bbox_end + padding, mask.shape)
print("BBox: {}-{}".format(bbox_start, bbox_end))
volume = nib.load(
os.path.join(original_data_folder, subject_id, 'volume.nii'))
orig_volume_shape = np.array(volume.get_data().shape)
volume = cut_bounding_box(volume, bbox_start,
bbox_end).get_data().astype(np.float)
if preprocess_method is not None:
print('Applying preprocess by {}...'.format(preprocess_method))
if preprocess_method == 'window_1_99':
volume = window_intensities_data(volume)
else:
raise Exception(
'Unknown preprocess: {}'.format(preprocess_method))
if norm_params is not None and any(norm_params.values()):
volume = normalize_data(volume,
mean=norm_params['mean'],
std=norm_params['std'])
prediction = patch_wise_prediction(model=model,
data=np.expand_dims(volume, 0),
patch_shape=config["patch_shape"] +
[config["patch_depth"]],
overlap_factor=overlap_factor)
prediction = prediction.squeeze()
padding2 = list(zip(bbox_start, orig_volume_shape - bbox_end))
print(padding2)
prediction = np.pad(prediction,
padding2,
mode='constant',
constant_values=0)
assert all(
[s1 == s2 for s1, s2 in zip(prediction.shape, orig_volume_shape)])
prediction = get_image(prediction)
nib.save(prediction, os.path.join(dest_folder, Path(mask_path).name))
test_data = np.asarray([data_file.root.data[sid]])
test_truth_data = np.asarray([data_file.root.truth[sid]])
# Get all predictions
for i, model in enumerate(step_1_networks):
config = step_1_configs[i]
# step 1 - no use of context for prediction
if config["use_augmentations"]: # TODO - add this key to configs, default False
prediction = predict_augment(data=test_data, model=model, overlap_factor=overlap_factor,
patch_shape=config["patch_shape"])
else:
prediction, _ = \
patch_wise_prediction(model=model, data=test_data, overlap_factor=overlap_factor,
patch_shape=config["patch_shape"], permute=config["augment"]["permute"])
prediction = prediction.squeeze()
prediction_image = get_image(prediction) # NIB format
filename = os.path.join(output_prediction_dir, subject_id, "prediction_{}.nii.gz".format(i))
prediction_image.to_filename(filename)
if i > 0:
avg_pred += prediction
else:
avg_pred = prediction
# Get the final averaged prediction and write to file
avg_pred /= n_step_1_models
nib.save(nib.Nifti1Image(avg_pred, prediction_image.affine, header=prediction_image.header),
os.path.join(output_prediction_dir, subject_id, f'averaged_prediction.nii'))
pred_storage[sid] = np.asarray(avg_pred).astype(np.float)
bin_avg_pred = adapted_postprocess_pred(os.path.join(output_prediction_dir, subject_id, f'averaged_prediction.nii'))
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,
is3d=False,
pred_index=None,
pred_size=None,
use_augmentations=False,
scale_xy=None,
resolution_file=''):
"""
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 scale_xy: wether to attempt and scale the image to main resolution
:param resolution_file: a file containing a dict of all existing scans' resolutions
:param model:
"""
cur_subject_id = data_file.root.subject_ids[data_index].decode()
if not os.path.exists(output_dir):
os.makedirs(output_dir)
test_data = np.asarray([data_file.root.data[data_index]])
#test_data = scale_data(test_data, cur_subject_id, dict_pkl=resolution_file, scale_xy=scale_xy)
if prev_truth_index is not None:
test_truth_data = np.asarray([data_file.root.truth[data_index]])
#test_truth_data = scale_data(test_truth_data, cur_subject_id, dict_pkl=resolution_file, scale_xy=scale_xy)
else:
test_truth_data = None
if pred_index is not None:
test_pred_data = np.asarray([data_file.root.pred[data_index]])
#test_pred_data = scale_data(test_pred_data, cur_subject_id, dict_pkl=resolution_file, scale_xy=scale_xy)
else:
test_pred_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)))
try:
test_truth = np.asarray([data_file.root.truth[data_index]])
test_truth = get_image(
scale_data(test_truth,
cur_subject_id,
dict_pkl=resolution_file,
scale_xy=scale_xy))
test_truth.to_filename(os.path.join(output_dir, "truth.nii.gz"))
except:
pass
if patch_shape == test_data.shape[-3:]:
print("Warning - went in where it wasn't expected!!!!!")
prediction = predict(model, test_data, permute=permute)
else:
if use_augmentations:
prediction = predict_augment(model=model,
data=test_data,
overlap_factor=overlap_factor,
patch_shape=patch_shape,
permute=permute,
truth_data=test_truth_data,
prev_truth_index=prev_truth_index,
prev_truth_size=prev_truth_size,
pred_data=test_pred_data,
pred_index=pred_index,
pred_size=pred_size,
is3d=is3d)
else:
prediction, prediction_var = \
patch_wise_prediction(model=model, data=test_data, overlap_factor=overlap_factor,
patch_shape=patch_shape, permute=permute, truth_data=test_truth_data,
prev_truth_index=prev_truth_index, prev_truth_size=prev_truth_size,
pred_data=test_pred_data, pred_index=pred_index, pred_size=pred_size, is3d=is3d)
# if prediction.shape[-1] > 1:
# prediction = prediction[..., 1]
prediction = prediction.squeeze()
prediction_image = get_image(prediction)
# prediction_var = prediction_var.squeeze()
# prediction_var_image = get_image(prediction_var)
name_counter = 0
if isinstance(prediction_image, list):
for i, image in enumerate(prediction_image):
filename = os.path.join(output_dir,
"prediction_{0}.nii.gz".format(i + 1))
while os.path.exists(filename):
name_counter += 1
filename = os.path.join(
output_dir,
"prediction_{0}_{1}.nii.gz".format(i + 1, name_counter))
image.to_filename(filename)
else:
filename = os.path.join(output_dir, "prediction.nii.gz")
# var_fname = os.path.join(output_dir, "prediction_variance.nii.gz")
while os.path.exists(filename):
name_counter += 1
filename = os.path.join(
output_dir, "prediction_{0}.nii.gz".format(name_counter))
var_fname = os.path.join(
output_dir,
"prediction_variance_{0}.nii.gz".format(name_counter))
print("Saving to {}".format(filename))
prediction_image.to_filename(filename)
# prediction_var_image.to_filename(var_fname)
return filename
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 predict_augment(model,
data,
patch_shape,
num_augments=5,
overlap_factor=0,
batch_size=32,
is3d=False,
permute=False,
truth_data=None,
prev_truth_index=None,
prev_truth_size=None,
pred_data=None,
pred_index=None,
pred_size=None,
specific_slice=None):
data_max = data.max()
data_min = data.min()
data = data.squeeze()
curr_truth_data = None
curr_pred_data = None
order = 2
cval = np.percentile(data, q=1)
init_prediction, _ = \
patch_wise_prediction(model=model, data=data[np.newaxis, ...], overlap_factor=overlap_factor,
batch_size=batch_size, patch_shape=patch_shape, permute=permute,
truth_data=truth_data, prev_truth_index=prev_truth_index,
prev_truth_size=prev_truth_size, pred_data=pred_data, pred_index=pred_index,
pred_size=pred_size, is3d=is3d)
predictions = [init_prediction.squeeze()]
for i in range(num_augments):
### pixel-wise augmentations - don't need to apply to truth or predictions
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)
curr_data = data.copy()
### spatial augmentations - need to apply to truth or predictions
# rotate_factor = np.random.uniform(-30, 30)
rotate_factor = 0
to_flip = np.arange(0, 3)[np.random.choice([True, False], size=3)]
to_transpose = np.random.choice([True, False])
scale_factors = np.random.normal(1, [0.15, 0.15, 0], 3)
curr_data = flip_it(curr_data, to_flip)
curr_affine = scale_image(np.eye(4), scale_factors)
curr_data = resample_to_img(get_image(curr_data, curr_affine),
get_image(curr_data),
interpolation="continuous",
copy=False,
clip=True).get_fdata()
curr_truth_data = None
curr_pred_data = None
if truth_data is not None:
curr_truth_data = truth_data.copy().squeeze()
curr_truth_data = flip_it(curr_truth_data, to_flip)
curr_truth_data = resample_to_img(get_image(
curr_truth_data, curr_affine),
get_image(curr_truth_data),
interpolation="continuous",
copy=False,
clip=True).get_fdata()
if pred_data is not None:
curr_pred_data = pred_data.copy().squeeze()
curr_pred_data = flip_it(curr_pred_data, to_flip)
curr_pred_data = resample_to_img(get_image(curr_pred_data,
curr_affine),
get_image(curr_pred_data),
interpolation="continuous",
copy=False,
clip=True).get_fdata()
if to_transpose:
curr_data = curr_data.transpose([1, 0, 2])
if truth_data is not None:
curr_truth_data = curr_truth_data.transpose([1, 0, 2])
if pred_data is not None:
curr_pred_data = curr_pred_data.transpose([1, 0, 2])
curr_data = ndimage.rotate(curr_data,
rotate_factor,
order=order,
reshape=False,
mode='constant',
cval=cval)
if truth_data is not None:
curr_truth_data = ndimage.rotate(curr_truth_data,
rotate_factor,
order=order,
reshape=False,
mode='constant',
cval=cval)
curr_truth_data = curr_truth_data[np.newaxis, ...]
if pred_data is not None:
curr_pred_data = ndimage.rotate(curr_pred_data,
rotate_factor,
order=order,
reshape=False,
mode='constant',
cval=cval)
curr_pred_data = curr_pred_data[np.newaxis, ...]
curr_prediction, _ = \
patch_wise_prediction(model=model, data=curr_data[np.newaxis, ...], overlap_factor=overlap_factor,
batch_size=batch_size, patch_shape=patch_shape, permute=permute,
truth_data=curr_truth_data, prev_truth_index=prev_truth_index,
prev_truth_size=prev_truth_size, pred_data=curr_pred_data, pred_index=pred_index,
pred_size=pred_size, is3d=is3d)
curr_prediction = curr_prediction.squeeze()
curr_prediction = ndimage.rotate(curr_prediction,
-rotate_factor,
order=0,
reshape=False,
mode='constant',
cval=0)
if to_transpose:
curr_prediction = curr_prediction.transpose([1, 0, 2])
rev_scale_factors = [1 / x for x in scale_factors]
curr_affine = scale_image(np.eye(4), rev_scale_factors)
curr_prediction = resample_to_img(get_image(curr_prediction,
curr_affine),
get_image(curr_prediction),
interpolation="continuous",
copy=False,
clip=True).get_fdata()
curr_prediction = flip_it(curr_prediction, to_flip)
predictions += [curr_prediction.squeeze()]
res = np.stack(predictions, axis=0)
return res
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
