--- a +++ b/prod/predict_mat2.py @@ -0,0 +1,155 @@ +import argparse +import json +import os + +import numpy as np + +from brats.utils import get_last_model_path +from fetal_net.normalize import normalize_data +from fetal_net.postprocess import postprocess_prediction as process_pred +from fetal_net.prediction import patch_wise_prediction +from fetal_net.training import load_old_model +from brats.preprocess import window_intensities_data + +from scipy.io import loadmat, savemat + +from fetal_net.utils.cut_relevant_areas import find_bounding_box, check_bounding_box + + +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 preproc_and_norm(data, preprocess_method, norm_params): + 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']) + return data + + +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 get_params(config_dir): + with open(os.path.join(config_dir, 'config.json'), 'r') as f: + __config = json.load(f) + with open(os.path.join(config_dir, 'norm_params.json'), 'r') as f: + __norm_params = json.load(f) + __model_path = os.path.join(config_dir, os.path.basename(__config['model_file'])) + return __config, __norm_params, __model_path + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("--input_mat", help="specifies mat file dir path", + type=str, required=True) + parser.add_argument("--output_mat", help="specifies mat file dir path", + type=str, required=True) + parser.add_argument("--overlap_factor", help="specifies overlap between prediction patches", + type=float, default=0.9) + + # Params for primary prediction + parser.add_argument("--config_dir", help="specifies config dir path", + type=str, required=True) + parser.add_argument("--preprocess", help="what preprocess to do", + type=str, required=False, default=None) + + # Params for secondary prediction + parser.add_argument("--config2_dir", help="specifies config dir path", + type=str, required=False, default=None) + parser.add_argument("--preprocess2", help="what preprocess to do", + type=str, required=False, default=None) + + opts = parser.parse_args() + + # 1 + _config, _norm_params, _model_path = get_params(opts.config_dir) + # 2 + if opts.config2_dir is not None: + _config2, _norm_params2, _model2_path = get_params(opts.config2_dir) + else: + _config2, _norm_params2, _model2_path = None, None, None + + main(opts.input_mat, opts.output_mat, overlap_factor=opts.overlap_factor, + config=_config, model_path=_model_path, preprocess_method=opts.preprocess, norm_params=_norm_params, + config2=_config2, model2_path=_model2_path, preprocess_method2=opts.preprocess2, norm_params2=_norm_params2)