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)
