def get_model(project_dir, train_seq, hparams, num_GPUs, logger, args):
"""
Initializes a tf.keras Model from MultiPlanarUNet.models as specified in
hparams['build']. If args.continue_training, the best previous model stored
in [project_dir]/models will be loaded.
If hparams["build"]["biased_output_layer"] is True, sets the bias weights
on the final conv. layer so that a zero-input gives an output of class
probabilities equal to the class frequencies of the training set.
If num_GPUs > 1, distribute the model on multiple GPUs
Args:
project_dir: A path to a MultiPlanarUNet project folder
train_seq: A MultiPlanarUNet.sequences object for the training data
hparams: A MultiPlanarUNet YAMLHParams object
num_GPUs: An integer, the number of GPUs to train the model on
logger: A MultiPlanarUNet logging object
args: argparse arguments
Returns:
model: The model to fit
org_model: The original, non-GPU-distributed model
(Same as model if num_GPUs==1)
"""
from MultiPlanarUNet.models import model_initializer
# Build new model (or continue training an existing one)
org_model = model_initializer(hparams=hparams,
continue_training=args.continue_training,
project_dir=project_dir,
logger=logger)
# Initialize weights in final layer?
if not args.continue_training and hparams["build"].get(
"biased_output_layer"):
from MultiPlanarUNet.utils.utils import set_bias_weights_on_all_outputs
set_bias_weights_on_all_outputs(org_model, train_seq, hparams, logger)
# Multi-GPU?
if num_GPUs > 1:
from tensorflow.keras.utils import multi_gpu_model
model = multi_gpu_model(org_model,
gpus=num_GPUs,
cpu_merge=False,
cpu_relocation=False)
logger("Creating multi-GPU model: N=%i" % num_GPUs)
else:
model = org_model
return model, org_model
def run(base_path, gpu_mon, num_GPUs, continue_training, force_GPU, just_one,
no_val, no_images, debug, wait_for, logger, train_images_per_epoch,
val_images_per_epoch, **kwargs):
from MultiPlanarUNet.train import Trainer, YAMLHParams
from MultiPlanarUNet.models import model_initializer
from MultiPlanarUNet.preprocessing import get_preprocessing_func
# Read in hyperparameters from YAML file
hparams = YAMLHParams(base_path + "/train_hparams.yaml", logger=logger)
validate_hparams(hparams)
# Wait for PID?
if wait_for:
from MultiPlanarUNet.utils import await_PIDs
await_PIDs(wait_for)
# Prepare Sequence generators and potential model specific hparam changes
f = get_preprocessing_func(hparams["build"].get("model_class_name"))
train, val, hparams = f(hparams,
logger=logger,
just_one=just_one,
no_val=no_val,
continue_training=continue_training,
base_path=base_path)
if gpu_mon:
# Wait for free GPU
if not force_GPU:
gpu_mon.await_and_set_free_GPU(N=num_GPUs, sleep_seconds=120)
else:
gpu_mon.set_GPUs = force_GPU
num_GPUs = len(force_GPU.split(","))
gpu_mon.stop()
# Build new model (or continue training an existing one)
org_model = model_initializer(hparams=hparams,
continue_training=continue_training,
project_dir=base_path,
logger=logger)
# Initialize weights in final layer?
if not continue_training and hparams["build"].get("biased_output_layer"):
from MultiPlanarUNet.utils.utils import set_bias_weights_on_all_outputs
set_bias_weights_on_all_outputs(org_model, train, hparams, logger)
# Multi-GPU?
if num_GPUs > 1:
from tensorflow.keras.utils import multi_gpu_model
model = multi_gpu_model(org_model,
gpus=num_GPUs,
cpu_merge=False,
cpu_relocation=False)
logger("Creating multi-GPU model: N=%i" % num_GPUs)
else:
model = org_model
# Init trainer
trainer = Trainer(model, logger=logger)
trainer.org_model = org_model
# Compile model
trainer.compile_model(n_classes=hparams["build"].get("n_classes"),
**hparams["fit"])
# Debug mode?
if debug:
from tensorflow.python import debug as tfdbg
from tensorflow.keras import backend as k
k.set_session(tfdbg.LocalCLIDebugWrapperSession(k.get_session()))
# Fit the model
_ = trainer.fit(train=train,
val=val,
train_im_per_epoch=train_images_per_epoch,
val_im_per_epoch=val_images_per_epoch,
hparams=hparams,
no_im=no_images,
**hparams["fit"])
# Save final model weights (usually not used, but maybe....?)
if not os.path.exists("%s/model" % base_path):
os.mkdir("%s/model" % base_path)
model_path = "%s/model/model_weights.h5" % base_path
logger("Saving current model to: %s" % model_path)
org_model.save_weights(model_path)
def entry_func(args=None):
# Get command line arguments
args = vars(get_argparser().parse_args(args))
base_dir = os.path.abspath(args["project_dir"])
_file = args["f"]
label = args["l"]
N_extra = args["extra"]
try:
N_extra = int(N_extra)
except ValueError:
pass
# Get settings from YAML file
from MultiPlanarUNet.train.hparams import YAMLHParams
hparams = YAMLHParams(os.path.join(base_dir, "train_hparams.yaml"))
# Set strides
hparams["fit"]["strides"] = args["strides"]
if not _file:
try:
# Data specified from command line?
data_dir = os.path.abspath(args["data_dir"])
# Set with default sub dirs
hparams["test_data"] = {
"base_dir": data_dir,
"img_subdir": "images",
"label_subdir": "labels"
}
except (AttributeError, TypeError):
data_dir = hparams["test_data"]["base_dir"]
else:
data_dir = False
out_dir = os.path.abspath(args["out_dir"])
overwrite = args["overwrite"]
predict_mode = args["no_eval"]
save_only_pred = args["save_only_pred"]
# Check if valid dir structures
validate_folders(base_dir, data_dir, out_dir, overwrite)
# Import all needed modules (folder is valid at this point)
import numpy as np
from MultiPlanarUNet.image import ImagePairLoader, ImagePair
from MultiPlanarUNet.utils import get_best_model, create_folders, \
pred_to_class, await_and_set_free_gpu, set_gpu
from MultiPlanarUNet.utils.fusion import predict_3D_patches, predict_3D_patches_binary, pred_3D_iso
from MultiPlanarUNet.logging import init_result_dict_3D, save_all_3D
from MultiPlanarUNet.evaluate import dice_all
from MultiPlanarUNet.bin.predict import save_nii_files
# Fetch GPU(s)
num_GPUs = args["num_GPUs"]
force_gpu = args["force_GPU"]
# Wait for free GPU
if force_gpu == -1:
await_and_set_free_gpu(N=num_GPUs, sleep_seconds=240)
else:
set_gpu(force_gpu)
# Read settings from the project hyperparameter file
dim = hparams["build"]["dim"]
n_classes = hparams["build"]["n_classes"]
mode = hparams["fit"]["intrp_style"]
# Set ImagePairLoader object
if not _file:
image_pair_loader = ImagePairLoader(predict_mode=predict_mode,
**hparams["test_data"])
else:
predict_mode = not bool(label)
image_pair_loader = ImagePairLoader(predict_mode=predict_mode,
initialize_empty=True)
image_pair_loader.add_image(ImagePair(_file, label))
# Put them into a dict and remove from image_pair_loader to gain more control with
# garbage collection
all_images = {image.id: image for image in image_pair_loader.images}
image_pair_loader.images = None
""" Define UNet model """
from MultiPlanarUNet.models import model_initializer
hparams["build"]["batch_size"] = 1
unet = model_initializer(hparams, False, base_dir)
model_path = get_best_model(base_dir + "/model")
unet.load_weights(model_path)
# Evaluate?
if not predict_mode:
# Prepare dictionary to store results in pd df
results, detailed_res = init_result_dict_3D(all_images, n_classes)
# Save to check correct format
save_all_3D(results, detailed_res, out_dir)
# Define result paths
nii_res_dir = os.path.join(out_dir, "nii_files")
create_folders(nii_res_dir)
image_ids = sorted(all_images)
for n_image, image_id in enumerate(image_ids):
print("\n[*] Running on: %s" % image_id)
# Set image_pair_loader object with only the given file
image = all_images[image_id]
image_pair_loader.images = [image]
seq = image_pair_loader.get_sequencer(n_classes=n_classes,
no_log=True,
**hparams["fit"])
if mode.lower() == "iso_live_3d":
pred = pred_3D_iso(model=unet,
sequence=seq,
image=image,
extra_boxes=N_extra,
min_coverage=None)
else:
# Predict on volume using model
if n_classes > 1:
pred = predict_3D_patches(model=unet,
patches=seq,
image=image,
N_extra=N_extra)
else:
pred = predict_3D_patches_binary(model=unet,
patches=seq,
image_id=image_id,
N_extra=N_extra)
if not predict_mode:
# Get patches for the current image
y = image.labels
# Calculate dice score
print("Mean dice: ", end="", flush=True)
p = pred_to_class(pred, img_dims=3, has_batch_dim=False)
dices = dice_all(y, p, n_classes=n_classes, ignore_zero=True)
mean_dice = dices[~np.isnan(dices)].mean()
print("Dices: ", dices)
print("%s (n=%i)" % (mean_dice, len(dices)))
# Add to results
results[image_id] = [mean_dice]
detailed_res[image_id] = dices
# Overwrite with so-far results
save_all_3D(results, detailed_res, out_dir)
# Save results
save_nii_files(p, image, nii_res_dir, save_only_pred)
# Remove image from dictionary and image_pair_loader to free memory
del all_images[image_id]
image_pair_loader.images.remove(image)
if not predict_mode:
# Write final results
save_all_3D(results, detailed_res, out_dir)