| import os |
| import shutil |
| from nnunetv2.analysis.revert_normalisation import get_ct_normalisation_values, revert_normalisation |
| from nnunetv2.utilities.file_path_utilities import get_output_folder |
|
|
| def nnsyn_predict_entry(): |
| import argparse |
| parser = argparse.ArgumentParser(description="nnUNet prediction for nnSyn") |
| parser = argparse.ArgumentParser(description='Use this to run inference with nnU-Net. This function is used when ' |
| 'you want to manually specify a folder containing a trained nnU-Net ' |
| 'model. This is useful when the nnunet environment variables ' |
| '(nnUNet_results) are not set.') |
| parser.add_argument('-i', type=str, required=True, |
| help='input folder. Remember to use the correct channel numberings for your files (_0000 etc). ' |
| 'File endings must be the same as the training dataset!') |
| parser.add_argument('-o', type=str, required=True, |
| help='Output folder. If it does not exist it will be created. Predicted segmentations will ' |
| 'have the same name as their source images.') |
| parser.add_argument('-m', type=str, required=False, default=None, |
| help='input mask folder. If it does not exist it will be created. Predicted segmentations will ' |
| 'have the same name as their source images.') |
| parser.add_argument('-d', type=str, required=True, |
| help='Dataset with which you would like to predict. You can specify either dataset name or id') |
| parser.add_argument('-p', type=str, required=False, default='nnUNetPlans', |
| help='Plans identifier. Specify the plans in which the desired configuration is located. ' |
| 'Default: nnUNetPlans') |
| parser.add_argument('-tr', type=str, required=False, default='nnUNetTrainer', |
| help='What nnU-Net trainer class was used for training? Default: nnUNetTrainer') |
| parser.add_argument('-c', type=str, required=True, |
| help='nnU-Net configuration that should be used for prediction. Config must be located ' |
| 'in the plans specified with -p') |
| parser.add_argument('-f', nargs='+', type=str, required=False, default=(0, 1, 2, 3, 4), |
| help='Specify the folds of the trained model that should be used for prediction. ' |
| 'Default: (0, 1, 2, 3, 4)') |
| parser.add_argument('--revert_norm', action='store_true', required=False, default=False, |
| help='Set this flag if you want to undo the z-score normalization that is applied to ' |
| 'non-CT data. This will not give you the original intensities but will undo the ' |
| 'z-score normalization (mean and std are stored in the plans file). Note that this ' |
| 'only makes sense if you trained your model with revert_norm=True!') |
| parser.add_argument('-step_size', type=float, required=False, default=0.5, |
| help='Step size for sliding window prediction. The larger it is the faster but less accurate ' |
| 'the prediction. Default: 0.5. Cannot be larger than 1. We recommend the default.') |
| parser.add_argument('--disable_tta', action='store_true', required=False, default=False, |
| help='Set this flag to disable test time data augmentation in the form of mirroring. Faster, ' |
| 'but less accurate inference. Not recommended.') |
| parser.add_argument('--verbose', action='store_true', help="Set this if you like being talked to. You will have " |
| "to be a good listener/reader.") |
| parser.add_argument('--save_probabilities', action='store_true', |
| help='Set this to export predicted class "probabilities". Required if you want to ensemble ' |
| 'multiple configurations.') |
| parser.add_argument('--continue_prediction', action='store_true', |
| help='Continue an aborted previous prediction (will not overwrite existing files)') |
| parser.add_argument('-chk', type=str, required=False, default='checkpoint_final.pth', |
| help='Name of the checkpoint you want to use. Default: checkpoint_final.pth') |
| parser.add_argument('-npp', type=int, required=False, default=3, |
| help='Number of processes used for preprocessing. More is not always better. Beware of ' |
| 'out-of-RAM issues. Default: 3') |
| parser.add_argument('-nps', type=int, required=False, default=3, |
| help='Number of processes used for segmentation export. More is not always better. Beware of ' |
| 'out-of-RAM issues. Default: 3') |
| parser.add_argument('-prev_stage_predictions', type=str, required=False, default=None, |
| help='Folder containing the predictions of the previous stage. Required for cascaded models.') |
| parser.add_argument('-num_parts', type=int, required=False, default=1, |
| help='Number of separate nnUNetv2_predict call that you will be making. Default: 1 (= this one ' |
| 'call predicts everything)') |
| parser.add_argument('-part_id', type=int, required=False, default=0, |
| help='If multiple nnUNetv2_predict exist, which one is this? IDs start with 0 can end with ' |
| 'num_parts - 1. So when you submit 5 nnUNetv2_predict calls you need to set -num_parts ' |
| '5 and use -part_id 0, 1, 2, 3 and 4. Simple, right? Note: You are yourself responsible ' |
| 'to make these run on separate GPUs! Use CUDA_VISIBLE_DEVICES (google, yo!)') |
| parser.add_argument('-device', type=str, default='cuda', required=False, |
| help="Use this to set the device the inference should run with. Available options are 'cuda' " |
| "(GPU), 'cpu' (CPU) and 'mps' (Apple M1/M2). Do NOT use this to set which GPU ID! " |
| "Use CUDA_VISIBLE_DEVICES=X nnUNetv2_predict [...] instead!") |
| parser.add_argument('--disable_progress_bar', action='store_true', required=False, default=False, |
| help='Set this flag to disable progress bar. Recommended for HPC environments (non interactive ' |
| 'jobs)') |
| parser.add_argument('--rec', type=str, default='mean', choices=['mean', 'median'], |
| help='Method of reconstruction: mean or median. Default is mean.') |
| |
| args = parser.parse_args() |
| |
| args.f = ' '.join(args.f) |
|
|
| os.system(f"nnUNetv2_predict -d {args.d} -i {args.i} -o {args.o} -c {args.c} -p {args.p} -tr {args.tr} -f {args.f} -chk {args.chk} ") |
|
|
| |
| if args.revert_norm: |
| print("Reverting normalisation...") |
|
|
| model_folder = get_output_folder(args.d, args.tr, args.p, args.c) |
| dataset_name = model_folder.split(os.sep)[-2] |
| ct_plan_path = os.path.join(os.environ["nnUNet_preprocessed"], dataset_name, f"gt_{args.p}.json") |
| if not os.path.exists(ct_plan_path): |
| ct_plan_path = os.path.join(os.environ["nnUNet_preprocessed"], dataset_name, 'gt_plan', f"{args.p}.json") |
| ct_mean, ct_std = get_ct_normalisation_values(ct_plan_path) |
| revert_normalisation(args.o, ct_mean, ct_std, save_path=args.o + "_revert_norm", \ |
| mask_path=args.m, mask_outside_value=-1000) |
| |
| if os.path.exists(args.o): |
| shutil.move(args.o, args.o + "_revert_norm/backup_normalised") |
|
|
| if __name__ == '__main__': |
| nnsyn_predict_entry() |
| |
