From 5db96042779fe720dc6cef7ba4b32d2f9d127d31 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Fri, 26 Apr 2024 12:04:06 +0200 Subject: [PATCH] add error message when trainer is not found in inference --- nnunetv2/inference/predict_from_raw_data.py | 4 +++- nnunetv2/utilities/find_class_by_name.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/nnunetv2/inference/predict_from_raw_data.py b/nnunetv2/inference/predict_from_raw_data.py index c28b41e1b..8d4096482 100644 --- a/nnunetv2/inference/predict_from_raw_data.py +++ b/nnunetv2/inference/predict_from_raw_data.py @@ -96,7 +96,9 @@ def initialize_from_trained_model_folder(self, model_training_output_dir: str, num_input_channels = determine_num_input_channels(plans_manager, configuration_manager, dataset_json) trainer_class = recursive_find_python_class(join(nnunetv2.__path__[0], "training", "nnUNetTrainer"), trainer_name, 'nnunetv2.training.nnUNetTrainer') - + if trainer_class is None: + raise RuntimeError(f'Unable to locate trainer class {trainer_name} in nnunetv2.training.nnUNetTrainer. ' + f'Please place it there (in any .py file)!') network = trainer_class.build_network_architecture( configuration_manager.network_arch_class_name, configuration_manager.network_arch_init_kwargs, diff --git a/nnunetv2/utilities/find_class_by_name.py b/nnunetv2/utilities/find_class_by_name.py index a345d99a7..223b3acc3 100644 --- a/nnunetv2/utilities/find_class_by_name.py +++ b/nnunetv2/utilities/find_class_by_name.py @@ -21,4 +21,4 @@ def recursive_find_python_class(folder: str, class_name: str, current_module: st tr = recursive_find_python_class(join(folder, modname), class_name, current_module=next_current_module) if tr is not None: break - return tr \ No newline at end of file + return tr