def save_nii_files(combined, image, nii_res_dir, save_input_files): from MultiPlanarUNet.utils import create_folders import nibabel as nib import os # Extract data if nii files try: combined = combined.get_data() except AttributeError: combined = nib.Nifti1Image(combined, affine=image.affine) volumes = [combined, image.image_obj, image.labels_obj] labels = [ "%s_PRED.nii.gz" % image.id, "%s_IMAGE.nii.gz" % image.id, "%s_LABELS.nii.gz" % image.id ] if not save_input_files: volumes = volumes[:1] labels = labels[:1] p = os.path.abspath(nii_res_dir) # Save file directly in nii_res_dir else: # Create sub-folder under nii_res_dir p = os.path.join(nii_res_dir, image.id) create_folders(p) # Save for nii, fname in zip(volumes, labels): try: nib.save(nii, "%s/%s" % (p, fname)) except AttributeError: # No labels file? pass
def run_sub_experiment(split_dir, out_dir, script, hparams, no_hparams, GPUs, GPU_queue, lock, logger): # Create sub-directory split = os.path.split(split_dir)[-1] out_dir = os.path.join(out_dir, split) out_hparams = os.path.join(out_dir, "train_hparams.yaml") create_folders(out_dir) # Get list of commands commands = parse_script(script, GPUs) # Move hparams and script files into folder if not no_hparams: copy_yaml_and_set_data_dirs(in_path=hparams, out_path=out_hparams, data_dir=split_dir) # Change directory and file permissions os.chdir(out_dir) # Log lock.acquire() s = "[*] Running experiment: %s" % split logger("\n%s\n%s" % ("-" * len(s), s)) logger("Data dir:", split_dir) logger("Out dir:", out_dir) logger("Using GPUs:", GPUs) logger("\nRunning commands:") for i, command in enumerate(commands): logger(" %i) %s" % (i + 1, " ".join(command))) logger("-" * len(s)) lock.release() # Run the commands run_next_command = True for command in commands: if not run_next_command: break lock.acquire() logger("[%s - STARTING] %s" % (split, " ".join(command))) lock.release() p = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE) _, err = p.communicate() rc = p.returncode lock.acquire() if rc != 0: logger("[%s - ERROR - Exit code %i] %s" % (split, rc, " ".join(command))) logger("\n----- START error message -----\n%s\n" "----- END error message -----\n" % err.decode("utf-8")) run_next_command = False else: logger("[%s - FINISHED] %s" % (split, " ".join(command))) lock.release() # Add the GPUs back into the queue GPU_queue.put(GPUs)
def __init__(self, base_path, print_to_screen=True, active_file=None, overwrite_existing=False, print_calling_method=True, no_sub_folder=False, log_prefix=""): self.base_path = os.path.abspath(base_path) if not no_sub_folder: self.path = os.path.join(self.base_path, "logs") else: self.path = self.base_path create_folders([self.path]) self.overwrite_existing = overwrite_existing self._enabled = True # Get built in print function # (if overwritten globally, Logger still maintains a reference to the # true print function) self.print_f = __builtins__["print"] # Print options self.separator = "-" * 80 self.print_to_screen = print_to_screen self.print_calling_method = print_calling_method # Set paths to log files self.log_files = {} self.currently_logging = {} self.prefix = "" if log_prefix is None else str(log_prefix) self.active_log_file = active_file or "log" # For using the logger from multiple threads self.lock = Lock()
def entry_func(args=None): parser = get_parser() args = vars(parser.parse_args(args)) # Get arguments project_folder = os.path.abspath(args["project_folder"]) copy_weights = args["copy_weights"] weights = os.path.abspath(args["weights_file"] or get_best_model(project_folder + "/model")) out_dir = os.path.join(project_folder, args["out_dir"]) # Get main hyperparamter file and check if correct modelt ype hparams = YAMLHParams(project_folder + "/train_hparams.yaml", no_log=True) tasks = hparams.get("tasks", False) if not tasks: print("[ERROR] Project must be of type 'MultiTask'.") sys.exit(0) # Branch out each sub-task create_folders(out_dir) for name, hparams_file in zip(tasks["task_names"], tasks["hparam_files"]): print("\n[*] Branching task %s" % name) # Reload the hparams in each iteration as we overwrite fields each time hparams = YAMLHParams(project_folder + "/train_hparams.yaml", no_log=False) # Get task specific parameters task_hparams = YAMLHParams(project_folder + "/%s" % hparams_file) branch(task_name=name, out_dir=os.path.join(out_dir, name), task_hparams=task_hparams, task_hparams_file=hparams_file, shared_hparams=hparams, weights=weights, copy_weights=copy_weights, views_file=os.path.join(project_folder, "views.npz"))
def branch(task_name, out_dir, task_hparams, task_hparams_file, shared_hparams, weights, copy_weights, views_file): create_folders(out_dir) # Set the task fields to only the current task under 'build' shared_hparams.set_value("build", "task_names", value='["%s"]' % task_name, overwrite=True) shared_hparams.delete_group("tasks") # No longer necessary # Create a map defining where each task specific parameter should go mapping = { "task_specifics/n_classes": "build/n_classes", "task_specifics/n_channels": "build/n_channels", "task_specifics/dim": ["build/dim", "fit/dim"], "task_specifics/out_activation": "build/out_activation", "task_specifics/real_space_span": "fit/real_space_span" } for soruce, targets in mapping.items(): in_key1, in_key2 = soruce.split("/") value = task_hparams[in_key1][in_key2] # Set the value at all targets targets = [targets] if isinstance(targets, str) else targets for target in targets: out_key1, out_key2 = target.split("/") shared_hparams.set_value(out_key1, out_key2, value, overwrite=True) # Add all data folders data_folders = ("train_data", "val_data", "test_data", "aug_data") for df in data_folders: yaml = task_hparams.get_group(df) shared_hparams.add_group(yaml_string=yaml) # Save the updates parameters to a new location shared_hparams.save_current(os.path.join(out_dir, "train_hparams.yaml")) # Add weights to folder weights_folder = os.path.join(out_dir, "model") create_folders(weights_folder) out_weights_name = os.path.split(weights)[1] out_weights_path = os.path.join(weights_folder, out_weights_name) func = copy if copy_weights else os.symlink if copy_weights: print("Copying weights...") else: print("Symlinking weights...") func(weights, out_weights_path) # Add views (check for existence for future compatibility with 3D models) if os.path.exists(views_file): func(views_file, os.path.join(out_dir, "views.npz"))
def save_all_3D(results, detailed_res, out_dir): txt_res_dir = os.path.join(out_dir, "txt") csv_res_dir = os.path.join(out_dir, "csv") # Create folders create_folders([txt_res_dir, csv_res_dir]) # Save main text/csv results files results_to_txt(results, txt_res_dir, transpose=True) results_to_csv(results, csv_res_dir, transpose=True) results_to_txt(detailed_res, txt_res_dir, fname="detailed") results_to_csv(detailed_res, csv_res_dir, fname="detailed")
def run_on_split(split_path, test_split, train_val_data, n_val, args): """ Add the train/val/test files of a single split to the split directories Depending on the arguments parsed (--copy, --file_list etc.) either copies, symlinks or creates a LIST_OF_FILES.txt file of absolute paths in each split sub-directory. Args: split_path: (string) Path to the split directory test_split: (list) List of paths pointing to split test data train_val_data: (list) List of paths pointing to the remaining data n_val: (int) Number of samples in 'train_val_data' to use for validation - rest is used for training. args: (tuple) Parsed arguments, see argparser. """ # Define train, val and test sub-dirs train_path = os.path.join(split_path, "train") val_path = os.path.join(split_path, "val") if n_val else None test_path = os.path.join(split_path, "test") # Create folders if not existing create_folders([train_path, val_path, test_path]) # Copy or symlink? if args.copy: from shutil import copyfile move_func = copyfile elif args.file_list: move_func = _add_to_file_list_fallback else: move_func = os.symlink # Extract validation data from the remaining random.shuffle(train_val_data) validation = train_val_data[:n_val] training = train_val_data[n_val:] # Add training add_files(training, train_path, move_func) # Add test data add_files(test_split, test_path, move_func) if n_val: # Add validation add_files(validation, val_path, move_func)
def save_all(results, pc_results, out_dir): # Get output paths txt_res_dir = os.path.join(out_dir, "txt") csv_res_dir = os.path.join(out_dir, "csv") # Create folders create_folders([txt_res_dir, csv_res_dir]) # Save main text/csv results files results_to_txt(results, txt_res_dir) results_to_csv(results, csv_res_dir) # Write detailed results for view in pc_results: r = pc_results[view] view_str = str(view).replace("[", "").strip().replace("]", "").replace(" ", "_") results_to_txt(r, txt_res_dir, fname=view_str) results_to_csv(r, csv_res_dir, fname=view_str)
def entry_func(args=None): # Get parser parser = vars(get_parser().parse_args(args)) # Get parser arguments cv_dir = os.path.abspath(parser["CV_dir"]) out_dir = os.path.abspath(parser["out_dir"]) create_folders(out_dir) await_PID = parser["wait_for"] run_split = parser["run_on_split"] start_from = parser["start_from"] or 0 num_jobs = parser["num_jobs"] or 1 # GPU settings num_GPUs = parser["num_GPUs"] force_GPU = parser["force_GPU"] ignore_GPU = parser["ignore_GPU"] monitor_GPUs_every = parser["monitor_GPUs_every"] # User input assertions _assert_force_and_ignore_gpus(force_GPU, ignore_GPU) if run_split: _assert_run_split(start_from, monitor_GPUs_every, num_jobs) # Wait for PID? if await_PID: from MultiPlanarUNet.utils import await_PIDs await_PIDs(await_PID) # Get file paths script = os.path.abspath(parser["script_prototype"]) hparams = os.path.abspath(parser["hparams_prototype"]) no_hparams = parser["no_hparams"] # Get list of folders of CV data to run on cv_folders = get_CV_folders(cv_dir) if run_split is not None: if run_split < 0 or run_split >= len(cv_folders): raise ValueError("--run_on_split should be in range [0-{}], " "got {}".format(len(cv_folders) - 1, run_split)) cv_folders = [cv_folders[run_split]] log_appendix = "_split{}".format(run_split) else: log_appendix = "" # Get a logger object logger = Logger(base_path="./", active_file="output" + log_appendix, print_calling_method=False, overwrite_existing=True) if force_GPU: # Only these GPUs fill be chosen from from MultiPlanarUNet.utils import set_gpu set_gpu(force_GPU) if num_GPUs: # Get GPU sets (up to the number of splits) gpu_sets = get_free_GPU_sets(num_GPUs, ignore_GPU)[:len(cv_folders)] elif not num_jobs or num_jobs < 0: raise ValueError("Should specify a number of jobs to run in parallel " "with the --num_jobs flag when using 0 GPUs pr. " "process (--num_GPUs=0 was set).") else: gpu_sets = ["''"] * parser["num_jobs"] # Get process pool, lock and GPU queue objects lock = Lock() gpu_queue = Queue() for gpu in gpu_sets: gpu_queue.put(gpu) procs = [] if monitor_GPUs_every is not None and monitor_GPUs_every: logger("\nOBS: Monitoring GPU pool every %i seconds\n" % monitor_GPUs_every) # Start a process monitoring new GPU availability over time stop_event = Event() t = Process(target=monitor_GPUs, args=(monitor_GPUs_every, gpu_queue, num_GPUs, ignore_GPU, gpu_sets, stop_event)) t.start() procs.append(t) else: stop_event = None try: for cv_folder in cv_folders[start_from:]: gpus = gpu_queue.get() t = Process(target=run_sub_experiment, args=(cv_folder, out_dir, script, hparams, no_hparams, gpus, gpu_queue, lock, logger)) t.start() procs.append(t) for t in procs: if not t.is_alive(): t.join() except KeyboardInterrupt: for t in procs: t.terminate() if stop_event is not None: stop_event.set() for t in procs: t.join()
def entry_func(args=None): # Get parser parser = vars(get_parser().parse_args(args)) # Get arguments data_dir = os.path.abspath(parser["data_dir"]) n_splits = int(parser["CV"]) if n_splits > 1: out_dir = os.path.join(data_dir, parser["out_dir"], "%i_CV" % n_splits) else: out_dir = os.path.join(data_dir, parser["out_dir"], "fixed_split") im_dir = os.path.join(data_dir, parser["im_sub_dir"]) lab_dir = os.path.join(data_dir, parser["lab_sub_dir"]) if parser["aug_sub_dir"]: aug_dir = os.path.join(data_dir, parser["aug_sub_dir"]) else: aug_dir = None copy = parser["copy"] file_list = parser["file_list"] regex = parser["file_regex"] val_frac = parser["validation_fraction"] test_frac = parser["test_fraction"] if n_splits == 1 and not test_frac: raise ValueError("Must specify --test_fraction with --CV=1.") if copy and file_list: raise ValueError("Only one of --copy and --file_list " "flags must be set.") # Assert suitable folders assert_dir_structure(data_dir, im_dir, lab_dir, out_dir) # Create sub-folders create_view_folders(out_dir, n_splits) # Get images images = glob(os.path.join(im_dir, regex)) if aug_dir: aug_im_dir = os.path.join(aug_dir, parser["im_sub_dir"]) aug_lab_dir = os.path.join(aug_dir, parser["lab_sub_dir"]) aug_images = glob(os.path.join(aug_im_dir, regex)) aug_labels = glob(os.path.join(aug_lab_dir, regex)) assert (len(aug_images) == len(aug_labels)) # Get validation size N_total = len(images) if n_splits > 1: N_test = N_total // n_splits else: N_test = int(N_total * test_frac) N_val = int(N_total * val_frac) if N_val + N_test >= N_total: raise ValueError("Too large validation_fraction - " "No training samples left!") N_train = N_total - N_test - N_val print("-----") print("Total images:".ljust(40), N_total) print("Train images pr. split:".ljust(40), N_train) print("Validation images pr. split:".ljust(40), N_val) print("Test images pr. split:".ljust(40), N_test) # Shuffle and split the images into CV parts random.shuffle(images) splits = np.array_split(images, n_splits) # Symlink / copy files for i, split in enumerate(splits): print(" Split %i/%i" % (i + 1, n_splits), end="\r", flush=True) # Set root path to split folder if n_splits > 1: split_path = os.path.join(out_dir, "split_%i" % i) else: split_path = out_dir # Here we kind of hacky force the following code to work with CV=1 # Define a test set and overwrite the current split (which stores # add the data, as splits was never split with n_splits=1 split = splits[0][:N_test] # Overwrite the splits variable to a length 2 array with the # remaining data which will be used as val+train. The loop still # refers to the old split and thus will only execute once splits = [split, splits[0][N_test:]] # Define train, val and test sub-dirs train_path = os.path.join(split_path, "train") train_im_path = os.path.join(train_path, parser["im_sub_dir"]) train_label_path = os.path.join(train_path, parser["lab_sub_dir"]) val_path = os.path.join(split_path, "val") val_im_path = os.path.join(val_path, parser["im_sub_dir"]) val_label_path = os.path.join(val_path, parser["lab_sub_dir"]) test_path = os.path.join(split_path, "test") test_im_path = os.path.join(test_path, parser["im_sub_dir"]) test_label_path = os.path.join(test_path, parser["lab_sub_dir"]) if aug_dir: aug_path = os.path.join(split_path, "aug") aug_im_path = os.path.join(aug_path, parser["im_sub_dir"]) aug_label_path = os.path.join(aug_path, parser["lab_sub_dir"]) else: aug_path, aug_im_path, aug_label_path = None, None, None # Create folders if not existing create_folders([ train_path, val_path, train_im_path, train_label_path, val_im_path, val_label_path, test_path, test_im_path, test_label_path, aug_path, aug_im_path, aug_label_path ]) # Copy or symlink? if copy: from shutil import copyfile move_func = copyfile elif file_list: move_func = _add_to_file_list_fallback else: move_func = os.symlink # Add test data to test folder add_images(split, test_im_path, test_label_path, im_dir, lab_dir, move_func) # Join remaining splits into train+val remaining = [x for ind, x in enumerate(splits) if ind != i] remaining = [item for sublist in remaining for item in sublist] # Extract validation data from the remaining random.shuffle(remaining) validation = remaining[:N_val] training = remaining[N_val:] # Add add_images(validation, val_im_path, val_label_path, im_dir, lab_dir, move_func) add_images(training, train_im_path, train_label_path, im_dir, lab_dir, move_func) # Add augmented images to training dir? if aug_dir: # Get list of train file names train_fnames = [os.path.split(f)[-1] for f in training] augmented = [] for aug_im in aug_images: aug_im_fname = os.path.split(aug_im)[-1] if any([t_fn in aug_im_fname for t_fn in train_fnames]): augmented.append(aug_im) # Add the images add_images(augmented, aug_im_path, aug_label_path, aug_im_dir, aug_lab_dir, move_func)
def entry_func(args=None): # Project base path args = vars(get_argparser().parse_args(args)) basedir = os.path.abspath(args["project_dir"]) overwrite = args["overwrite"] continue_training = args["continue_training"] eval_prob = args["eval_prob"] await_PID = args["wait_for"] dice_weight = args["dice_weight"] print("Fitting fusion model for project-folder: %s" % basedir) # Minimum images in validation set before also using training images min_val_images = 15 # Fusion model training params epochs = args['epochs'] fm_batch_size = args["batch_size"] # Early stopping params early_stopping = args["early_stopping"] # Wait for PID? if await_PID: from MultiPlanarUNet.utils import await_PIDs await_PIDs(await_PID) # Fetch GPU(s) num_GPUs = args["num_GPUs"] force_gpu = args["force_GPU"] # Wait for free GPU if not force_gpu: await_and_set_free_gpu(N=num_GPUs, sleep_seconds=120) num_GPUs = 1 else: set_gpu(force_gpu) num_GPUs = len(force_gpu.split(",")) # Get logger logger = Logger(base_path=basedir, active_file="train_fusion", overwrite_existing=overwrite) # Get YAML hyperparameters hparams = YAMLHParams(os.path.join(basedir, "train_hparams.yaml")) # Get some key settings n_classes = hparams["build"]["n_classes"] if hparams["build"]["out_activation"] == "linear": # Trained with logit targets? hparams["build"][ "out_activation"] = "softmax" if n_classes > 1 else "sigmoid" # Get views views = np.load("%s/views.npz" % basedir)["arr_0"] del hparams["fit"]["views"] # Get weights and set fusion (output) path weights = get_best_model("%s/model" % basedir) weights_name = os.path.splitext(os.path.split(weights)[-1])[0] fusion_weights = "%s/model/fusion_weights/" \ "%s_fusion_weights.h5" % (basedir, weights_name) create_folders(os.path.split(fusion_weights)[0]) # Log a few things log(logger, hparams, views, weights, fusion_weights) # Check if exists already... if not overwrite and os.path.exists(fusion_weights): from sys import exit print("\n[*] A fusion weights file already exists at '%s'." "\n Use the --overwrite flag to overwrite." % fusion_weights) exit(0) # Load validation data images = ImagePairLoader(**hparams["val_data"], logger=logger) is_validation = {m.id: True for m in images} # Define random sets of images to train on simul. (cant be all due # to memory constraints) image_IDs = [m.id for m in images] if len(images) < min_val_images: # Pick N random training images diff = min_val_images - len(images) logger("Adding %i training images to set" % diff) # Load the training data and pick diff images train = ImagePairLoader(**hparams["train_data"], logger=logger) indx = np.random.choice(np.arange(len(train)), diff, replace=diff > len(train)) # Add the images to the image set set train_add = [train[i] for i in indx] for m in train_add: is_validation[m.id] = False image_IDs.append(m.id) images.add_images(train_add) # Append to length % sub_size == 0 sub_size = args["images_per_round"] rest = int(sub_size * np.ceil(len(image_IDs) / sub_size)) - len(image_IDs) if rest: image_IDs += list(np.random.choice(image_IDs, rest, replace=False)) # Shuffle and split random.shuffle(image_IDs) sets = [ set(s) for s in np.array_split(image_IDs, len(image_IDs) / sub_size) ] assert (contains_all_images(sets, image_IDs)) # Define fusion model (named 'org' to store reference to orgiginal model if # multi gpu model is created below) fusion_model_org = FusionModel(n_inputs=len(views), n_classes=n_classes, weight=dice_weight, logger=logger, verbose=False) if continue_training: fusion_model_org.load_weights(fusion_weights) print("\n[OBS] CONTINUED TRAINING FROM:\n", fusion_weights) # Define model unet = init_model(hparams["build"], logger) print("\n[*] Loading weights: %s\n" % weights) unet.load_weights(weights, by_name=True) if num_GPUs > 1: from tensorflow.keras.utils import multi_gpu_model # Set for predictor model n_classes = n_classes unet = multi_gpu_model(unet, gpus=num_GPUs) unet.n_classes = n_classes # Set for fusion model fusion_model = multi_gpu_model(fusion_model_org, gpus=num_GPUs) else: fusion_model = fusion_model_org # Compile the model logger("Compiling...") metrics = [ "sparse_categorical_accuracy", sparse_fg_precision, sparse_fg_recall ] fusion_model.compile(optimizer=Adam(lr=1e-3), loss=fusion_model_org.loss, metrics=metrics) fusion_model_org._log() try: _run_fusion_training(sets, logger, hparams, min_val_images, is_validation, views, n_classes, unet, fusion_model_org, fusion_model, early_stopping, fm_batch_size, epochs, eval_prob, fusion_weights) except KeyboardInterrupt: pass finally: if not os.path.exists(os.path.split(fusion_weights)[0]): os.mkdir(os.path.split(fusion_weights)[0]) # Save fusion model weights # OBS: Must be original model if multi-gpu is performed! fusion_model_org.save_weights(fusion_weights)
def entry_func(args=None): # Get parser parser = vars(get_parser().parse_args(args)) # Get parser arguments cv_dir = os.path.abspath(parser["CV_dir"]) out_dir = os.path.abspath(parser["out_dir"]) create_folders(out_dir) start_from = parser["start_from"] await_PID = parser["wait_for"] monitor_GPUs_every = parser["monitor_GPUs_every"] # Get a logger object logger = Logger(base_path="./", active_file="output", print_calling_method=False, overwrite_existing=True) # Wait for PID? if await_PID: from MultiPlanarUNet.utils import await_PIDs await_PIDs(await_PID) # Get number of GPUs per process num_GPUs = parser["num_GPUs"] # Get file paths script = os.path.abspath(parser["script_prototype"]) hparams = os.path.abspath(parser["hparams_prototype"]) # Get list of folders of CV data to run on cv_folders = get_CV_folders(cv_dir) # Get GPU sets gpu_sets = get_free_GPU_sets(num_GPUs) # Get process pool, lock and GPU queue objects lock = Lock() gpu_queue = Queue() for gpu in gpu_sets: gpu_queue.put(gpu) procs = [] if monitor_GPUs_every is not None and monitor_GPUs_every: logger("\nOBS: Monitoring GPU pool every %i seconds\n" % monitor_GPUs_every) # Start a process monitoring new GPU availability over time stop_event = Event() t = Process(target=monitor_GPUs, args=(monitor_GPUs_every, gpu_queue, num_GPUs, gpu_sets, stop_event)) t.start() procs.append(t) else: stop_event = None try: for cv_folder in cv_folders[start_from:]: gpus = gpu_queue.get() t = Process(target=run_sub_experiment, args=(cv_folder, out_dir, script, hparams, gpus, gpu_queue, lock, logger)) t.start() procs.append(t) for t in procs: if not t.is_alive(): t.join() except KeyboardInterrupt: for t in procs: t.terminate() if stop_event is not None: stop_event.set() for t in procs: t.join()
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)
def entry_func(args=None): # Get command line arguments args = vars(get_argparser().parse_args(args)) base_dir = os.path.abspath(args["project_dir"]) analytical = args["analytical"] majority = args["majority"] _file = args["f"] label = args["l"] await_PID = args["wait_for"] eval_prob = args["eval_prob"] _continue = args["continue"] if analytical and majority: raise ValueError("Cannot specify both --analytical and --majority.") # Get settings from YAML file from MultiPlanarUNet.train.hparams import YAMLHParams hparams = YAMLHParams(os.path.join(base_dir, "train_hparams.yaml")) 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_input_files = args["save_input_files"] no_argmax = args["no_argmax"] on_val = args["on_val"] # Check if valid dir structures validate_folders(base_dir, out_dir, overwrite, _continue) # Import all needed modules (folder is valid at this point) import numpy as np from MultiPlanarUNet.image import ImagePairLoader, ImagePair from MultiPlanarUNet.models import FusionModel from MultiPlanarUNet.models.model_init import init_model from MultiPlanarUNet.utils import await_and_set_free_gpu, get_best_model, \ create_folders, pred_to_class, set_gpu from MultiPlanarUNet.logging import init_result_dicts, save_all, load_result_dicts from MultiPlanarUNet.evaluate import dice_all from MultiPlanarUNet.utils.fusion import predict_volume, map_real_space_pred from MultiPlanarUNet.interpolation.sample_grid import get_voxel_grid_real_space # Wait for PID? if await_PID: from MultiPlanarUNet.utils import await_PIDs await_PIDs(await_PID) # Set GPU device # Fetch GPU(s) num_GPUs = args["num_GPUs"] force_gpu = args["force_GPU"] # Wait for free GPU if not force_gpu: await_and_set_free_gpu(N=num_GPUs, sleep_seconds=120) num_GPUs = 1 else: set_gpu(force_gpu) num_GPUs = len(force_gpu.split(",")) # Read settings from the project hyperparameter file n_classes = hparams["build"]["n_classes"] # Get views views = np.load("%s/views.npz" % base_dir)["arr_0"] # Force settings hparams["fit"]["max_background"] = 1 hparams["fit"]["test_mode"] = True hparams["fit"]["mix_planes"] = False hparams["fit"]["live_intrp"] = False if "use_bounds" in hparams["fit"]: del hparams["fit"]["use_bounds"] del hparams["fit"]["views"] if hparams["build"]["out_activation"] == "linear": # Trained with logit targets? hparams["build"][ "out_activation"] = "softmax" if n_classes > 1 else "sigmoid" # Set ImagePairLoader object if not _file: data = "test_data" if not on_val else "val_data" image_pair_loader = ImagePairLoader(predict_mode=predict_mode, **hparams[data]) else: predict_mode = not bool(label) image_pair_loader = ImagePairLoader(predict_mode=predict_mode, single_file_mode=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 if _continue: all_images = remove_already_predicted(all_images, out_dir) # Evaluate? if not predict_mode: if _continue: csv_dir = os.path.join(out_dir, "csv") results, detailed_res = load_result_dicts(csv_dir=csv_dir, views=views) else: # Prepare dictionary to store results in pd df results, detailed_res = init_result_dicts(views, all_images, n_classes) # Save to check correct format save_all(results, detailed_res, out_dir) # Define result paths nii_res_dir = os.path.join(out_dir, "nii_files") create_folders(nii_res_dir) """ Define UNet model """ model_path = get_best_model(base_dir + "/model") unet = init_model(hparams["build"]) unet.load_weights(model_path, by_name=True) if num_GPUs > 1: from tensorflow.keras.utils import multi_gpu_model n_classes = unet.n_classes unet = multi_gpu_model(unet, gpus=num_GPUs) unet.n_classes = n_classes weights_name = os.path.splitext(os.path.split(model_path)[1])[0] if not analytical and not majority: # Get Fusion model fm = FusionModel(n_inputs=len(views), n_classes=n_classes) weights = base_dir + "/model/fusion_weights/%s_fusion_weights.h5" % weights_name print("\n[*] Loading weights:\n", weights) # Load fusion weights fm.load_weights(weights) print("\nLoaded weights:\n\n%s\n%s\n---" % tuple(fm.layers[-1].get_weights())) # Multi-gpu? if num_GPUs > 1: print("Using multi-GPU model (%i GPUs)" % num_GPUs) fm = multi_gpu_model(fm, gpus=num_GPUs) """ Finally predict on the images """ image_ids = sorted(all_images) N_images = len(image_ids) for n_image, image_id in enumerate(image_ids): print("\n[*] (%i/%s) Running on: %s" % (n_image + 1, N_images, image_id)) # Set image_pair_loader object with only the given file image = all_images[image_id] image_pair_loader.images = [image] # Load views kwargs = hparams["fit"] kwargs.update(hparams["build"]) seq = image_pair_loader.get_sequencer(views=views, **kwargs) # Get voxel grid in real space voxel_grid_real_space = get_voxel_grid_real_space(image) # Prepare tensor to store combined prediction d = image.image.shape[:-1] if not majority: combined = np.empty(shape=(len(views), d[0], d[1], d[2], n_classes), dtype=np.float32) else: combined = np.empty(shape=(d[0], d[1], d[2], n_classes), dtype=np.float32) print("Predicting on brain hyper-volume of shape:", combined.shape) # Predict for each view for n_view, v in enumerate(views): print("\n[*] (%i/%i) View: %s" % (n_view + 1, len(views), v)) # for each view, predict on all voxels and map the predictions # back into the original coordinate system # Sample planes from the image at grid_real_space grid # in real space (scanner RAS) coordinates. X, y, grid, inv_basis = seq.get_view_from(image.id, v, n_planes="same+20") # Predict on volume using model pred = predict_volume(unet, X, axis=2, batch_size=seq.batch_size) # Map the real space coordiante predictions to nearest # real space coordinates defined on voxel grid mapped_pred = map_real_space_pred(pred, grid, inv_basis, voxel_grid_real_space, method="nearest") if not majority: combined[n_view] = mapped_pred else: combined += mapped_pred if n_classes == 1: # Set to background if outside pred domain combined[n_view][np.isnan(combined[n_view])] = 0. if not predict_mode and np.random.rand() <= eval_prob: view_dices = dice_all(y, pred_to_class(pred, img_dims=3, has_batch_dim=False), ignore_zero=False, n_classes=n_classes, skip_if_no_y=False) mapped_dices = dice_all(image.labels, pred_to_class(mapped_pred, img_dims=3, has_batch_dim=False), ignore_zero=False, n_classes=n_classes, skip_if_no_y=False) mean_dice = mapped_dices[~np.isnan(mapped_dices)][1:].mean() # Print dice scores print("View dice scores: ", view_dices) print("Mapped dice scores: ", mapped_dices) print("Mean dice (n=%i): " % (len(mapped_dices) - 1), mean_dice) # Add to results results.loc[image_id, str(v)] = mean_dice detailed_res[str(v)][image_id] = mapped_dices[1:] # Overwrite with so-far results save_all(results, detailed_res, out_dir) else: print("Skipping evaluation for this view... " "(eval_prob=%.3f, predict_mode=%s)" % (eval_prob, predict_mode)) if not analytical and not majority: # Combine predictions across views using Fusion model print("\nFusing views...") combined = np.moveaxis(combined, 0, -2).reshape( (-1, len(views), n_classes)) combined = fm.predict(combined, batch_size=10**4, verbose=1).reshape( (d[0], d[1], d[2], n_classes)) elif analytical: print("\nFusing views (analytical)...") combined = np.sum(combined, axis=0) if not no_argmax: print("\nComputing majority vote...") combined = pred_to_class(combined.squeeze(), img_dims=3).astype(np.uint8) if not predict_mode: if no_argmax: # MAP only for dice calculation c_temp = pred_to_class(combined, img_dims=3).astype(np.uint8) else: c_temp = combined # Calculate combined prediction dice dices = dice_all(image.labels, c_temp, n_classes=n_classes, ignore_zero=True, skip_if_no_y=False) mean_dice = dices[~np.isnan(dices)].mean() detailed_res["MJ"][image_id] = dices print("Combined dices: ", dices) print("Combined mean dice: ", mean_dice) results.loc[image_id, "MJ"] = mean_dice # Overwrite with so-far results save_all(results, detailed_res, out_dir) # Save combined prediction volume as .nii file print("Saving .nii files...") save_nii_files(combined, image, nii_res_dir, save_input_files) # 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(results, detailed_res, out_dir)