From 43349fa5f0680e8109a78dca7215c19e258c9dd7 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Tue, 3 Dec 2024 09:49:53 +0100 Subject: [PATCH] add predict_from_files_sequential to nnUNetPredictor --- nnunetv2/inference/predict_from_raw_data.py | 140 +++++++++++++++----- 1 file changed, 106 insertions(+), 34 deletions(-) diff --git a/nnunetv2/inference/predict_from_raw_data.py b/nnunetv2/inference/predict_from_raw_data.py index 1f5ede64f..a900ef1d1 100644 --- a/nnunetv2/inference/predict_from_raw_data.py +++ b/nnunetv2/inference/predict_from_raw_data.py @@ -658,6 +658,92 @@ def predict_sliding_window_return_logits(self, input_image: torch.Tensor) \ predicted_logits = predicted_logits[(slice(None), *slicer_revert_padding[1:])] return predicted_logits + def predict_from_files_sequential(self, + list_of_lists_or_source_folder: Union[str, List[List[str]]], + output_folder_or_list_of_truncated_output_files: Union[str, None, List[str]], + save_probabilities: bool = False, + overwrite: bool = True, + folder_with_segs_from_prev_stage: str = None): + """ + Just like predict_from_files but doesn't use any multiprocessing. Slow, but sometimes necessary + """ + if isinstance(output_folder_or_list_of_truncated_output_files, str): + output_folder = output_folder_or_list_of_truncated_output_files + elif isinstance(output_folder_or_list_of_truncated_output_files, list): + output_folder = os.path.dirname(output_folder_or_list_of_truncated_output_files[0]) + else: + output_folder = None + + ######################## + # let's store the input arguments so that its clear what was used to generate the prediction + if output_folder is not None: + my_init_kwargs = {} + for k in inspect.signature(self.predict_from_files_sequential).parameters.keys(): + my_init_kwargs[k] = locals()[k] + my_init_kwargs = deepcopy( + my_init_kwargs) # let's not unintentionally change anything in-place. Take this as a + recursive_fix_for_json_export(my_init_kwargs) + maybe_mkdir_p(output_folder) + save_json(my_init_kwargs, join(output_folder, 'predict_from_raw_data_args.json')) + + # we need these two if we want to do things with the predictions like for example apply postprocessing + save_json(self.dataset_json, join(output_folder, 'dataset.json'), sort_keys=False) + save_json(self.plans_manager.plans, join(output_folder, 'plans.json'), sort_keys=False) + ####################### + + # check if we need a prediction from the previous stage + if self.configuration_manager.previous_stage_name is not None: + assert folder_with_segs_from_prev_stage is not None, \ + f'The requested configuration is a cascaded network. It requires the segmentations of the previous ' \ + f'stage ({self.configuration_manager.previous_stage_name}) as input. Please provide the folder where' \ + f' they are located via folder_with_segs_from_prev_stage' + + # sort out input and output filenames + list_of_lists_or_source_folder, output_filename_truncated, seg_from_prev_stage_files = \ + self._manage_input_and_output_lists(list_of_lists_or_source_folder, + output_folder_or_list_of_truncated_output_files, + folder_with_segs_from_prev_stage, overwrite, 0, 1, + save_probabilities) + if len(list_of_lists_or_source_folder) == 0: + return + + label_manager = self.plans_manager.get_label_manager(self.dataset_json) + preprocessor = self.configuration_manager.preprocessor_class(verbose=self.verbose) + + if output_filename_truncated is None: + output_filename_truncated = [None] * len(list_of_lists_or_source_folder) + if seg_from_prev_stage_files is None: + seg_from_prev_stage_files = [None] * len(seg_from_prev_stage_files) + + ret = [] + for li, of, sps in zip(list_of_lists_or_source_folder, output_filename_truncated, seg_from_prev_stage_files): + data, seg, data_properties = preprocessor.run_case( + li, + sps, + self.plans_manager, + self.configuration_manager, + self.dataset_json + ) + + print(f'perform_everything_on_device: {self.perform_everything_on_device}') + + prediction = self.predict_logits_from_preprocessed_data(torch.from_numpy(data)).cpu() + + if of is not None: + export_prediction_from_logits(prediction, data_properties, self.configuration_manager, self.plans_manager, + self.dataset_json, of, save_probabilities) + else: + ret.append(convert_predicted_logits_to_segmentation_with_correct_shape(prediction, self.plans_manager, + self.configuration_manager, self.label_manager, + data_properties, + save_probabilities)) + + # clear lru cache + compute_gaussian.cache_clear() + # clear device cache + empty_cache(self.device) + return ret + def predict_entry_point_modelfolder(): import argparse @@ -891,7 +977,7 @@ def predict_entry_point(): if __name__ == '__main__': - # predict a bunch of files + ########################## predict a bunch of files from nnunetv2.paths import nnUNet_results, nnUNet_raw predictor = nnUNetPredictor( @@ -905,42 +991,28 @@ def predict_entry_point(): allow_tqdm=True ) predictor.initialize_from_trained_model_folder( - join(nnUNet_results, 'Dataset003_Liver/nnUNetTrainer__nnUNetPlans__3d_lowres'), + join(nnUNet_results, 'Dataset004_Hippocampus/nnUNetTrainer_5epochs__nnUNetPlans__3d_fullres'), use_folds=(0,), checkpoint_name='checkpoint_final.pth', ) - predictor.predict_from_files(join(nnUNet_raw, 'Dataset003_Liver/imagesTs'), - join(nnUNet_raw, 'Dataset003_Liver/imagesTs_predlowres'), - save_probabilities=False, overwrite=False, - num_processes_preprocessing=2, num_processes_segmentation_export=2, - folder_with_segs_from_prev_stage=None, num_parts=1, part_id=0) - - # predict a numpy array - from nnunetv2.imageio.simpleitk_reader_writer import SimpleITKIO - - img, props = SimpleITKIO().read_images([join(nnUNet_raw, 'Dataset003_Liver/imagesTr/liver_63_0000.nii.gz')]) - ret = predictor.predict_single_npy_array(img, props, None, None, False) - - iterator = predictor.get_data_iterator_from_raw_npy_data([img], None, [props], None, 1) - ret = predictor.predict_from_data_iterator(iterator, False, 1) - - # predictor = nnUNetPredictor( - # tile_step_size=0.5, - # use_gaussian=True, - # use_mirroring=True, - # perform_everything_on_device=True, - # device=torch.device('cuda', 0), - # verbose=False, - # allow_tqdm=True - # ) - # predictor.initialize_from_trained_model_folder( - # join(nnUNet_results, 'Dataset003_Liver/nnUNetTrainer__nnUNetPlans__3d_cascade_fullres'), - # use_folds=(0,), - # checkpoint_name='checkpoint_final.pth', - # ) # predictor.predict_from_files(join(nnUNet_raw, 'Dataset003_Liver/imagesTs'), - # join(nnUNet_raw, 'Dataset003_Liver/imagesTs_predCascade'), + # join(nnUNet_raw, 'Dataset003_Liver/imagesTs_predlowres'), # save_probabilities=False, overwrite=False, # num_processes_preprocessing=2, num_processes_segmentation_export=2, - # folder_with_segs_from_prev_stage='/media/isensee/data/nnUNet_raw/Dataset003_Liver/imagesTs_predlowres', - # num_parts=1, part_id=0) + # folder_with_segs_from_prev_stage=None, num_parts=1, part_id=0) + # + # # predict a numpy array + # from nnunetv2.imageio.simpleitk_reader_writer import SimpleITKIO + # + # img, props = SimpleITKIO().read_images([join(nnUNet_raw, 'Dataset003_Liver/imagesTr/liver_63_0000.nii.gz')]) + # ret = predictor.predict_single_npy_array(img, props, None, None, False) + # + # iterator = predictor.get_data_iterator_from_raw_npy_data([img], None, [props], None, 1) + # ret = predictor.predict_from_data_iterator(iterator, False, 1) + + ret = predictor.predict_from_files_sequential( + [['/media/isensee/raw_data/nnUNet_raw/Dataset004_Hippocampus/imagesTs/hippocampus_002_0000.nii.gz'], ['/media/isensee/raw_data/nnUNet_raw/Dataset004_Hippocampus/imagesTs/hippocampus_005_0000.nii.gz']], + '/home/isensee/temp/tmp', False, True, None + ) + +