def secondary_prediction(mask, vol, config2, model2_path=None, preprocess_method2=None, norm_params2=None, overlap_factor=0.9, augment2=None, num_augment=32, return_all_preds=False): model2 = load_old_model(get_last_model_path(model2_path), config=config2) 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 = get_prediction(data, model2, augment=augment2, num_augments=num_augment, return_all_preds=return_all_preds, overlap_factor=overlap_factor, config=config2) padding2 = list(zip(bbox_start, np.array(vol.shape) - bbox_end)) if return_all_preds: padding2 = [(0, 0)] + padding2 print(padding2) print(prediction.shape) prediction = np.pad(prediction, padding2, mode='constant', constant_values=0) return prediction
def main(input_path, output_path, overlap_factor, config, model_path, preprocess_method=None, norm_params=None, augment=None, num_augment=0, config2=None, model2_path=None, preprocess_method2=None, norm_params2=None, augment2=None, num_augment2=0, z_scale=None, xy_scale=None, return_all_preds=False): print(model_path) model = load_old_model(get_last_model_path(model_path), config=config) print('Loading nifti from {}...'.format(input_path)) nifti = read_img(input_path) print('Predicting mask...') data = nifti.get_fdata().astype(np.float).squeeze() print('original_shape: ' + str(data.shape)) scan_name = Path(input_path).name.split('.')[0] if (z_scale is None): z_scale = 1.0 if (xy_scale is None): xy_scale = 1.0 if z_scale != 1.0 or xy_scale != 1.0: data = ndimage.zoom(data, [xy_scale, xy_scale, z_scale]) data = preproc_and_norm(data, preprocess_method, norm_params, scale=config.get('scale_data', None), preproc=config.get('preproc', None)) save_nifti(data, os.path.join(output_path, scan_name + '_data.nii.gz')) data = np.pad(data, 3, 'constant', constant_values=data.min()) print('Shape: ' + str(data.shape)) prediction = get_prediction(data=data, model=model, augment=augment, num_augments=num_augment, return_all_preds=return_all_preds, overlap_factor=overlap_factor, config=config) # unpad prediction = prediction[3:-3, 3:-3, 3:-3] # revert to original size if config.get('scale_data', None) is not None: prediction = ndimage.zoom(prediction.squeeze(), np.divide([1, 1, 1], config.get('scale_data', None)), order=0)[..., np.newaxis] save_nifti(prediction, os.path.join(output_path, scan_name + '_pred.nii.gz')) if z_scale != 1.0 or xy_scale != 1.0: prediction = ndimage.zoom(prediction.squeeze(), [1.0 / xy_scale, 1.0 / xy_scale, 1.0 / z_scale], order=1)[..., np.newaxis] # if prediction.shape[-1] > 1: # prediction = prediction[..., 1] if config2 is not None: prediction = prediction.squeeze() mask = process_pred(prediction, gaussian_std=0.5, threshold=0.5) # .astype(np.uint8) nifti = read_img(input_path) 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=overlap_factor, augment2=augment2, num_augment=num_augment2, return_all_preds=return_all_preds) save_nifti(prediction, os.path.join(output_path, scan_name + 'pred_roi.nii.gz')) print('Saving to {}'.format(output_path)) print('Finished.')
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 main(overwrite=False): # convert input images into an hdf5 file if overwrite or not os.path.exists(config["data_file"]): create_data_file(config) data_file_opened = open_data_file(config["data_file"]) seg_loss_func = getattr(fetal_net.metrics, config['loss']) dis_loss_func = getattr(fetal_net.metrics, config['dis_loss']) # instantiate new model seg_model_func = getattr(fetal_net.model, config['model_name']) gen_model = seg_model_func( input_shape=config["input_shape"], initial_learning_rate=config["initial_learning_rate"], **{ 'dropout_rate': config['dropout_rate'], 'loss_function': seg_loss_func, 'mask_shape': None if config["weight_mask"] is None else config["input_shape"], 'old_model_path': config['old_model'] }) dis_model_func = getattr(fetal_net.model, config['dis_model_name']) dis_model = dis_model_func( input_shape=[config["input_shape"][0] + config["n_labels"]] + config["input_shape"][1:], initial_learning_rate=config["initial_learning_rate"], **{ 'dropout_rate': config['dropout_rate'], 'loss_function': dis_loss_func }) if not overwrite \ and len(glob.glob(config["model_file"] + 'g_*.h5')) > 0: # dis_model_path = get_last_model_path(config["model_file"] + 'dis_') gen_model_path = get_last_model_path(config["model_file"] + 'g_') # print('Loading dis model from: {}'.format(dis_model_path)) print('Loading gen model from: {}'.format(gen_model_path)) # dis_model = load_old_model(dis_model_path) # gen_model = load_old_model(gen_model_path) # dis_model.load_weights(dis_model_path) gen_model.load_weights(gen_model_path) gen_model.summary() dis_model.summary() # Build "frozen discriminator" frozen_dis_model = Network(dis_model.inputs, dis_model.outputs, name='frozen_discriminator') frozen_dis_model.trainable = False inputs_real = Input(shape=config["input_shape"]) inputs_fake = Input(shape=config["input_shape"]) segs_real = Activation(None, name='seg_real')(gen_model(inputs_real)) segs_fake = Activation(None, name='seg_fake')(gen_model(inputs_fake)) valid = Activation(None, name='dis')(frozen_dis_model( Concatenate(axis=1)([segs_fake, inputs_fake]))) combined_model = Model(inputs=[inputs_real, inputs_fake], outputs=[segs_real, valid]) combined_model.compile(loss=[seg_loss_func, 'binary_crossentropy'], loss_weights=[1, config["gd_loss_ratio"]], optimizer=Adam(config["initial_learning_rate"])) combined_model.summary() # get training and testing generators train_generator, validation_generator, n_train_steps, n_validation_steps = get_training_and_validation_generators( data_file_opened, batch_size=config["batch_size"], data_split=config["validation_split"], overwrite=overwrite, validation_keys_file=config["validation_file"], training_keys_file=config["training_file"], test_keys_file=config["test_file"], n_labels=config["n_labels"], labels=config["labels"], patch_shape=(*config["patch_shape"], config["patch_depth"]), validation_batch_size=config["validation_batch_size"], augment=config["augment"], skip_blank_train=config["skip_blank_train"], skip_blank_val=config["skip_blank_val"], truth_index=config["truth_index"], truth_size=config["truth_size"], prev_truth_index=config["prev_truth_index"], prev_truth_size=config["prev_truth_size"], truth_downsample=config["truth_downsample"], truth_crop=config["truth_crop"], patches_per_epoch=config["patches_per_epoch"], categorical=config["categorical"], is3d=config["3D"], drop_easy_patches_train=config["drop_easy_patches_train"], drop_easy_patches_val=config["drop_easy_patches_val"]) # get training and testing generators _, semi_generator, _, _ = get_training_and_validation_generators( data_file_opened, batch_size=config["batch_size"], data_split=config["validation_split"], overwrite=overwrite, validation_keys_file=config["validation_file"], training_keys_file=config["training_file"], test_keys_file=config["test_file"], n_labels=config["n_labels"], labels=config["labels"], patch_shape=(*config["patch_shape"], config["patch_depth"]), validation_batch_size=config["validation_batch_size"], val_augment=config["augment"], skip_blank_train=config["skip_blank_train"], skip_blank_val=config["skip_blank_val"], truth_index=config["truth_index"], truth_size=config["truth_size"], prev_truth_index=config["prev_truth_index"], prev_truth_size=config["prev_truth_size"], truth_downsample=config["truth_downsample"], truth_crop=config["truth_crop"], patches_per_epoch=config["patches_per_epoch"], categorical=config["categorical"], is3d=config["3D"], drop_easy_patches_train=config["drop_easy_patches_train"], drop_easy_patches_val=config["drop_easy_patches_val"]) # start training scheduler = Scheduler(config["dis_steps"], config["gen_steps"], init_lr=config["initial_learning_rate"], lr_patience=config["patience"], lr_decay=config["learning_rate_drop"]) best_loss = np.inf for epoch in range(config["n_epochs"]): postfix = {'g': None, 'd': None} # , 'val_g': None, 'val_d': None} with tqdm(range(n_train_steps // config["gen_steps"]), dynamic_ncols=True, postfix={ 'gen': None, 'dis': None, 'val_gen': None, 'val_dis': None, None: None }) as pbar: for n_round in pbar: # train D outputs = np.zeros(dis_model.metrics_names.__len__()) for i in range(scheduler.get_dsteps()): real_patches, real_segs = next(train_generator) semi_patches, _ = next(semi_generator) d_x_batch, d_y_batch = input2discriminator( real_patches, real_segs, semi_patches, gen_model.predict(semi_patches, batch_size=config["batch_size"]), dis_model.output_shape) outputs += dis_model.train_on_batch(d_x_batch, d_y_batch) if scheduler.get_dsteps(): outputs /= scheduler.get_dsteps() postfix['d'] = build_dsc(dis_model.metrics_names, outputs) pbar.set_postfix(**postfix) # train G (freeze discriminator) outputs = np.zeros(combined_model.metrics_names.__len__()) for i in range(scheduler.get_gsteps()): real_patches, real_segs = next(train_generator) semi_patches, _ = next(validation_generator) g_x_batch, g_y_batch = input2gan(real_patches, real_segs, semi_patches, dis_model.output_shape) outputs += combined_model.train_on_batch( g_x_batch, g_y_batch) outputs /= scheduler.get_gsteps() postfix['g'] = build_dsc(combined_model.metrics_names, outputs) pbar.set_postfix(**postfix) # evaluate on validation set dis_metrics = np.zeros(dis_model.metrics_names.__len__(), dtype=float) gen_metrics = np.zeros(gen_model.metrics_names.__len__(), dtype=float) evaluation_rounds = n_validation_steps for n_round in range(evaluation_rounds): # rounds_for_evaluation: val_patches, val_segs = next(validation_generator) # D if scheduler.get_dsteps() > 0: d_x_test, d_y_test = input2discriminator( val_patches, val_segs, val_patches, gen_model.predict( val_patches, batch_size=config["validation_batch_size"]), dis_model.output_shape) dis_metrics += dis_model.evaluate( d_x_test, d_y_test, batch_size=config["validation_batch_size"], verbose=0) # G # gen_x_test, gen_y_test = input2gan(val_patches, val_segs, dis_model.output_shape) gen_metrics += gen_model.evaluate( val_patches, val_segs, batch_size=config["validation_batch_size"], verbose=0) dis_metrics /= float(evaluation_rounds) gen_metrics /= float(evaluation_rounds) # save the model and weights with the best validation loss if gen_metrics[0] < best_loss: best_loss = gen_metrics[0] print('Saving Model...') with open( os.path.join( config["base_dir"], "g_{}_{:.3f}.json".format(epoch, gen_metrics[0])), 'w') as f: f.write(gen_model.to_json()) gen_model.save_weights( os.path.join( config["base_dir"], "g_{}_{:.3f}.h5".format(epoch, gen_metrics[0]))) postfix['val_d'] = build_dsc(dis_model.metrics_names, dis_metrics) postfix['val_g'] = build_dsc(gen_model.metrics_names, gen_metrics) # pbar.set_postfix(**postfix) print('val_d: ' + postfix['val_d'], end=' | ') print('val_g: ' + postfix['val_g']) # pbar.refresh() # update step sizes, learning rates scheduler.update_steps(epoch, gen_metrics[0]) K.set_value(dis_model.optimizer.lr, scheduler.get_lr()) K.set_value(combined_model.optimizer.lr, scheduler.get_lr()) data_file_opened.close()
def main(overwrite=False): # convert input images into an hdf5 file if overwrite or not os.path.exists(config["data_file"]): create_data_file(config) data_file_opened = open_data_file(config["data_file"]) if not overwrite and len(glob.glob(config["model_file"] + '*.h5')) > 0: model_path = get_last_model_path(config["model_file"]) print('Loading model from: {}'.format(model_path)) model = load_old_model(model_path) else: # instantiate new model loss_func = getattr(fetal_net.metrics, config['loss']) model_func = getattr(fetal_net.model, config['model_name']) model = model_func( input_shape=config["input_shape"], initial_learning_rate=config["initial_learning_rate"], **{ 'dropout_rate': config['dropout_rate'], 'loss_function': loss_func, 'mask_shape': None if config["weight_mask"] is None else config["input_shape"], # TODO: change to output shape 'old_model_path': config['old_model'] }) if not overwrite and len(glob.glob(config["model_file"] + '*.h5')) > 0: model_path = get_last_model_path(config["model_file"]) print('Loading model from: {}'.format(model_path)) model.load_weights(model_path) model.summary() # get training and testing generators train_generator, validation_generator, n_train_steps, n_validation_steps = get_training_and_validation_generators( data_file_opened, batch_size=config["batch_size"], data_split=config["validation_split"], overwrite=overwrite, validation_keys_file=config["validation_file"], training_keys_file=config["training_file"], test_keys_file=config["test_file"], n_labels=config["n_labels"], labels=config["labels"], patch_shape=(*config["patch_shape"], config["patch_depth"]), validation_batch_size=config["validation_batch_size"], augment=config["augment"], skip_blank_train=config["skip_blank_train"], skip_blank_val=config["skip_blank_val"], truth_index=config["truth_index"], truth_size=config["truth_size"], prev_truth_index=config["prev_truth_index"], prev_truth_size=config["prev_truth_size"], truth_downsample=config["truth_downsample"], truth_crop=config["truth_crop"], patches_per_epoch=config["patches_per_epoch"], categorical=config["categorical"], is3d=config["3D"], drop_easy_patches_train=config["drop_easy_patches_train"], drop_easy_patches_val=config["drop_easy_patches_val"]) # run training train_model(model=model, model_file=config["model_file"], training_generator=train_generator, validation_generator=validation_generator, steps_per_epoch=n_train_steps, validation_steps=n_validation_steps, initial_learning_rate=config["initial_learning_rate"], learning_rate_drop=config["learning_rate_drop"], learning_rate_patience=config["patience"], early_stopping_patience=config["early_stop"], n_epochs=config["n_epochs"], output_folder=config["base_dir"]) data_file_opened.close()