def get_image_pair_loader(args, hparams, out_dir): """ TODO """ from mpunet.image import ImagePairLoader, ImagePair if not args.f: # No single file was specified with -f flag, load the desired dataset dataset = args.dataset.replace("_data", "") + "_data" image_pair_loader = ImagePairLoader(predict_mode=args.no_eval, **hparams[dataset]) else: predict_mode = not bool(args.l) image_pair_loader = ImagePairLoader(predict_mode=predict_mode, initialize_empty=True) image_pair_loader.add_image(ImagePair(args.f, args.l)) # Put image pairs into a dict and remove from image_pair_loader to gain # more control with garbage collection image_pair_dict = { image.identifier: image for image in image_pair_loader.images } if vars(args)["continue"]: # Remove images that were already predicted image_pair_dict = remove_already_predicted(image_pair_dict, out_dir) return image_pair_loader, image_pair_dict
def predict_single(image, model, hparams, verbose=1): """ A generic prediction function that sets up a ImagePairLoader object for the given image, prepares the image and predicts. Note that this function should only be used for convinience in scripts that work on single images at a time anyway, as batch-preparing the entire ImagePairLoader object prior to prediction is faster. NOTE: Only works with iso_live intrp modes at this time """ mode = hparams["fit"]["intrp_style"].lower() assert mode in ("iso_live", "iso_live_3d") # Prepare image for prediction kwargs = hparams["fit"] kwargs.update(hparams["build"]) # Set verbose memory verb_mem = kwargs["verbose"] kwargs["verbose"] = verbose # Create a ImagePairLoader with only the given file from mpunet.image import ImagePairLoader image_pair_loader = ImagePairLoader(predict_mode=True, initialize_empty=True, no_log=bool(verbose)) image_pair_loader.add_image(image) # Get N classes n_classes = kwargs["n_classes"] if mode == "iso_live": # Add views if SMMV model kwargs["views"] = np.load(hparams.project_path + "/views.npz")["arr_0"] # Get sequence object sequence = image_pair_loader.get_sequencer(**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 predicted = np.empty(shape=(len(kwargs["views"]), d[0], d[1], d[2], n_classes), dtype=np.float32) print("Predicting on brain hyper-volume of shape:", predicted.shape) for n_view, v in enumerate(kwargs["views"]): print("\nView %i/%i: %s" % (n_view + 1, len(kwargs["views"]), v)) # Sample the volume along the view X, y, grid, inv_basis = sequence.get_view_from(image.id, v, n_planes="same+20") # Predict on volume using model pred = predict_volume(model, X, axis=2) # Map the real space coordiante predictions to nearest # real space coordinates defined on voxel grid predicted[n_view] = map_real_space_pred(pred, grid, inv_basis, voxel_grid_real_space, method="nearest") else: predicted = pred_3D_iso( model=model, sequence=image_pair_loader.get_sequencer(**kwargs), image=image, extra_boxes="3x", min_coverage=None) # Revert verbose mem kwargs["verbose"] = verb_mem return predicted
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 mpunet.hyperparameters 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 mpunet.image import ImagePairLoader, ImagePair from mpunet.utils import get_best_model, create_folders, \ pred_to_class, await_and_set_free_gpu, set_gpu from mpunet.utils.fusion import predict_3D_patches, predict_3D_patches_binary, pred_3D_iso from mpunet.logging import init_result_dict_3D, save_all_3D from mpunet.evaluate import dice_all from mpunet.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)) all_images = { image.identifier: image for image in image_pair_loader.images } # Set scaler and bg values image_pair_loader.set_scaler_and_bg_values( bg_value=hparams.get_from_anywhere('bg_value'), scaler=hparams.get_from_anywhere('scaler'), compute_now=False) # Init LazyQueue and get its sequencer from mpunet.sequences.utils import get_sequence seq = get_sequence(data_queue=image_pair_loader, is_validation=True, **hparams["fit"], **hparams["build"]) """ Define UNet model """ from mpunet.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) with seq.image_pair_queue.get_image_by_id(image_id) as image_pair: if mode.lower() == "iso_live_3d": pred = pred_3D_iso(model=unet, sequence=seq, image=image_pair, 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_pair, 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_pair.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_pair, nii_res_dir, save_only_pred) if not predict_mode: # Write final results save_all_3D(results, detailed_res, out_dir)
def predict_single(image, model, hparams, verbose=1): """ A generic prediction function that sets up a ImagePairLoader object for the given image, prepares the image and predicts. Note that this function should only be used for convinience in scripts that work on single images at a time anyway, as batch-preparing the entire ImagePairLoader object prior to prediction is faster. NOTE: Only works with iso_live intrp modes at this time """ mode = hparams["fit"]["intrp_style"].lower() assert mode in ("iso_live", "iso_live_3d") # Create a ImagePairLoader with only the given file from mpunet.image import ImagePairLoader image_pair_loader = ImagePairLoader(predict_mode=True, initialize_empty=True, no_log=bool(verbose)) image_pair_loader.add_image(image) # Set scaler and bg values image_pair_loader.set_scaler_and_bg_values( bg_value=hparams.get_from_anywhere('bg_value'), scaler=hparams.get_from_anywhere('scaler'), compute_now=False) if mode == "iso_live": # Init LazyQueue and get its sequencer seq = get_sequence(data_queue=image_pair_loader, views=np.load(hparams.project_path + "/views.npz")["arr_0"], is_validation=True, **hparams["fit"], **hparams["build"]) with seq.image_pair_queue.get_image_by_id(image.identifier) as image: # 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 predicted = np.empty(shape=(len(seq.views), d[0], d[1], d[2], seq.n_classes), dtype=np.float32) print("Predicting on brain hyper-volume of shape:", predicted.shape) for n_view, v in enumerate(seq.views): print("\nView %i/%i: %s" % (n_view + 1, len(seq.views), v)) # Sample the volume along the view X, y, grid, inv_basis = seq.get_view_from(image, v, n_planes="same+20") # Predict on volume using model pred = predict_volume(model, X, axis=2) # Map the real space coordiante predictions to nearest # real space coordinates defined on voxel grid predicted[n_view] = map_real_space_pred(pred, grid, inv_basis, voxel_grid_real_space, method="nearest") else: # Init LazyQueue and get its sequencer seq = get_sequence(data_queue=image_pair_loader, is_validation=True, **hparams["fit"], **hparams["build"]) predicted = pred_3D_iso(model=model, sequence=seq, image=image, extra_boxes="3x", min_coverage=None) return predicted