def main():
parser = argparse.ArgumentParser()
parser.add_argument("-i", "--input_folder", type=str, required=True,
help="folder with input files, files must be named PATIENTID_0000.nii.gz, "
"PATIENTID_0001.nii.gz, PATIENTID_0002.nii.gz, PATIENTID_0003.nii.gz for T1, "
"T1c, T2 and Flair, respectively. There can be an arbitrary number of patients "
"in the folder (PATIENTID can be anything). CAREFUL: The files MUST fullfill the "
"following criteria: 1) They must be brain extracted with the non-brain region being "
"0 (you can achieve that by using hd-bet (https://github.com/MIC-DKFZ/HD-BET); 2) "
"They must be coregistered and in the same co-ordinate system (pixels arrays must be "
"aligned) 3) makse sure the T1, T1c, T2 and FLAIR file always have the correct "
"file ending (_0000.nii.gz, ...)")
parser.add_argument("-o", "--output_folder", type=str, required=True,
help="output folder. This is there the resulting segmentations will be saved (as PATIENT_ID."
"nii.gz). Cannot be the same folder as the input folder. If output_folder does not exist "
"it will be created")
parser.add_argument("-p", "--processes", default=4, type=str, required=False,
help="number of processes for data preprocessing and nifti export. You should not have to "
"touch this. So don't unless there is a clear indication that it is required. Default: 4")
parser.add_argument('--overwrite_existing', default=True, type=str, required=False,
help="set to False to keep segmentations in output_folder and continue where you left off "
"(useful if something crashes). If True then all segmentations that may already be "
"present in output_folder will be overwritten. Default: True")
args = parser.parse_args()
input_folder = args.input_folder
output_folder = args.output_folder
processes = args.processes
overwrite_existing = args.overwrite_existing
maybe_download_weights()
predict_from_folder(folder_with_parameter_files, input_folder, output_folder, (0, ), False, processes, processes,
None, 0, 1, True, overwrite_existing=overwrite_existing, checkpoint_name='model_final_checkpoint')
def _run_interface(self, runtime):
input_folder = self.inputs.input_folder
output_folder = os.path.abspath(self.inputs.output_folder)
part_id = self.inputs.part_id
num_parts = self.inputs.num_parts
folds = self.inputs.folds
save_npz = self.inputs.save_npz
lowres_segmentations = self.inputs.lowres_segmentations
num_threads_preprocessing = self.inputs.threads_preprocessing
num_threads_nifti_save = self.inputs.threads_save
tta = self.inputs.tta
overwrite = self.inputs.overwrite
model_folder = self.inputs.model_folder
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 == 6:
folds = None
else:
raise ValueError("Unexpected value for argument folds")
if tta == 0:
tta = False
elif tta == 1:
tta = True
else:
raise ValueError("Unexpected value for tta, Use 1 or 0")
if overwrite == 0:
overwrite = False
elif overwrite == 1:
overwrite = True
else:
raise ValueError("Unexpected value for overwrite, Use 1 or 0")
predict_from_folder(model_folder,
input_folder,
output_folder,
folds,
save_npz,
num_threads_preprocessing,
num_threads_nifti_save,
lowres_segmentations,
part_id,
num_parts,
tta,
overwrite_existing=overwrite)
torch.cuda.empty_cache()
return runtime
copyfile(input_folder, join(tmp_nifti_folder, nifti_filename))
elif input_folder.endswith('.nrrd'):
_nrrd = nrrd.read(input_folder)
data = _nrrd[0]
header = _nrrd[1]
print(input_folder)
print('DATA:', data)
print("HEADER:", header)
#nrrd = sitk.ReadImage(input_folder, sitk.sitkFloat64)
#data = sitk.GetArrayFromImage(nrrd)
x = list(map(float, header['srow_x'].split(' ')))
y = list(map(float, header['srow_y'].split(' ')))
z = list(map(float, header['srow_z'].split(' ')))
affine = np.vstack([x, y, z])
affine = np.concatenate([affine, np.expand_dims(np.array([0, 0, 0, 1]), axis=0)], axis=0)
data = data.astype(np.float64)
img = nib.Nifti1Image(data, affine)
nib.save(img,os.path.join(tmp_nifti_folder, nifti_filename))
else:
print('Error - unrecognized file format.')
predict_from_folder(output_folder_name, tmp_nifti_folder, output_folder, folds, save_npz, num_threads_preprocessing,
num_threads_nifti_save, lowres_segmentations, part_id, num_parts, tta,
overwrite_existing=overwrite)
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)
if overwrite == 0:
overwrite = False
elif overwrite == 1:
overwrite = True
else:
raise ValueError("Unexpected value for overwrite, Use 1 or 0")
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
predict_from_folder(output_folder_name,
input_folder,
output_folder,
folds,
save_npz,
num_threads_preprocessing,
num_threads_nifti_save,
lowres_segmentations,
part_id,
num_parts,
tta,
overwrite_existing=overwrite,
mode=mode,
overwrite_all_in_gpu=all_in_gpu)
