def split_4d(input_folder, num_processes=default_num_threads, overwrite_task_output_id=None): assert isdir(join(input_folder, "imagesTr")) and isdir(join(input_folder, "labelsTr")) and \ isfile(join(input_folder, "dataset.json")), \ "The input folder must be a valid Task folder from the Medical Segmentation Decathlon with at least the " \ "imagesTr and labelsTr subfolders and the dataset.json file" while input_folder.endswith("/") or input_folder.endswith("\\"): input_folder = input_folder[:-1] full_task_name = Path(input_folder).parts[-1] assert full_task_name.startswith( "Task" ), "The input folder must point to a folder that starts with TaskXX_" first_underscore = full_task_name.find("_") assert first_underscore == 6, "Input folder start with TaskXX with XX being a 3-digit id: 00, 01, 02 etc" input_task_id = int(full_task_name[4:6]) if overwrite_task_output_id is None: overwrite_task_output_id = input_task_id task_name = full_task_name[7:] output_folder = join(nnUNet_raw_data, "Task%03.0d_" % overwrite_task_output_id + task_name) if isdir(output_folder): shutil.rmtree(output_folder) files = [] output_dirs = [] maybe_mkdir_p(output_folder) for subdir in ["imagesTr", "imagesTs"]: curr_out_dir = join(output_folder, subdir) if not isdir(curr_out_dir): maybe_mkdir_p(curr_out_dir) curr_dir = join(input_folder, subdir) nii_files = [ join(curr_dir, i) for i in os.listdir(curr_dir) if i.endswith(".nii.gz") ] nii_files.sort() for n in nii_files: files.append(n) output_dirs.append(curr_out_dir) shutil.copytree(join(input_folder, "labelsTr"), join(output_folder, "labelsTr")) p = Pool(num_processes) p.starmap(split_4d_nifti, zip(files, output_dirs)) p.close() p.join() shutil.copy(join(input_folder, "dataset.json"), output_folder)
def crop(task_string, override=False, num_threads=default_num_threads): cropped_out_dir = join(nnUNet_cropped_data, task_string) os.makedirs(cropped_out_dir, exist_ok=True) if override and isdir(cropped_out_dir): shutil.rmtree(cropped_out_dir) os.makedirs(cropped_out_dir, exist_ok=True) splitted_4d_output_dir_task = join(nnUNet_raw_data, task_string) lists, _ = create_lists_from_splitted_dataset(splitted_4d_output_dir_task) imgcrop = ImageCropper(num_threads, cropped_out_dir) imgcrop.run_cropping(lists, overwrite_existing=override) shutil.copy(join(nnUNet_raw_data, task_string, "dataset.json"), cropped_out_dir)
def crop(task_string, override=False, num_threads=default_num_threads): cropped_out_dir = join(nnUNet_cropped_data, task_string) maybe_mkdir_p(cropped_out_dir) if override and isdir(cropped_out_dir): shutil.rmtree(cropped_out_dir) maybe_mkdir_p(cropped_out_dir) splitted_4d_output_dir_task = join(nnUNet_raw_data, task_string) lists, _ = create_lists_from_splitted_dataset(splitted_4d_output_dir_task) # from nnunet.preprocessing.cropping import ImageCropper imgcrop = ImageCropper(num_threads, cropped_out_dir) # run_cropping imgcrop.run_cropping(lists, overwrite_existing=override) shutil.copy(join(nnUNet_raw_data, task_string, "dataset.json"), cropped_out_dir)
def crop(task_string, override=False, num_threads=default_num_threads): cropped_out_dir = tuFramework_cropped_data + "/" + task_string if not os.path.isdir(cropped_out_dir): os.makedirs(cropped_out_dir) if override and isdir(cropped_out_dir): shutil.rmtree(cropped_out_dir) if not os.path.isdir(cropped_out_dir): os.makedirs(cropped_out_dir) splitted_4d_output_dir_task = tuFramework_raw_data + "/" + task_string lists, _ = create_lists_from_splitted_dataset(splitted_4d_output_dir_task) imgcrop = ImageCropper(num_threads, cropped_out_dir) imgcrop.run_cropping(lists, overwrite_existing=override) shutil.copy( tuFramework_raw_data + "/" + task_string + "/" + "dataset.json", cropped_out_dir)
def main(): parser = argparse.ArgumentParser() parser.add_argument( "-i", '--input_folder', help="Must contain all modalities for each patient in the correct" " order (same as training). Files must be named " "CASENAME_XXXX.nii.gz where XXXX is the modality " "identifier (0000, 0001, etc)", required=True) parser.add_argument('-o', "--output_folder", required=True, help="folder for saving predictions") parser.add_argument('-t', '--task_name', help='task name or task ID, required.', default=default_plans_identifier, required=True) parser.add_argument( '-tr', '--trainer_class_name', help= 'Name of the nnUNetTrainer used for 2D U-Net, full resolution 3D U-Net and low resolution ' 'U-Net. The default is %s. If you are running inference with the cascade and the folder ' 'pointed to by --lowres_segmentations does not contain the segmentation maps generated by ' 'the low resolution U-Net then the low resolution segmentation maps will be automatically ' 'generated. For this case, make sure to set the trainer class here that matches your ' '--cascade_trainer_class_name (this part can be ignored if defaults are used).' % default_trainer, required=False, default=default_trainer) parser.add_argument( '-ctr', '--cascade_trainer_class_name', help= "Trainer class name used for predicting the 3D full resolution U-Net part of the cascade." "Default is %s" % default_cascade_trainer, required=False, default=default_cascade_trainer) parser.add_argument( '-m', '--model', help= "2d, 3d_lowres, 3d_fullres or 3d_cascade_fullres. Default: 3d_fullres", default="3d_fullres", required=False) parser.add_argument( '-p', '--plans_identifier', help='do not touch this unless you know what you are doing', default=default_plans_identifier, required=False) parser.add_argument( '-f', '--folds', nargs='+', default='None', help= "folds to use for prediction. Default is None which means that folds will be detected " "automatically in the model output folder") parser.add_argument( '-z', '--save_npz', required=False, action='store_true', help= "use this if you want to ensemble these predictions with those of other models. Softmax " "probabilities will be saved as compressed numpy arrays in output_folder and can be " "merged between output_folders with nnUNet_ensemble_predictions") parser.add_argument( '-l', '--lowres_segmentations', required=False, default='None', help= "if model is the highres stage of the cascade then you can use this folder to provide " "predictions from the low resolution 3D U-Net. If this is left at default, the " "predictions will be generated automatically (provided that the 3D low resolution U-Net " "network weights are present") parser.add_argument("--part_id", type=int, required=False, default=0, help="Used to parallelize the prediction of " "the folder over several GPUs. If you " "want to use n GPUs to predict this " "folder you need to run this command " "n times with --part_id=0, ... n-1 and " "--num_parts=n (each with a different " "GPU (for example via " "CUDA_VISIBLE_DEVICES=X)") parser.add_argument("--num_parts", type=int, required=False, default=1, help="Used to parallelize the prediction of " "the folder over several GPUs. If you " "want to use n GPUs to predict this " "folder you need to run this command " "n times with --part_id=0, ... n-1 and " "--num_parts=n (each with a different " "GPU (via " "CUDA_VISIBLE_DEVICES=X)") parser.add_argument( "--num_threads_preprocessing", required=False, default=6, type=int, help= "Determines many background processes will be used for data preprocessing. Reduce this if you " "run into out of memory (RAM) problems. Default: 6") parser.add_argument( "--num_threads_nifti_save", required=False, default=2, type=int, help= "Determines many background processes will be used for segmentation export. Reduce this if you " "run into out of memory (RAM) problems. Default: 2") parser.add_argument( "--disable_tta", required=False, default=False, action="store_true", help= "set this flag to disable test time data augmentation via mirroring. Speeds up inference " "by roughly factor 4 (2D) or 8 (3D)") parser.add_argument( "--overwrite_existing", required=False, default=False, action="store_true", help= "Set this flag if the target folder contains predictions that you would like to overwrite" ) parser.add_argument("--mode", type=str, default="normal", required=False, help="Hands off!") parser.add_argument("--all_in_gpu", type=str, default="None", required=False, help="can be None, False or True. " "Do not touch.") parser.add_argument("--step_size", type=float, default=0.5, required=False, help="don't touch") # parser.add_argument("--interp_order", required=False, default=3, type=int, # help="order of interpolation for segmentations, has no effect if mode=fastest. Do not touch this.") # parser.add_argument("--interp_order_z", required=False, default=0, type=int, # help="order of interpolation along z is z is done differently. Do not touch this.") # parser.add_argument("--force_separate_z", required=False, default="None", type=str, # help="force_separate_z resampling. Can be None, True or False, has no effect if mode=fastest. " # "Do not touch this.") parser.add_argument( '-chk', help='checkpoint name, default: model_final_checkpoint', required=False, default='model_final_checkpoint') parser.add_argument( '--disable_mixed_precision', default=False, action='store_true', required=False, help= 'Predictions are done with mixed precision by default. This improves speed and reduces ' 'the required vram. If you want to disable mixed precision you can set this flag. Note ' 'that yhis is not recommended (mixed precision is ~2x faster!)') ### ----------- added by Camila parser.add_argument( '--disable_sliding_window', default=False, action='store_true', required=False, help='Disable sliding window to predict the whole image') ### ----------- end added by Camila args = parser.parse_args() input_folder = args.input_folder output_folder = args.output_folder part_id = args.part_id num_parts = args.num_parts folds = args.folds save_npz = args.save_npz lowres_segmentations = args.lowres_segmentations num_threads_preprocessing = args.num_threads_preprocessing num_threads_nifti_save = args.num_threads_nifti_save disable_tta = args.disable_tta step_size = args.step_size # interp_order = args.interp_order # interp_order_z = args.interp_order_z # force_separate_z = args.force_separate_z overwrite_existing = args.overwrite_existing mode = args.mode all_in_gpu = args.all_in_gpu model = args.model trainer_class_name = args.trainer_class_name cascade_trainer_class_name = args.cascade_trainer_class_name ### ----------- added by Camila disable_sliding_window = args.disable_sliding_window ### ----------- end added by Camila task_name = args.task_name if not task_name.startswith("Task"): task_id = int(task_name) task_name = convert_id_to_task_name(task_id) assert model in ["2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"], "-m must be 2d, 3d_lowres, 3d_fullres or " \ "3d_cascade_fullres" # if force_separate_z == "None": # force_separate_z = None # elif force_separate_z == "False": # force_separate_z = False # elif force_separate_z == "True": # force_separate_z = True # else: # raise ValueError("force_separate_z must be None, True or False. Given: %s" % force_separate_z) if lowres_segmentations == "None": lowres_segmentations = None if isinstance(folds, list): if folds[0] == 'all' and len(folds) == 1: pass else: folds = [int(i) for i in folds] elif folds == "None": folds = None else: raise ValueError("Unexpected value for argument folds") assert all_in_gpu in ['None', 'False', 'True'] if all_in_gpu == "None": all_in_gpu = None elif all_in_gpu == "True": all_in_gpu = True elif all_in_gpu == "False": all_in_gpu = False # we need to catch the case where model is 3d cascade fullres and the low resolution folder has not been set. # In that case we need to try and predict with 3d low res first if model == "3d_cascade_fullres" and lowres_segmentations is None: print( "lowres_segmentations is None. Attempting to predict 3d_lowres first..." ) assert part_id == 0 and num_parts == 1, "if you don't specify a --lowres_segmentations folder for the " \ "inference of the cascade, custom values for part_id and num_parts " \ "are not supported. If you wish to have multiple parts, please " \ "run the 3d_lowres inference first (separately)" model_folder_name = join( network_training_output_dir, "3d_lowres", task_name, trainer_class_name + "__" + args.plans_identifier) assert isdir( model_folder_name ), "model output folder not found. Expected: %s" % model_folder_name lowres_output_folder = join(output_folder, "3d_lowres_predictions") predict_from_folder(model_folder_name, input_folder, lowres_output_folder, folds, False, num_threads_preprocessing, num_threads_nifti_save, None, part_id, num_parts, not disable_tta, overwrite_existing=overwrite_existing, mode=mode, overwrite_all_in_gpu=all_in_gpu, mixed_precision=not args.disable_mixed_precision, step_size=step_size, disable_sliding_window=disable_sliding_window) lowres_segmentations = lowres_output_folder torch.cuda.empty_cache() print("3d_lowres done") if model == "3d_cascade_fullres": trainer = cascade_trainer_class_name else: trainer = trainer_class_name model_folder_name = join(network_training_output_dir, model, task_name, trainer + "__" + args.plans_identifier) print("using model stored in ", model_folder_name) assert isdir( model_folder_name ), "model output folder not found. Expected: %s" % model_folder_name predict_from_folder(model_folder_name, input_folder, output_folder, folds, save_npz, num_threads_preprocessing, num_threads_nifti_save, lowres_segmentations, part_id, num_parts, not disable_tta, overwrite_existing=overwrite_existing, mode=mode, overwrite_all_in_gpu=all_in_gpu, mixed_precision=not args.disable_mixed_precision, step_size=step_size, checkpoint_name=args.chk, disable_sliding_window=disable_sliding_window)
input_folder = args.input_folder output_folder = args.output_folder part_id = args.part_id num_parts = args.num_parts folds = args.folds save_npz = args.save_npz lowres_segmentations = args.lowres_segmentations num_threads_preprocessing = args.num_threads_preprocessing num_threads_nifti_save = args.num_threads_nifti_save tta = args.tta overwrite = args.overwrite_existing network_training_output_dir = args.training_dir output_folder_name = args.model print("using model stored in ", output_folder_name) assert isdir(output_folder_name ), "model output folder not found: %s" % output_folder_name if lowres_segmentations == "None": lowres_segmentations = None if isinstance(folds, list): if folds[0] == 'all' and len(folds) == 1: pass else: folds = [int(i) for i in folds] elif folds == "None": folds = None else: raise ValueError("Unexpected value for argument folds") if tta == 0:
# -*- coding: utf-8 -*- """ Created on Sat Jul 10 16:11:54 2021 @author: linhai """ import sys import inspect import os from pathlib import Path from batchgenerators.utilities.file_and_folder_operations import join, isdir, maybe_mkdir_p, subfiles, subdirs, isfile #print (sys.path) curDir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) parDir = os.path.dirname(curDir) sys.path.insert(0, parDir) #sys.path.insert(0, ) p1 = 'C:\\Research\\IMA_on_segmentation\\nnUnet\\nnUNet\\rawData\\nnUNet_raw_data\\Task05_Prostate\\imagesTr' p2 = 'C:/Research/IMA_on_segmentation/nnUnet/nnUNet/rawData/nnUNet_raw_data\\Task05_Prostate' p3 = 'C:\\Research\\IMA_on_segmentation\\aaa' p4 = 'C:/Research/IMA_on_segmentation/333/aaab' print (os.path.join(p1, "aaa")+"\\") print (isdir(join(p1, "aaa")+"\\")) print(p1) print (isdir(p2)) #os.mkdir(p4) maybe_mkdir_p(p4) #os.makedirs(p4, exist_ok=True)
def predict_simple(input_folder, output_folder, gt_folder ,task_id, default_trainer, model, folds, save_npz, gpus, disable_mixed_precision, mode, using_pretrain, overwrite_existing, eval_flag): # default_trainer: nnUNetTrainerV2, can change to nnUNetTrainerV2_DP or nnUNetTrainerV2_DDP trainer_class_name = default_trainer cascade_trainer_class_name = default_cascade_trainer plans_identifier = default_plans_identifier num_threads_preprocessing = 6 num_threads_nifti_save = 2 input_folder = join(DATASET_DIR,input_folder) part_id = gpus-1 num_parts = gpus disable_tta = False step_size = 0.5 all_in_gpu = "None" task_name = task_id if not task_name.startswith("Task"): task_id = int(task_name) task_name = convert_id_to_task_name(task_id) assert model in ["2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"], "-m must be 2d, 3d_lowres, 3d_fullres or 3d_cascade_fullres" output_folder = join(DATASET_DIR,output_folder,model,task_name) lowres_segmentations = None if isinstance(folds, list): if folds[0] == 'all' and len(folds) == 1: pass else: folds = [int(i) for i in folds] elif folds == "None": folds = None else: raise ValueError("Unexpected value for argument folds") assert all_in_gpu in ['None', 'False', 'True'] if all_in_gpu == "None": all_in_gpu = None elif all_in_gpu == "True": all_in_gpu = True elif all_in_gpu == "False": all_in_gpu = False if model == "3d_cascade_fullres" and lowres_segmentations is None: print("lowres_segmentations is None. Attempting to predict 3d_lowres first...") assert part_id == 0 and num_parts == 1, "if you don't specify a --lowres_segmentations folder for the " \ "inference of the cascade, custom values for part_id and num_parts " \ "are not supported. If you wish to have multiple parts, please " \ "run the 3d_lowres inference first (separately)" lowres_output_folder = join(output_folder, "3d_lowres_predictions") if using_pretrain: model_folder_name = join(pre_training_output_dir, "3d_lowres", task_name, trainer_class_name + "__" + plans_identifier) predict_from_folder(model_folder_name, input_folder, lowres_output_folder, folds, False, num_threads_preprocessing, num_threads_nifti_save, None, part_id, num_parts, not disable_tta, overwrite_existing=overwrite_existing, mode=mode, overwrite_all_in_gpu=all_in_gpu, mixed_precision=not disable_mixed_precision, step_size=step_size, checkpoint_name="model_final_checkpoint") else: model_folder_name = join(network_training_output_dir, "3d_lowres", task_name, trainer_class_name + "__" + plans_identifier) predict_from_folder(model_folder_name, input_folder, lowres_output_folder, folds, False, num_threads_preprocessing, num_threads_nifti_save, None, part_id, num_parts, not disable_tta, overwrite_existing=overwrite_existing, mode=mode, overwrite_all_in_gpu=all_in_gpu, mixed_precision=not disable_mixed_precision, step_size=step_size, checkpoint_name ="model_best") assert isdir(model_folder_name), "model output folder not found. Expected: %s" % model_folder_name lowres_segmentations = lowres_output_folder torch.cuda.empty_cache() print("3d_lowres done") if model == "3d_cascade_fullres": trainer = cascade_trainer_class_name else: if model == "2d": assert mode == "normal", "please set mode = normal for 2D predict" trainer = trainer_class_name if using_pretrain: model_folder_name = join(pre_training_output_dir, model, task_name, trainer + "__" + plans_identifier) predict_from_folder(model_folder_name, input_folder, output_folder, folds, save_npz, num_threads_preprocessing, num_threads_nifti_save, lowres_segmentations, part_id, num_parts, not disable_tta, overwrite_existing=overwrite_existing, mode=mode, overwrite_all_in_gpu=all_in_gpu, mixed_precision=not disable_mixed_precision, step_size=step_size, checkpoint_name="model_final_checkpoint") else: model_folder_name = join(network_training_output_dir, model, task_name, trainer + "__" + plans_identifier) predict_from_folder(model_folder_name, input_folder, output_folder, folds, save_npz, num_threads_preprocessing, num_threads_nifti_save, lowres_segmentations, part_id, num_parts, not disable_tta, overwrite_existing=overwrite_existing, mode=mode, overwrite_all_in_gpu=all_in_gpu, mixed_precision=not disable_mixed_precision, step_size=step_size, checkpoint_name="model_best") print("using model stored in ", model_folder_name) assert isdir(model_folder_name), "model output folder not found. Expected: %s" % model_folder_name if eval_flag: # task = output_folder.split('/')[-1] gt_folder = join(DATASET_DIR,gt_folder) predict_val(output_folder,gt_folder)