def secondary_prediction(mask,
vol,
config2,
model2_path=None,
preprocess_method2=None,
norm_params2=None,
overlap_factor=0.9):
model2 = load_old_model(get_last_model_path(model2_path))
pred = mask
bbox_start, bbox_end = find_bounding_box(pred)
check_bounding_box(pred, bbox_start, bbox_end)
padding = [16, 16, 8]
if padding is not None:
bbox_start = np.maximum(bbox_start - padding, 0)
bbox_end = np.minimum(bbox_end + padding, mask.shape)
data = vol.astype(np.float)[bbox_start[0]:bbox_end[0],
bbox_start[1]:bbox_end[1],
bbox_start[2]:bbox_end[2]]
data = preproc_and_norm(data, preprocess_method2, norm_params2)
prediction = \
patch_wise_prediction(model=model2,
data=np.expand_dims(data, 0),
overlap_factor=overlap_factor,
patch_shape=config2["patch_shape"] + [config2["patch_depth"]])
prediction = prediction.squeeze()
padding2 = list(zip(bbox_start, np.array(vol.shape) - bbox_end))
print(padding2)
print(prediction.shape)
prediction = np.pad(prediction,
padding2,
mode='constant',
constant_values=0)
return prediction
def main(input_mat_path,
output_mat_path,
overlap_factor,
config,
model_path,
preprocess_method=None,
norm_params=None,
config2=None,
model2_path=None,
preprocess_method2=None,
norm_params2=None):
print(model_path)
model = load_old_model(get_last_model_path(model_path))
print('Loading mat from {}...'.format(input_mat_path))
mat = loadmat(input_mat_path)
print('Predicting mask...')
data = mat['volume'].astype(np.float)
data = preproc_and_norm(data, preprocess_method, norm_params)
prediction = \
patch_wise_prediction(model=model,
data=np.expand_dims(data, 0),
overlap_factor=overlap_factor,
patch_shape=config["patch_shape"] + [config["patch_depth"]])
print('Post-processing mask...')
if prediction.shape[-1] > 1:
prediction = prediction[..., 1]
prediction = prediction.squeeze()
print("Storing prediction in [7-9], 7 should be the best...")
mat['masks'][0, 9] = \
process_pred(prediction, gaussian_std=0, threshold=0.2) # .astype(np.uint8)
mat['masks'][0, 8] = \
process_pred(prediction, gaussian_std=1, threshold=0.5) # .astype(np.uint8)
mat['masks'][0, 7] = \
process_pred(prediction, gaussian_std=0.5, threshold=0.5) # .astype(np.uint8)
if config2 is not None:
print('Making secondary prediction... [6]')
prediction = secondary_prediction(
mat['masks'][0, 7],
vol=mat['volume'].astype(np.float),
config2=config2,
model2_path=model2_path,
preprocess_method2=preprocess_method2,
norm_params2=norm_params2,
overlap_factor=0.9)
mat['masks'][0, 6] = \
process_pred(prediction, gaussian_std=0, threshold=0.2) # .astype(np.uint8)
mat['masks'][0, 5] = \
process_pred(prediction, gaussian_std=1, threshold=0.5) # .astype(np.uint8)
mat['masks'][0, 4] = \
process_pred(prediction, gaussian_std=0.5, threshold=0.5) # .astype(np.uint8)
print('Saving mat to {}'.format(output_mat_path))
savemat(output_mat_path, mat)
print('Finished.')
def main(input_path,
output_path,
overlap_factor,
config,
model_path,
preprocess_method=None,
norm_params=None,
config2=None,
model2_path=None,
preprocess_method2=None,
norm_params2=None):
print(model_path)
model = load_old_model(get_last_model_path(model_path))
print('Loading mat from {}...'.format(input_path))
nifti = read_img(input_path)
print('Predicting mask...')
data = nifti.get_fdata().astype(np.float)
data = preproc_and_norm(data, preprocess_method, norm_params)
prediction = \
patch_wise_prediction(model=model,
data=np.expand_dims(data, 0),
overlap_factor=overlap_factor,
patch_shape=config["patch_shape"] + [config["patch_depth"]])
nib.save(prediction, os.path.join(output_path, 'pred.nii'))
print('Post-processing mask...')
if prediction.shape[-1] > 1:
prediction = prediction[..., 1]
prediction = prediction.squeeze()
print("Storing prediction in [7-9], 7 should be the best...")
mask = process_pred(prediction, gaussian_std=0.5,
threshold=0.5) # .astype(np.uint8)
if config2 is not None:
print('Making secondary prediction... [6]')
prediction = secondary_prediction(
mask,
vol=nifti.get_fdata().astype(np.float),
config2=config2,
model2_path=model2_path,
preprocess_method2=preprocess_method2,
norm_params2=norm_params2,
overlap_factor=0.9)
nib.save(prediction, os.path.join(output_path, 'pred_roi.nii'))
print('Saving to {}'.format(output_path))
print('Finished.')
def get_prediction(data, model, augment, num_augments, return_all_preds, overlap_factor, config):
if augment is not None:
patch_shape = config["patch_shape"] + [config["patch_depth"]]
if augment == 'all':
prediction = predict_augment(data, model=model, overlap_factor=overlap_factor, num_augments=num_augments, patch_shape=patch_shape)
elif augment == 'flip':
prediction = predict_flips(data, model=model, overlap_factor=overlap_factor, patch_shape=patch_shape, config=config)
else:
raise ("Unknown augmentation {}".format(augment))
if not return_all_preds:
prediction = np.median(prediction, axis=0)
else:
prediction = \
patch_wise_prediction(model=model,
data=np.expand_dims(data, 0),
overlap_factor=overlap_factor,
patch_shape=config["patch_shape"] + [config["patch_depth"]])
prediction = prediction.squeeze()
return prediction
def get_prediction(sid, model, config, slices_range=None):
# Assumes the data_file has been updated
test_data = np.asarray([data_file.root.data[sid]])
test_truth_data = np.asarray([data_file.root.truth[sid]])
test_pred_data = np.asarray([data_file.root.pred[sid]])
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"],
truth_data=test_truth_data, prev_truth_index=config["prev_truth_index"],
prev_truth_size=config["prev_truth_size"],
pred_data=test_pred_data, pred_index=config["pred_index"],
pred_size=config["pred_size"],
slices_range=slices_range
)
prediction = prediction.squeeze()
return prediction
def main(input_mat_path, output_mat_path, config, overlap_factor, model_path, preprocess_method=None, norm_params=None):
print(model_path)
model = load_old_model(get_last_model_path(model_path))
print('Loading mat from {}...'.format(input_mat_path))
mat = loadmat(input_mat_path)
print('Predicting mask...')
data = mat['volume'].astype(np.float)
if preprocess_method is not None:
print('Applying preprocess by {}...'.format(preprocess_method))
if preprocess_method == 'window_1_99':
data = window_intensities_data(data)
else:
raise Exception('Unknown preprocess: {}'.format(preprocess_method))
if norm_params is not None and any(norm_params.values()):
data = normalize_data(data, mean=norm_params['mean'], std=norm_params['std'])
prediction = \
patch_wise_prediction(model=model,
data=np.expand_dims(data, 0),
overlap_factor=overlap_factor,
patch_shape=config["patch_shape"] + [config["patch_depth"]])
print('Post-processing mask...')
if prediction.shape[-1] > 1:
prediction = prediction[..., 1]
prediction = prediction.squeeze()
mat['masks'][0, 9] = \
process_pred(prediction, gaussian_std=0, threshold=0.2) # .astype(np.uint8)
mat['masks'][0, 8] = \
process_pred(prediction, gaussian_std=1, threshold=0.5) # .astype(np.uint8)
mat['masks'][0, 7] = \
process_pred(prediction, gaussian_std=0.5, threshold=0.5) # .astype(np.uint8)
print('Saving mat to {}'.format(output_mat_path))
savemat(output_mat_path, mat)
print('Finished.')
def main(configs_folder,
data_folder,
model_folder,
subject_id,
validated_ind,
overlap_factor,
propagate,
pred_path,
flip=False):
vol, truth, affine, header = get_subject_id_data(data_folder, subject_id)
if flip:
vol = vol[:, :, ::-1]
truth = truth[:, :, ::-1]
if validated_ind >= 0:
validated_ind = truth.shape[-1] - validated_ind - 1
inflated_vol = np.expand_dims(vol, axis=0).astype(np.float)
inflated_vol, _, _ = normalize_data_storage_each(inflated_vol)
inflated_truth = np.expand_dims(truth, axis=0)
pred = nib.load(pred_path).get_data()
pred_bin = postprocess_prediction(pred)
if flip:
pred = pred[:, :, ::-1]
pred_bin = pred_bin[:, :, ::-1]
config, model = get_config_and_model(configs_folder, model_folder)
print(f"Initial DICE: {dice_coefficient_np(truth, pred_bin)}")
if validated_ind < 0:
validated_ind = get_placenta_start_ind(truth)
starting_slice = get_starting_slice_ind(validated_ind,
config["prev_truth_index"])
while not starting_slice:
validated_ind += 1
starting_slice = get_starting_slice_ind(validated_ind,
config["prev_truth_index"])
fixing_slice = starting_slice + config["truth_index"]
imsave('before.png', pred[:, :, fixing_slice])
if not starting_slice or starting_slice >= vol.shape[-1]:
print(f"Error! Validated slice {validated_ind} is out of range")
sys.exit(0)
new_pred, new_var = patch_wise_prediction(
model,
inflated_vol,
config["patch_shape"] + [config["patch_depth"]],
overlap_factor,
is3d=config["3D"],
truth_data=inflated_truth,
prev_truth_index=config["prev_truth_index"],
batch_size=16,
prev_truth_size=config["prev_truth_size"],
specific_slice=starting_slice)
new_pred = new_pred.squeeze()
# d_before = dice_coefficient_np(truth[:,:,fixing_slice], pred_bin[:,:,fixing_slice])
pred[:, :, fixing_slice] = new_pred[:, :, fixing_slice]
imsave('after.png', pred[:, :, fixing_slice])
pred_bin = postprocess_prediction(pred)
#print(f"DICE after updating slice {fixing_slice}: {dice_coefficient_np(truth, pred_bin)}")
validated_ind = fixing_slice
inflated_truth[0][:, :, validated_ind] = pred_bin[:, :, fixing_slice]
starting_slice = get_starting_slice_ind(validated_ind,
config["prev_truth_index"])
if propagate:
while starting_slice <= (vol.shape[-1] - config["patch_depth"]):
fixing_slice = starting_slice + config["truth_index"]
new_pred, new_var = patch_wise_prediction(
model,
inflated_vol,
config["patch_shape"] + [config["patch_depth"]],
overlap_factor,
is3d=config["3D"],
truth_data=inflated_truth,
prev_truth_index=config["prev_truth_index"],
prev_truth_size=config["prev_truth_size"],
specific_slice=starting_slice,
batch_size=16)
new_pred = new_pred.squeeze()
pred[:, :, fixing_slice] = new_pred[:, :, fixing_slice]
pred_bin = postprocess_prediction(pred)
#print(f"DICE after updating slice {fixing_slice}: {dice_coefficient_np(truth, pred_bin)}")
validated_ind = fixing_slice
inflated_truth[0][:, :, validated_ind] = pred_bin[:, :,
fixing_slice]
starting_slice = get_starting_slice_ind(validated_ind,
config["prev_truth_index"])
print(
f"DICE after updating slice {fixing_slice}: {dice_coefficient_np(truth, pred_bin)}"
)
if flip:
pred = pred[:, :, ::-1]
truth = truth[:, :, ::-1]
return pred, truth
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))
subject_id = data_file.root.subject_ids[sid]
print("In ID {}".format(subject_id))
os.mkdir(os.path.join(output_prediction_dir, subject_id))
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)
