def predict_next_stage(trainer, stage_to_be_predicted_folder): output_folder = join(pardir(trainer.output_folder), "pred_next_stage") maybe_mkdir_p(output_folder) process_manager = Pool(2) results = [] for pat in trainer.dataset_val.keys(): print(pat) data_file = trainer.dataset_val[pat]['data_file'] data_preprocessed = np.load(data_file)['data'][:-1] predicted = trainer.predict_preprocessed_data_return_softmax( data_preprocessed, True, 1, False, 1, trainer.data_aug_params['mirror_axes'], True, True, 2, trainer.patch_size, True) data_file_nofolder = data_file.split("/")[-1] data_file_nextstage = join(stage_to_be_predicted_folder, data_file_nofolder) data_nextstage = np.load(data_file_nextstage)['data'] target_shp = data_nextstage.shape[1:] output_file = join( output_folder, data_file_nextstage.split("/")[-1][:-4] + "_segFromPrevStage.npz") results.append( process_manager.starmap_async( resample_and_save, [(predicted, target_shp, output_file)])) _ = [i.get() for i in results]
def initialize(self, training=True, force_load_plans=False): if not self.was_initialized: maybe_mkdir_p(self.output_folder) if force_load_plans or (self.plans is None): self.load_plans_file() self.process_plans(self.plans) self.setup_DA_params() ################# Here we wrap the loss for deep supervision ############ # we need to know the number of outputs of the network net_numpool = len(self.net_num_pool_op_kernel_sizes) # we give each output a weight which decreases exponentially (division by 2) as the resolution decreases # this gives higher resolution outputs more weight in the loss weights = np.array([1 / (2 ** i) for i in range(net_numpool)]) # we don't use the lowest 2 outputs. Normalize weights so that they sum to 1 mask = np.array([True if i < net_numpool - 1 else False for i in range(net_numpool)]) weights[~mask] = 0 weights = weights / weights.sum() # now wrap the loss self.loss = MultipleOutputLoss2(self.loss, weights) ################# END ################### self.folder_with_preprocessed_data = join(self.dataset_directory, self.plans['data_identifier'] + "_stage%d" % self.stage) if training: self.dl_tr, self.dl_val = self.get_basic_generators() if self.unpack_data: print("unpacking dataset") unpack_dataset(self.folder_with_preprocessed_data) print("done") else: print( "INFO: Not unpacking data! Training may be slow due to that. Pray you are not using 2d or you " "will wait all winter for your model to finish!") self.tr_gen, self.val_gen = get_no_augmentation(self.dl_tr, self.dl_val, self.data_aug_params[ 'patch_size_for_spatialtransform'], self.data_aug_params, deep_supervision_scales=self.deep_supervision_scales) self.print_to_log_file("TRAINING KEYS:\n %s" % (str(self.dataset_tr.keys())), also_print_to_console=False) self.print_to_log_file("VALIDATION KEYS:\n %s" % (str(self.dataset_val.keys())), also_print_to_console=False) else: pass self.initialize_network() self.initialize_optimizer_and_scheduler() assert isinstance(self.network, (SegmentationNetwork, nn.DataParallel)) else: self.print_to_log_file('self.was_initialized is True, not running self.initialize again') self.was_initialized = True
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 plan_and_preprocess(task_string, processes_lowres=default_num_threads, processes_fullres=3, no_preprocessing=False): from nnunet.experiment_planning.experiment_planner_baseline_2DUNet import ExperimentPlanner2D from nnunet.experiment_planning.experiment_planner_baseline_3DUNet import ExperimentPlanner preprocessing_output_dir_this_task_train = join(preprocessing_output_dir, task_string) cropped_out_dir = join(nnUNet_cropped_data, task_string) maybe_mkdir_p(preprocessing_output_dir_this_task_train) shutil.copy(join(cropped_out_dir, "dataset_properties.pkl"), preprocessing_output_dir_this_task_train) shutil.copy(join(nnUNet_raw_data, task_string, "dataset.json"), preprocessing_output_dir_this_task_train) exp_planner = ExperimentPlanner(cropped_out_dir, preprocessing_output_dir_this_task_train) exp_planner.plan_experiment() if not no_preprocessing: exp_planner.run_preprocessing((processes_lowres, processes_fullres)) exp_planner = ExperimentPlanner2D( cropped_out_dir, preprocessing_output_dir_this_task_train) exp_planner.plan_experiment() if not no_preprocessing: exp_planner.run_preprocessing(processes_fullres) # write which class is in which slice to all training cases (required to speed up 2D Dataloader) # This is done for all data so that if we wanted to use them with 2D we could do so if not no_preprocessing: p = Pool(default_num_threads) # if there is more than one my_data_identifier (different brnaches) then this code will run for all of them if # they start with the same string. not problematic, but not pretty stages = [ i for i in subdirs( preprocessing_output_dir_this_task_train, join=True, sort=True) if Path(i).parts[-1].find("stage") != -1 ] for s in stages: print(Path(s).parts[-1]) list_of_npz_files = subfiles(s, True, None, ".npz", True) list_of_pkl_files = [i[:-4] + ".pkl" for i in list_of_npz_files] all_classes = [] for pk in list_of_pkl_files: with open(pk, 'rb') as f: props = pickle.load(f) all_classes_tmp = np.array(props['classes']) all_classes.append(all_classes_tmp[all_classes_tmp >= 0]) p.map(add_classes_in_slice_info, zip(list_of_npz_files, list_of_pkl_files, all_classes)) p.close() p.join()
def run_inference_and_store_results(self,output_file_tag=''): output_file_base_name = output_file_tag + "_nnunet_seg.nii.gz" # passing only lists of length one to predict_cases for inner_list in self.data.inference_loader: list_of_lists = [inner_list] # output filenames (list of one) include information about patient folder name # infering patient folder name from all file paths for a sanity check # (should give the same answer) folder_names = [fpath.split('/')[-2] for fpath in inner_list] if set(folder_names) != set(folder_names[:1]): raise RuntimeError('Patient file paths: {} were found to come from different folders against expectation.'.format(inner_list)) patient_folder_name = folder_names[0] output_filename = patient_folder_name + output_file_base_name final_out_folder = join(self.intermediate_out_folder, patient_folder_name) intermediate_output_folders = [] for model_name, folds in zip(self.model_list, self.folds_list): output_model = join(self.intermediate_out_folder, model_name) intermediate_output_folders.append(output_model) intermediate_output_filepaths = [join(output_model, output_filename)] maybe_mkdir_p(output_model) params_folder_model = join(self.params_folder, model_name) predict_cases(model=params_folder_model, list_of_lists=list_of_lists, output_filenames=intermediate_output_filepaths, folds=folds, save_npz=True, num_threads_preprocessing=1, num_threads_nifti_save=1, segs_from_prev_stage=None, do_tta=True, mixed_precision=True, overwrite_existing=True, all_in_gpu=False, step_size=0.5) merge(folders=intermediate_output_folders, output_folder=final_out_folder, threads=1, override=True, postprocessing_file=None, store_npz=False) f = join(final_out_folder, output_filename) apply_brats_threshold(f, f, self.threshold, self.replace_with) load_convert_save(f) _ = [shutil.rmtree(i) for i in intermediate_output_folders]
def predict_next_stage(trainer, stage_to_be_predicted_folder): output_folder = join(pardir(trainer.output_folder), "pred_next_stage") maybe_mkdir_p(output_folder) if 'segmentation_export_params' in trainer.plans.keys(): force_separate_z = trainer.plans['segmentation_export_params'][ 'force_separate_z'] interpolation_order = trainer.plans['segmentation_export_params'][ 'interpolation_order'] interpolation_order_z = trainer.plans['segmentation_export_params'][ 'interpolation_order_z'] else: force_separate_z = None interpolation_order = 1 interpolation_order_z = 0 export_pool = Pool(2) results = [] for pat in trainer.dataset_val.keys(): print(pat) data_file = trainer.dataset_val[pat]['data_file'] data_preprocessed = np.load(data_file)['data'][:-1] predicted_probabilities = trainer.predict_preprocessed_data_return_seg_and_softmax( data_preprocessed, do_mirroring=trainer.data_aug_params["do_mirror"], mirror_axes=trainer.data_aug_params['mirror_axes'], mixed_precision=trainer.fp16)[1] data_file_nofolder = data_file.split("/")[-1] data_file_nextstage = join(stage_to_be_predicted_folder, data_file_nofolder) data_nextstage = np.load(data_file_nextstage)['data'] target_shp = data_nextstage.shape[1:] output_file = join( output_folder, data_file_nextstage.split("/")[-1][:-4] + "_segFromPrevStage.npz") if np.prod(predicted_probabilities.shape) > ( 2e9 / 4 * 0.85): # *0.85 just to be save np.save(output_file[:-4] + ".npy", predicted_probabilities) predicted_probabilities = output_file[:-4] + ".npy" results.append( export_pool.starmap_async(resample_and_save, [ (predicted_probabilities, target_shp, output_file, force_separate_z, interpolation_order, interpolation_order_z) ])) _ = [i.get() for i in results] export_pool.close() export_pool.join()
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) 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 initialize(self, training=True, force_load_plans=False): """ For prediction of test cases just set training=False, this will prevent loading of training data and training batchgenerator initialization :param training: :return: """ maybe_mkdir_p(self.output_folder) if force_load_plans or (self.plans is None): self.load_plans_file() self.process_plans(self.plans) self.setup_DA_params() self.folder_with_preprocessed_data = join( self.dataset_directory, self.plans['data_identifier'] + "_stage%d" % self.stage) if training: self.dl_tr, self.dl_val = self.get_basic_generators() if self.unpack_data: print("unpacking dataset") unpack_dataset(self.folder_with_preprocessed_data) print("done") else: print( "INFO: Not unpacking data! Training may be slow due to that. Pray you are not using 2d or you " "will wait all winter for your model to finish!") self.tr_gen, self.val_gen = get_no_augmentation( self.dl_tr, self.dl_val, self.data_aug_params['patch_size_for_spatialtransform'], self.data_aug_params) self.print_to_log_file("TRAINING KEYS:\n %s" % (str(self.dataset_tr.keys())), also_print_to_console=False) self.print_to_log_file("VALIDATION KEYS:\n %s" % (str(self.dataset_val.keys())), also_print_to_console=False) else: pass self.initialize_network() assert isinstance(self.network, (SegmentationNetwork, nn.DataParallel)) self.was_initialized = True self.data_aug_params['mirror_axes'] = ()
def predict_next_stage(trainer, stage_to_be_predicted_folder, force_separate_z=False, interpolation_order=1, interpolation_order_z=0): output_folder = join(pardir(trainer.output_folder), "pred_next_stage") maybe_mkdir_p(output_folder) export_pool = Pool(2) results = [] for pat in trainer.dataset_val.keys(): print(pat) data_file = trainer.dataset_val[pat]['data_file'] data_preprocessed = np.load(data_file)['data'][:-1] predicted = trainer.predict_preprocessed_data_return_softmax( data_preprocessed, True, 1, False, 1, trainer.data_aug_params['mirror_axes'], True, True, 2, trainer.patch_size, True) data_file_nofolder = data_file.split("/")[-1] data_file_nextstage = join(stage_to_be_predicted_folder, data_file_nofolder) data_nextstage = np.load(data_file_nextstage)['data'] target_shp = data_nextstage.shape[1:] output_file = join( output_folder, data_file_nextstage.split("/")[-1][:-4] + "_segFromPrevStage.npz") if np.prod( predicted.shape) > (2e9 / 4 * 0.85): # *0.85 just to be save np.save(output_file[:-4] + ".npy", predicted) predicted = output_file[:-4] + ".npy" results.append( export_pool.starmap_async( resample_and_save, [(predicted, target_shp, output_file, force_separate_z, interpolation_order, interpolation_order_z)])) _ = [i.get() for i in results] export_pool.close() export_pool.join()
def testing(unet, test_loader, batch_size, device, output_dir): start = time() start_eval_test = 0 idx = 0 maybe_mkdir_p(os.path.join(output_dir, 'images')) maybe_mkdir_p(os.path.join(output_dir, 'preds')) maybe_mkdir_p(os.path.join(output_dir, 'labels')) for batch in test_loader: image, label = batch pred = unet(image.to(device)) pad = int((pred.shape[-1] - label.shape[-1]) / 2) pred = pred[:, :, pad:label.shape[-1] + pad, pad:label.shape[-1] + pad].argmax(dim=1) save_image( image[0, 0, pad:label.shape[-1] + pad, pad:label.shape[-1] + pad], os.path.join(output_dir, 'images', f'image{idx}.tif')) save_image(label[0, 0, :, :].float(), os.path.join(output_dir, 'labels', f'label{idx}.tif')) save_image(pred[0, :, :].float(), os.path.join(output_dir, 'preds', f'pred{idx}.tif')) idx += 1 if start_eval_test == 0: test_eval = evaluation_metrics(pred[0, :, :].detach(), label[0, 0, :, :].detach()) start_eval_test += 1 else: np.concatenate((test_eval, evaluation_metrics(pred[0, :, :].detach(), label[0, 0, :, :].detach())), axis=1) test = np.mean(test_eval, axis=1) test_std = np.std(test_eval, axis=1) test_iou = [test[0], test_std[0]] test_pe = [test[1], test_std[1]] np.savetxt(os.path.join(output_dir, 'test_iou.out'), test_iou) np.savetxt(os.path.join(output_dir, 'test_pe.out'), test_pe) print('Mean IoU testing:', "{:.6f}".format(test[0])) print('Mean PE testing :', "{:.6f}".format(test[1])) print('Testing took :', "{:.6f}".format(time() - start), 's') print(' ') print('Testing is finished')
samp_tr = int(np.round(tr_per * len(train_dataset))) samp_val = int(np.round(val_per * len(train_dataset))) # Round numbers so that we do not exceed total number of samples while samp_tr + samp_val > len(train_dataset): samp_val += -1 # Generate an order vector to shuffle the samples before each fold for the cross validation np.random.seed(SEED) order = np.arange(len(train_dataset)) np.random.shuffle(order) if FOLDS is None: # Training with all available samples # Make directory where we save model and data all_dir = os.path.join(CUR_DIR, 'models', f'{DATASET}', 'all') maybe_mkdir_p(all_dir) val_dataset = ImageDataset_test(root_dir, ISBI2012=ISBI2012) # Suffle and load the training set train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=True) torch.cuda.empty_cache() unet = Unet().to(device) if START_FROM == -1: # load latest model epoch_id = max([
# raw_data_base_dir = "/data0/mzs/zhx/lung_lobe_seg/galaNet_raw_data" # 原始数据保存文件夹 # preprocessed_output_dir = "/data0/mzs/zhx/lung_lobe_seg/galaNet_preprocessed" # 预处理后数据存放处 # network_output_dir_base = "/data0/mzs/zhx/lung_lobe_seg/galaNet_trained_models" # 网络存放处 # raw_data_base_dir = "/home/zenghexiang/data/zenghexiang/lung_lobe_seg/galaNet_raw_data" # 原始数据保存文件夹 # preprocessed_output_dir = "/home/zenghexiang/data/zenghexiang/lung_lobe_seg/galaNet_preprocessed" # 预处理后数据存放处 # network_output_dir_base = "/home/zenghexiang/data/zenghexiang/lung_lobe_seg/galaNet_trained_models" # 网络存放处 if raw_data_base_dir is not None: raw_dicom_data_dir = join(raw_data_base_dir, "dicom_data") # dicom原始数据存放文件夹 raw_cropped_data_dir = join(raw_data_base_dir, "cropped_data") # 原始数据被crop后存放的文件夹 raw_splited_dir = join(raw_data_base_dir, "splited_data") maybe_mkdir_p(raw_data_base_dir) maybe_mkdir_p(raw_cropped_data_dir) else: raise AssertionError( "Attention! raw_data_base_dir is not defined! Please set raw_data_base_dir in paths.py." ) if preprocessed_output_dir is not None: maybe_mkdir_p(preprocessed_output_dir) maybe_mkdir_p(join(preprocessed_output_dir, preprocessed_data_identifer)) maybe_mkdir_p(join(preprocessed_output_dir, preprocessed_net_inputs)) else: raise AssertionError("Attention! preprocessed_output_dir is not defined! " "Please set preprocessed_output_dir in paths.py.") if network_output_dir_base is not None:
def validate(self, do_mirroring: bool = True, use_sliding_window: bool = True, step_size: float = 0.5, save_softmax: bool = True, use_gaussian: bool = True, overwrite: bool = True, validation_folder_name: str = 'validation_raw', debug: bool = False, all_in_gpu: bool = False, segmentation_export_kwargs: dict = None, run_postprocessing_on_folds: bool = True): if isinstance(self.network, DDP): net = self.network.module else: net = self.network ds = net.do_ds net.do_ds = False current_mode = self.network.training self.network.eval() assert self.was_initialized, "must initialize, ideally with checkpoint (or train first)" if self.dataset_val is None: self.load_dataset() self.do_split() if segmentation_export_kwargs is None: if 'segmentation_export_params' in self.plans.keys(): force_separate_z = self.plans['segmentation_export_params'][ 'force_separate_z'] interpolation_order = self.plans['segmentation_export_params'][ 'interpolation_order'] interpolation_order_z = self.plans[ 'segmentation_export_params']['interpolation_order_z'] else: force_separate_z = None interpolation_order = 1 interpolation_order_z = 0 else: force_separate_z = segmentation_export_kwargs['force_separate_z'] interpolation_order = segmentation_export_kwargs[ 'interpolation_order'] interpolation_order_z = segmentation_export_kwargs[ 'interpolation_order_z'] # predictions as they come from the network go here output_folder = join(self.output_folder, validation_folder_name) maybe_mkdir_p(output_folder) # this is for debug purposes my_input_args = { 'do_mirroring': do_mirroring, 'use_sliding_window': use_sliding_window, 'step_size': step_size, 'save_softmax': save_softmax, 'use_gaussian': use_gaussian, 'overwrite': overwrite, 'validation_folder_name': validation_folder_name, 'debug': debug, 'all_in_gpu': all_in_gpu, 'segmentation_export_kwargs': segmentation_export_kwargs, } save_json(my_input_args, join(output_folder, "validation_args.json")) if do_mirroring: if not self.data_aug_params['do_mirror']: raise RuntimeError( "We did not train with mirroring so you cannot do inference with mirroring enabled" ) mirror_axes = self.data_aug_params['mirror_axes'] else: mirror_axes = () pred_gt_tuples = [] export_pool = Pool(default_num_threads) results = [] all_keys = list(self.dataset_val.keys()) my_keys = all_keys[self.local_rank::dist.get_world_size()] # we cannot simply iterate over all_keys because we need to know pred_gt_tuples and valid_labels of all cases # for evaluation (which is done by local rank 0) for k in my_keys: properties = load_pickle(self.dataset[k]['properties_file']) fname = properties['list_of_data_files'][0].split("/")[-1][:-12] pred_gt_tuples.append([ join(output_folder, fname + ".nii.gz"), join(self.gt_niftis_folder, fname + ".nii.gz") ]) if k in my_keys: if overwrite or (not isfile(join(output_folder, fname + ".nii.gz"))) or \ (save_softmax and not isfile(join(output_folder, fname + ".npz"))): data = np.load(self.dataset[k]['data_file'])['data'] print(k, data.shape) data[-1][data[-1] == -1] = 0 softmax_pred = self.predict_preprocessed_data_return_seg_and_softmax( data[:-1], do_mirroring=do_mirroring, mirror_axes=mirror_axes, use_sliding_window=use_sliding_window, step_size=step_size, use_gaussian=use_gaussian, all_in_gpu=all_in_gpu, mixed_precision=self.fp16)[1] softmax_pred = softmax_pred.transpose( [0] + [i + 1 for i in self.transpose_backward]) if save_softmax: softmax_fname = join(output_folder, fname + ".npz") else: softmax_fname = None """There is a problem with python process communication that prevents us from communicating obejcts larger than 2 GB between processes (basically when the length of the pickle string that will be sent is communicated by the multiprocessing.Pipe object then the placeholder (\%i I think) does not allow for long enough strings (lol). This could be fixed by changing i to l (for long) but that would require manually patching system python code. We circumvent that problem here by saving softmax_pred to a npy file that will then be read (and finally deleted) by the Process. save_segmentation_nifti_from_softmax can take either filename or np.ndarray and will handle this automatically""" if np.prod(softmax_pred.shape) > ( 2e9 / 4 * 0.85): # *0.85 just to be save np.save(join(output_folder, fname + ".npy"), softmax_pred) softmax_pred = join(output_folder, fname + ".npy") results.append( export_pool.starmap_async( save_segmentation_nifti_from_softmax, ((softmax_pred, join(output_folder, fname + ".nii.gz"), properties, interpolation_order, self.regions_class_order, None, None, softmax_fname, None, force_separate_z, interpolation_order_z), ))) _ = [i.get() for i in results] self.print_to_log_file("finished prediction") distributed.barrier() if self.local_rank == 0: # evaluate raw predictions self.print_to_log_file("evaluation of raw predictions") task = self.dataset_directory.split("/")[-1] job_name = self.experiment_name _ = aggregate_scores(pred_gt_tuples, labels=list(range(self.num_classes)), json_output_file=join(output_folder, "summary.json"), json_name=job_name + " val tiled %s" % (str(use_sliding_window)), json_author="Fabian", json_task=task, num_threads=default_num_threads) if run_postprocessing_on_folds: # in the old nnunet we would stop here. Now we add a postprocessing. This postprocessing can remove everything # except the largest connected component for each class. To see if this improves results, we do this for all # classes and then rerun the evaluation. Those classes for which this resulted in an improved dice score will # have this applied during inference as well self.print_to_log_file("determining postprocessing") determine_postprocessing( self.output_folder, self.gt_niftis_folder, validation_folder_name, final_subf_name=validation_folder_name + "_postprocessed", debug=debug) # after this the final predictions for the vlaidation set can be found in validation_folder_name_base + "_postprocessed" # They are always in that folder, even if no postprocessing as applied! # detemining postprocesing on a per-fold basis may be OK for this fold but what if another fold finds another # postprocesing to be better? In this case we need to consolidate. At the time the consolidation is going to be # done we won't know what self.gt_niftis_folder was, so now we copy all the niftis into a separate folder to # be used later gt_nifti_folder = join(self.output_folder_base, "gt_niftis") maybe_mkdir_p(gt_nifti_folder) for f in subfiles(self.gt_niftis_folder, suffix=".nii.gz"): success = False attempts = 0 e = None while not success and attempts < 10: try: shutil.copy(f, gt_nifti_folder) success = True except OSError as e: attempts += 1 sleep(1) if not success: print("Could not copy gt nifti file %s into folder %s" % (f, gt_nifti_folder)) if e is not None: raise e self.network.train(current_mode) net.do_ds = ds
def run_training(self): """ if we run with -c then we need to set the correct lr for the first epoch, otherwise it will run the first continued epoch with self.initial_lr we also need to make sure deep supervision in the network is enabled for training, thus the wrapper :return: """ self.maybe_update_lr( self.epoch ) # if we dont overwrite epoch then self.epoch+1 is used which is not what we # want at the start of the training if isinstance(self.network, DDP): net = self.network.module else: net = self.network ds = net.do_ds net.do_ds = True _ = self.tr_gen.next() _ = self.val_gen.next() if torch.cuda.is_available(): torch.cuda.empty_cache() self._maybe_init_amp() maybe_mkdir_p(self.output_folder) self.plot_network_architecture() if cudnn.benchmark and cudnn.deterministic: warn( "torch.backends.cudnn.deterministic is True indicating a deterministic training is desired. " "But torch.backends.cudnn.benchmark is True as well and this will prevent deterministic training! " "If you want deterministic then set benchmark=False") if not self.was_initialized: self.initialize(True) while self.epoch < self.max_num_epochs: self.print_to_log_file("\nepoch: ", self.epoch) epoch_start_time = time() train_losses_epoch = [] # train one epoch self.network.train() if self.use_progress_bar: with trange(self.num_batches_per_epoch) as tbar: for b in tbar: tbar.set_description("Epoch {}/{}".format( self.epoch + 1, self.max_num_epochs)) l = self.run_iteration(self.tr_gen, True) tbar.set_postfix(loss=l) train_losses_epoch.append(l) else: for _ in range(self.num_batches_per_epoch): l = self.run_iteration(self.tr_gen, True) train_losses_epoch.append(l) self.all_tr_losses.append(np.mean(train_losses_epoch)) self.print_to_log_file("train loss : %.4f" % self.all_tr_losses[-1]) with torch.no_grad(): # validation with train=False self.network.eval() val_losses = [] for b in range(self.num_val_batches_per_epoch): l = self.run_iteration(self.val_gen, False, True) val_losses.append(l) self.all_val_losses.append(np.mean(val_losses)) self.print_to_log_file("validation loss: %.4f" % self.all_val_losses[-1]) if self.also_val_in_tr_mode: self.network.train() # validation with train=True val_losses = [] for b in range(self.num_val_batches_per_epoch): l = self.run_iteration(self.val_gen, False) val_losses.append(l) self.all_val_losses_tr_mode.append(np.mean(val_losses)) self.print_to_log_file( "validation loss (train=True): %.4f" % self.all_val_losses_tr_mode[-1]) self.update_train_loss_MA( ) # needed for lr scheduler and stopping of training continue_training = self.on_epoch_end() epoch_end_time = time() if not continue_training: # allows for early stopping break self.epoch += 1 self.print_to_log_file("This epoch took %f s\n" % (epoch_end_time - epoch_start_time)) self.epoch -= 1 # if we don't do this we can get a problem with loading model_final_checkpoint. if self.save_final_checkpoint: self.save_checkpoint( join(self.output_folder, "model_final_checkpoint.model")) if self.local_rank == 0: # now we can delete latest as it will be identical with final if isfile(join(self.output_folder, "model_latest.model")): os.remove(join(self.output_folder, "model_latest.model")) if isfile(join(self.output_folder, "model_latest.model.pkl")): os.remove(join(self.output_folder, "model_latest.model.pkl")) net.do_ds = ds
def initialize(self, training=True, force_load_plans=False): """ :param training: :return: """ if not self.was_initialized: maybe_mkdir_p(self.output_folder) if force_load_plans or (self.plans is None): self.load_plans_file() self.process_plans(self.plans) self.setup_DA_params() self.folder_with_preprocessed_data = join( self.dataset_directory, self.plans['data_identifier'] + "_stage%d" % self.stage) if training: self.dl_tr, self.dl_val = self.get_basic_generators() if self.unpack_data: if self.local_rank == 0: print("unpacking dataset") unpack_dataset(self.folder_with_preprocessed_data) print("done") distributed.barrier() else: print( "INFO: Not unpacking data! Training may be slow due to that. Pray you are not using 2d or you " "will wait all winter for your model to finish!") # setting weights for deep supervision losses net_numpool = len(self.net_num_pool_op_kernel_sizes) # we give each output a weight which decreases exponentially (division by 2) as the resolution decreases # this gives higher resolution outputs more weight in the loss weights = np.array([1 / (2**i) for i in range(net_numpool)]) # we don't use the lowest 2 outputs. Normalize weights so that they sum to 1 mask = np.array([ True if i < net_numpool - 1 else False for i in range(net_numpool) ]) weights[~mask] = 0 weights = weights / weights.sum() self.ds_loss_weights = weights seeds_train = np.random.random_integers( 0, 99999, self.data_aug_params.get('num_threads')) seeds_val = np.random.random_integers( 0, 99999, max(self.data_aug_params.get('num_threads') // 2, 1)) print("seeds train", seeds_train) print("seeds_val", seeds_val) self.tr_gen, self.val_gen = get_moreDA_augmentation( self.dl_tr, self.dl_val, self.data_aug_params['patch_size_for_spatialtransform'], self.data_aug_params, deep_supervision_scales=self.deep_supervision_scales, seeds_train=seeds_train, seeds_val=seeds_val, pin_memory=self.pin_memory) self.print_to_log_file("TRAINING KEYS:\n %s" % (str(self.dataset_tr.keys())), also_print_to_console=False) self.print_to_log_file("VALIDATION KEYS:\n %s" % (str(self.dataset_val.keys())), also_print_to_console=False) else: pass self.initialize_network() self.initialize_optimizer_and_scheduler() self.network = DDP(self.network, device_ids=[self.local_rank]) else: self.print_to_log_file( 'self.was_initialized is True, not running self.initialize again' ) self.was_initialized = True
def initialize(self, training=True, force_load_plans=False): """ For prediction of test cases just set training=False, this will prevent loading of training data and training batchgenerator initialization :param training: :return: """ if not self.was_initialized: maybe_mkdir_p(self.output_folder) if force_load_plans or (self.plans is None): self.load_plans_file() self.process_plans(self.plans) self.setup_DA_params() self.folder_with_preprocessed_data = join( self.dataset_directory, self.plans['data_identifier'] + "_stage%d" % self.stage) if training: self.dl_tr, self.dl_val = self.get_basic_generators() if self.unpack_data: if self.local_rank == 0: print("unpacking dataset") unpack_dataset(self.folder_with_preprocessed_data) print("done") else: # we need to wait until worker 0 has finished unpacking npz_files = subfiles( self.folder_with_preprocessed_data, suffix=".npz", join=False) case_ids = [i[:-4] for i in npz_files] all_present = all([ isfile( join(self.folder_with_preprocessed_data, i + ".npy")) for i in case_ids ]) while not all_present: print("worker", self.local_rank, "is waiting for unpacking") sleep(3) all_present = all([ isfile( join(self.folder_with_preprocessed_data, i + ".npy")) for i in case_ids ]) # there is some slight chance that there may arise some error because dataloader are loading a file # that is still being written by worker 0. We ignore this for now an address it only if it becomes # relevant # (this can occur because while worker 0 writes the file is technically present so the other workers # will proceed and eventually try to read it) else: print( "INFO: Not unpacking data! Training may be slow due to that. Pray you are not using 2d or you " "will wait all winter for your model to finish!") # setting weights for deep supervision losses net_numpool = len(self.net_num_pool_op_kernel_sizes) # we give each output a weight which decreases exponentially (division by 2) as the resolution decreases # this gives higher resolution outputs more weight in the loss weights = np.array([1 / (2**i) for i in range(net_numpool)]) # we don't use the lowest 2 outputs. Normalize weights so that they sum to 1 mask = np.array([ True if i < net_numpool - 1 else False for i in range(net_numpool) ]) weights[~mask] = 0 weights = weights / weights.sum() self.ds_loss_weights = weights seeds_train = np.random.random_integers( 0, 99999, self.data_aug_params.get('num_threads')) seeds_val = np.random.random_integers( 0, 99999, max(self.data_aug_params.get('num_threads') // 2, 1)) print("seeds train", seeds_train) print("seeds_val", seeds_val) self.tr_gen, self.val_gen = get_moreDA_augmentation( self.dl_tr, self.dl_val, self.data_aug_params['patch_size_for_spatialtransform'], self.data_aug_params, deep_supervision_scales=self.deep_supervision_scales, seeds_train=seeds_train, seeds_val=seeds_val) self.print_to_log_file("TRAINING KEYS:\n %s" % (str(self.dataset_tr.keys())), also_print_to_console=False) self.print_to_log_file("VALIDATION KEYS:\n %s" % (str(self.dataset_val.keys())), also_print_to_console=False) else: pass self.initialize_network() self.initialize_optimizer_and_scheduler() self._maybe_init_amp() self.network = DDP(self.network) else: self.print_to_log_file( 'self.was_initialized is True, not running self.initialize again' ) self.was_initialized = True
default_data_identifier = 'nnUNet' default_trainer = "nnUNetTrainerV2" default_cascade_trainer = "nnUNetTrainerV2CascadeFullRes" """ PLEASE READ paths.md FOR INFORMATION TO HOW TO SET THIS UP """ base = os.environ['nnUNet_raw_data_base'] if "nnUNet_raw_data_base" in os.environ.keys() else None preprocessing_output_dir = os.environ['nnUNet_preprocessed'] if "nnUNet_preprocessed" in os.environ.keys() else None network_training_output_dir_base = os.path.join(os.environ['RESULTS_FOLDER']) if "RESULTS_FOLDER" in os.environ.keys() else None if base is not None: nnUNet_raw_data = join(base, "nnUNet_raw_data") nnUNet_cropped_data = join(base, "nnUNet_cropped_data") maybe_mkdir_p(nnUNet_raw_data) maybe_mkdir_p(nnUNet_cropped_data) else: print("nnUNet_raw_data_base is not defined and nnU-Net can only be used on data for which preprocessed files " "are already present on your system. nnU-Net cannot be used for experiment planning and preprocessing like " "this. If this is not intended, please read nnunet/paths.md for information on how to set this up properly.") nnUNet_cropped_data = nnUNet_raw_data = None if preprocessing_output_dir is not None: maybe_mkdir_p(preprocessing_output_dir) else: print("nnUNet_preprocessed is not defined and nnU-Net can not be used for preprocessing " "or training. If this is not intended, please read nnunet/pathy.md for information on how to set this up.") preprocessing_output_dir = None if network_training_output_dir_base is not None:
def initialize(self, training=True, force_load_plans=False): ''' Print keys to visdom and set number of epochs ''' timestamp = datetime.now() if self.usevisdom and training: try: self.plotter = get_plotter(self.model_name) self.plotter.plot_text( "Initialising this model: %s <br> on %d_%d_%d_%02.0d_%02.0d_%02.0d" % (self.model_name, timestamp.year, timestamp.month, timestamp.day, timestamp.hour, timestamp.minute, timestamp.second), plot_name="Welcome") except: print("Unable to connect to visdom.") #super().initialize(training, force_load_plans) ## ------- nnunettrainerv2 nodeepsupervision """ removed deep supervision :return: """ if not self.was_initialized: maybe_mkdir_p(self.output_folder) if force_load_plans or (self.plans is None): self.load_plans_file() self.process_plans(self.plans) self.setup_DA_params() self.folder_with_preprocessed_data = join( self.dataset_directory, self.plans['data_identifier'] + "_stage%d" % self.stage) if training: self.dl_tr, self.dl_val = self.get_basic_generators() if self.unpack_data: print("unpacking dataset") unpack_dataset(self.folder_with_preprocessed_data) print("done") else: print( "INFO: Not unpacking data! Training may be slow due to that. Pray you are not using 2d or you " "will wait all winter for your model to finish!") assert self.deep_supervision_scales is None self.tr_gen, self.val_gen = get_moreDA_augmentation( self.dl_tr, self.dl_val, self.data_aug_params['patch_size_for_spatialtransform'], self.data_aug_params, deep_supervision_scales=self.deep_supervision_scales, classes=None, pin_memory=self.pin_memory) self.print_to_log_file("TRAINING KEYS:\n %s" % (str(self.dataset_tr.keys())), also_print_to_console=False) self.print_to_log_file("VALIDATION KEYS:\n %s" % (str(self.dataset_val.keys())), also_print_to_console=False) else: pass self.initialize_network() self.initialize_optimizer_and_scheduler() assert isinstance(self.network, (SegmentationNetwork, nn.DataParallel)) else: self.print_to_log_file( 'self.was_initialized is True, not running self.initialize again' ) self.was_initialized = True # ----------------------- if self.freeze: self.initialize_optimizer_and_scheduler_freezing() if training and self.usevisdom: try: self.plotter.plot_text( "EPOCHS: %s <br> LEARNING RATE: %s <br> TRAINING KEYS: %s <br> VALIDATION KEYS: %s" % (str(self.max_num_epochs), str(self.initial_lr), str(self.dataset_tr.keys()), str( self.dataset_val.keys())), plot_name="Dataset_Info") except: print("Unable to connect to visdom.")
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 output_folder_name = join( network_training_output_dir, args.model, args.task_name, args.nnunet_trainer + "__" + args.plans_identifier) maybe_mkdir_p(output_folder_name) 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:
def training(unet, train_loader, val_loader, epochs, batch_size, device, fold_dir, DATASET): # Set goals for training to end if DATASET is 'DIC-C2DH-HeLa': when_to_stop = 0 goal = 0.7756 # IoU value from table 2 in Ronneberger et al. (2015) elif DATASET is 'ISBI2012': when_to_stop = 1 goal = 0.0611 # PE value from table 1 in Ronneberger et al. (2015) elif DATASET is 'PhC-C2DH-U373': when_to_stop = 2 goal = 0.9203 # IoU value from table 2 in Ronneberger et al. (2015) else: when_to_stop = None optimizer = optim.SGD(unet.parameters(), lr=0.0001, momentum=0.99) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=30, threshold=1e-3, threshold_mode='rel', eps=1e-7) my_patience = 0 maybe_mkdir_p(os.path.join(fold_dir, 'progress')) maybe_mkdir_p(os.path.join(fold_dir, 'models')) loss_best_epoch = 100000.0 for epoch in range(epochs + 1): print(' ') print('Epoch:', epoch) start = time() total_loss = 0 total_loss_val = 0 start_eval_train = 0 start_eval_val = 0 torch.cuda.empty_cache() for batch in train_loader: optimizer.zero_grad() images, labels = batch preds = unet(images.to(device)) # pass batch to the unet pad = int((preds.shape[-1] - labels.shape[-1]) / 2) preds = preds[:, :, pad:labels.shape[-1] + pad, pad:labels.shape[-1] + pad] ll = torch.empty_like(preds) ll[:, 0, :, :] = 1 - labels[:, 0, :, :] # background ll[:, 1, :, :] = labels[:, 0, :, :] # cell ll = ll.to(device) if DATASET is 'DIC-C2DH-HeLa': weight_maps = weighted_map(labels.squeeze(1)).to(device) criterion = nn.BCEWithLogitsLoss(weight=weight_maps) else: weight_maps = class_balance(labels.squeeze(1)).to(device) criterion = nn.BCEWithLogitsLoss(weight=weight_maps) loss = criterion(preds, ll) loss.backward() # compute the gradients using backprop optimizer.step() # update the weights total_loss += loss preds = preds.argmax(dim=1) for idx in range(preds.shape[0]): if start_eval_train == 0 and idx == 0: # First time in epoch we initialize train_eval train_eval = evaluation_metrics( preds[idx, :, :].detach(), labels[idx, 0, :, :].detach()) start_eval_train += 1 else: np.concatenate( (train_eval, evaluation_metrics(preds[idx, :, :].detach(), labels[idx, 0, :, :].detach())), axis=1) train_eval_epoch = np.mean(train_eval, axis=1) torch.cuda.empty_cache() with torch.no_grad(): for batch in val_loader: images, labels = batch preds = unet(images.to(device)) pad = int((preds.shape[-1] - labels.shape[-1]) / 2) preds = preds[:, :, pad:labels.shape[-1] + pad, pad:labels.shape[-1] + pad] ll = torch.empty_like(preds) ll[:, 0, :, :] = 1 - labels[:, 0, :, :] # background ll[:, 1, :, :] = labels[:, 0, :, :] # cell ll = ll.to(device) if DATASET is 'DIC-C2DH-HeLa': weight_maps = weighted_map(labels.squeeze(1)).to(device) criterion = nn.BCEWithLogitsLoss(weight=weight_maps) else: weight_maps = class_balance(labels.squeeze(1)).to(device) criterion = nn.BCEWithLogitsLoss(weight=weight_maps) loss = criterion(preds, ll) total_loss_val += loss preds = preds.argmax(1) for idx in range(preds.shape[0]): if start_eval_val == 0 and idx == 0: # First time in epoch we initialize val_eval val_eval = evaluation_metrics( preds[idx, :, :].detach(), labels[idx, 0, :, :].detach()) start_eval_val += 1 else: np.concatenate((val_eval, evaluation_metrics( preds[idx, :, :].detach(), labels[idx, 0, :, :].detach())), axis=1) val_eval_epoch = np.mean(val_eval, axis=1) scheduler.step(total_loss_val / (len(val_loader) * batch_size)) # update the lr for param_group in optimizer.param_groups: l_rate = param_group['lr'] loss_epoch = total_loss / (len(train_loader) * batch_size) loss_epoch_val = total_loss_val / (len(val_loader) * batch_size) if loss_epoch_val < (loss_best_epoch * (1.0 - scheduler.threshold)): loss_best_epoch = loss_epoch_val print('New best epoch!') my_patience = 0 PATH = os.path.join(fold_dir, 'models', 'unet_weight_save_best.pth') torch.save(unet.state_dict(), PATH) print('Model has been saved:') print(PATH) else: my_patience += 1 print('Current lr is: ', l_rate) print('Patience is: {}/{}'.format( my_patience, scheduler.patience)) print('Mean IoU training: ', "{:.6f}".format(train_eval_epoch[0])) print('Mean PE training: ', "{:.6f}".format(train_eval_epoch[1])) print('Mean IoU validation: ', "{:.6f}".format(val_eval_epoch[0])) print('Mean PE validation: ', "{:.6f}".format(val_eval_epoch[1])) print('Total training loss: ', "{:.6f}".format(loss_epoch.item())) print('Total validation loss: ', "{:.6f}".format(loss_epoch_val.item())) print('Best epoch validation loss:', "{:.6f}".format(loss_best_epoch.item())) print('Epoch duration: ', "{:.6f}".format(time() - start), 's') print( ' ' ) # Save progress (evaluation metrics and loss) if epoch == 0: train_eval_progress_iou = [train_eval_epoch[0]] train_eval_progress_pe = [train_eval_epoch[1]] val_eval_progress_iou = [val_eval_epoch[0]] val_eval_progress_pe = [val_eval_epoch[1]] loss_progress = [loss_epoch.item()] loss_progress_val = [loss_epoch_val.item()] elif epoch > 0: train_eval_progress_iou = np.concatenate( (train_eval_progress_iou, [train_eval_epoch[0]])) train_eval_progress_pe = np.concatenate( (train_eval_progress_pe, [train_eval_epoch[1]])) val_eval_progress_iou = np.concatenate( (val_eval_progress_iou, [val_eval_epoch[0]])) val_eval_progress_pe = np.concatenate( (val_eval_progress_pe, [val_eval_epoch[1]])) loss_progress = np.append(loss_progress, [loss_epoch.item()]) loss_progress_val = np.append(loss_progress_val, [loss_epoch_val.item()]) np.savetxt(os.path.join(fold_dir, 'progress', 'train_eval_iou.out'), train_eval_progress_iou) np.savetxt(os.path.join(fold_dir, 'progress', 'train_eval_pe.out'), train_eval_progress_pe) np.savetxt(os.path.join(fold_dir, 'progress', 'val_eval_iou.out'), val_eval_progress_iou) np.savetxt(os.path.join(fold_dir, 'progress', 'val_eval_pe.out'), val_eval_progress_pe) np.savetxt(os.path.join(fold_dir, 'progress', 'loss.out'), loss_progress) np.savetxt(os.path.join(fold_dir, 'progress', 'loss_val.out'), loss_progress_val) if when_to_stop == 0: if val_eval_epoch[0] > goal: PATH = os.path.join(fold_dir, 'models', 'unet_weight_save_{}.pth'.format(DATASET)) torch.save(unet.state_dict(), PATH) print('The goal was reached in epoch {}!'.format(epoch)) print('Model has been saved:') print(PATH) # break when_to_stop = None continue elif when_to_stop == 1: if val_eval_epoch[0] > goal: PATH = os.path.join(fold_dir, 'models', 'unet_weight_save_{}.pth'.format(DATASET)) torch.save(unet.state_dict(), PATH) print('The goal was reached in epoch {}!'.format(epoch)) print('Model has been saved:') print(PATH) # break when_to_stop = None continue elif when_to_stop == 2: if val_eval_epoch[0] > goal: PATH = os.path.join(fold_dir, 'models', 'unet_weight_save_{}.pth'.format(DATASET)) torch.save(unet.state_dict(), PATH) print('The goal was reached in epoch {}!'.format(epoch)) print('Model has been saved:') print(PATH) # break when_to_stop = None continue # Save model every 50 epochs if epoch % 25 == 0: PATH = os.path.join(fold_dir, 'models', 'unet_weight_save_latest.pth') torch.save(unet.state_dict(), PATH) print('Model has been saved:') print(PATH) if l_rate < 10 * scheduler.eps and my_patience == scheduler.patience: print(f'LR dropped below {10 * scheduler.eps}!') print('Stopping training') print(' ') PATH = os.path.join(fold_dir, 'models', 'unet_weight_save_latest.pth') torch.save(unet.state_dict(), PATH) print('Model has been saved:') print(PATH) break if my_patience == scheduler.patience: my_patience = -1 print('Training is finished as epoch {} has been reached'.format(epoch)) print(' ')
# -*- 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)
# do not modify these unless you know what you are doing my_output_identifier = "nnUNet" default_plans_identifier = "nnUNetPlans" default_data_identifier = 'nnUNet' try: # base is the folder where the raw data is stored. You just need to set base only, the others will be created # automatically (they are subfolders of base). # Here I use environment variables to set the base folder. Environment variables allow me to use the same code on # different systems (and our compute cluster). You can replace this line with something like: # base = "/path/to/my/folder" base = os.environ['nnUNet_base'] raw_dataset_dir = join(base, "nnUNet_raw") splitted_4d_output_dir = join(base, "nnUNet_raw_splitted") cropped_output_dir = join(base, "nnUNet_raw_cropped") maybe_mkdir_p(splitted_4d_output_dir) maybe_mkdir_p(raw_dataset_dir) maybe_mkdir_p(cropped_output_dir) except KeyError: cropped_output_dir = splitted_4d_output_dir = raw_dataset_dir = base = None # preprocessing_output_dir is where the preprocessed data is stored. If you run a training I very strongly recommend # this is a SSD! try: # Here I use environment variables to set the folder. Environment variables allow me to use the same code on # different systems (and our compute cluster). You can replace this line with something like: # preprocessing_output_dir = "/path/to/my/folder_with_preprocessed_data" preprocessing_output_dir = os.environ['nnUNet_preprocessed'] except KeyError: preprocessing_output_dir = None
[join(nnunet.__path__[0], "training", "network_training")], trainerclass, "nnunet.training.network_training") if trainer_class is None: raise RuntimeError( "Could not find trainer class in nnunet.training.network_training") else: assert issubclass( trainer_class, nnUNetTrainer ), "network_trainer was found but is not derived from nnUNetTrainer" trainer = trainer_class(plans_file, fold, folder_with_preprocessed_data, output_folder=output_folder_name, dataset_directory=dataset_directory, batch_dice=batch_dice, stage=stage) trainer.initialize(False) trainer.load_dataset() trainer.do_split() trainer.load_best_checkpoint(train=False) stage_to_be_predicted_folder = join( dataset_directory, trainer.plans['data_identifier'] + "_stage%d" % 1) output_folder = join(pardir(trainer.output_folder), "pred_next_stage") maybe_mkdir_p(output_folder) predict_next_stage(trainer, stage_to_be_predicted_folder)