import SimpleITK as sitk import tensorflow as tf from tensorflow.keras.models import load_model from focal_loss import BinaryFocalLoss import numpy as np import multiprocessing from functools import partial import os from sfransen.utils_quintin import * from sfransen.DWI_exp.helpers import * from sfransen.DWI_exp.preprocessing_function import preprocess from sfransen.DWI_exp.callbacks import dice_coef from sfransen.FROC.blob_preprocess import * from sfransen.FROC.cal_froc_from_np import * from sfransen.load_images import load_images_parrallel parser = argparse.ArgumentParser( description='Calculate the froc metrics and store in froc_metrics.yml') parser.add_argument('-experiment', help='Title of experiment') parser.add_argument('--series', '-s', metavar='[series_name]', required=True, nargs='+', help='List of series to include') args = parser.parse_args() ######## CUDA ################ os.environ["CUDA_VISIBLE_DEVICES"] = "2" ######## constants ############# SERIES = args.series series_ = '_'.join(args.series) EXPERIMENT = args.experiment MODEL_PATH = f'./../train_output/{EXPERIMENT}_{series_}/models/{EXPERIMENT}_{series_}.h5' YAML_DIR = f'./../train_output/{EXPERIMENT}_{series_}' DATA_DIR = "./../data/Nijmegen paths/" TARGET_SPACING = (0.5, 0.5, 3) INPUT_SHAPE = (192, 192, 24, len(SERIES)) IMAGE_SHAPE = INPUT_SHAPE[:3] DATA_SPLIT_INDEX = read_yaml_to_dict('./../data/Nijmegen paths/train_val_test_idxs.yml') TEST_INDEX = DATA_SPLIT_INDEX['val_set0'] N_CPUS = 12 ########## load images in parrallel ############## print_(f"> Loading images into RAM...") # read paths from txt image_paths = {} for s in SERIES: with open(path.join(DATA_DIR, f"{s}.txt"), 'r') as f: image_paths[s] = [l.strip() for l in f.readlines()] with open(path.join(DATA_DIR, f"seg.txt"), 'r') as f: seg_paths = [l.strip() for l in f.readlines()] num_images = len(seg_paths) # create pool of workers pool = multiprocessing.Pool(processes=N_CPUS) partial_images = partial(load_images_parrallel, seq = 'images', target_shape=IMAGE_SHAPE, target_space = TARGET_SPACING) partial_seg = partial(load_images_parrallel, seq = 'seg', target_shape=IMAGE_SHAPE, target_space = TARGET_SPACING) #load images images = [] for s in SERIES: image_paths_seq = image_paths[s] image_paths_index = np.asarray(image_paths_seq)[TEST_INDEX] data_list = pool.map(partial_images,image_paths_index) data = np.stack(data_list, axis=0) images.append(data) images_list = np.transpose(images, (1, 2, 3, 4, 0)) #load segmentations seg_paths_index = np.asarray(seg_paths)[TEST_INDEX] data_list = pool.map(partial_seg,seg_paths_index) segmentations = np.stack(data_list, axis=0) ########### load module ################## print(' >>>>>>> LOAD MODEL <<<<<<<<<') dependencies = { 'dice_coef': dice_coef } reconstructed_model = load_model(MODEL_PATH, custom_objects=dependencies) # reconstructed_model.summary(line_length=120) # make predictions on all TEST_INDEX print(' >>>>>>> START prediction <<<<<<<<<') predictions_blur = reconstructed_model.predict(images_list, batch_size=1) ############# preprocess ################# # preprocess predictions by removing the blur and making individual blobs print('>>>>>>>> START preprocess') def move_dims(arr): # UMCG numpy dimensions convention: dims = (batch, width, heigth, depth) # Joeran numpy dimensions convention: dims = (batch, depth, heigth, width) arr = np.moveaxis(arr, 3, 1) arr = np.moveaxis(arr, 3, 2) return arr # Joeran has his numpy arrays ordered differently. predictions_blur = move_dims(np.squeeze(predictions_blur)) segmentations = move_dims(np.squeeze(segmentations)) predictions = [preprocess_softmax(pred, threshold="dynamic")[0] for pred in predictions_blur] # Remove outer edges zeros = np.zeros(np.shape(predictions)) test = np.squeeze(predictions)[:,:,2:190,2:190] zeros[:,:,2:190,2:190] = test predictions = zeros # perform Froc metrics = evaluate(y_true=segmentations, y_pred=predictions) dump_dict_to_yaml(metrics, YAML_DIR, "froc_metrics", verbose=True) ############## save image as example ################# # save image nmr 3 IMAGE_DIR = f'./../train_output/train_10h_{series_}' img_s = sitk.GetImageFromArray(predictions_blur[3].squeeze()) sitk.WriteImage(img_s, f"{IMAGE_DIR}/predictions_blur_001.nii.gz") img_s = sitk.GetImageFromArray(predictions[3].squeeze()) sitk.WriteImage(img_s, f"{IMAGE_DIR}/predictions_001.nii.gz") img_s = sitk.GetImageFromArray(segmentations[3].squeeze()) sitk.WriteImage(img_s, f"{IMAGE_DIR}/segmentations_001.nii.gz")