diff --git a/docs/source/apps.rst b/docs/source/apps.rst index 495ea06530..32c41d9ed2 100644 --- a/docs/source/apps.rst +++ b/docs/source/apps.rst @@ -297,3 +297,4 @@ FastMRIReader .. autofunction:: monai.apps.nnunet.convert_nnunet_to_monai_bundle .. autofunction:: monai.apps.nnunet.convert_monai_bundle_to_nnunet .. autofunction:: monai.apps.nnunet.get_network_from_nnunet_plans +.. autodata:: monai.apps.nnunet.DEFAULT_NNUNET_NETWORK_MODULES diff --git a/monai/apps/nnunet/__init__.py b/monai/apps/nnunet/__init__.py index 780ee7a861..5fa7af6deb 100644 --- a/monai/apps/nnunet/__init__.py +++ b/monai/apps/nnunet/__init__.py @@ -12,6 +12,7 @@ from __future__ import annotations from .nnunet_bundle import ( + DEFAULT_NNUNET_NETWORK_MODULES, ModelnnUNetWrapper, convert_monai_bundle_to_nnunet, convert_nnunet_to_monai_bundle, diff --git a/monai/apps/nnunet/nnunet_bundle.py b/monai/apps/nnunet/nnunet_bundle.py index 47a0755ddf..d71a2aa38f 100644 --- a/monai/apps/nnunet/nnunet_bundle.py +++ b/monai/apps/nnunet/nnunet_bundle.py @@ -12,7 +12,9 @@ import os import shutil +from collections.abc import Sequence from pathlib import Path +from pydoc import locate from typing import Any import numpy as np @@ -32,8 +34,60 @@ "convert_nnunet_to_monai_bundle", "convert_monai_bundle_to_nnunet", "ModelnnUNetWrapper", + "DEFAULT_NNUNET_NETWORK_MODULES", ] +#: Modules that nnUNet plans may reference for the network class and its imported keyword arguments. +DEFAULT_NNUNET_NETWORK_MODULES: tuple[str, ...] = ("dynamic_network_architectures", "torch.nn") + + +def _in_allowed_modules(dotted_name: str, allowed_modules: Sequence[str]) -> bool: + return any(dotted_name == m or dotted_name.startswith(m + ".") for m in allowed_modules) + + +def _check_nnunet_network_spec( + class_name: Any, init_kwargs: Any, kwargs_requiring_import: Any, allowed_modules: Sequence[str] | None = None +) -> None: + """ + Check the network specification read from an nnUNet ``plans.json`` before nnUNet resolves and builds it. + + nnUNet imports ``network_class_name`` and every keyword argument listed in ``_kw_requires_import`` by dotted + name and calls them, so each of these names must be under one of ``allowed_modules`` and, if it can be + imported, must refer to a ``torch.nn.Module`` subclass. The module prefix is checked before anything is + imported. + + Args: + class_name: the ``network_class_name`` from the plans. + init_kwargs: the ``arch_kwargs`` from the plans. + kwargs_requiring_import: the ``_kw_requires_import`` names from the plans. + allowed_modules: dotted module prefixes that may be referenced. Defaults to + ``DEFAULT_NNUNET_NETWORK_MODULES``. Add the module of a custom architecture here to load it. + + Raises: + ValueError: if a name is outside ``allowed_modules`` or does not resolve to a ``torch.nn.Module`` subclass. + """ + allowed = tuple(DEFAULT_NNUNET_NETWORK_MODULES if allowed_modules is None else allowed_modules) + hint = "Pass the module in `allowed_network_modules` if this architecture is trusted." + + def _check(name: Any, what: str) -> None: + if not isinstance(name, str) or not _in_allowed_modules(name, allowed): + raise ValueError(f"Refusing to load {what} {name!r} from nnUNet plans: not in {list(allowed)}. {hint}") + resolved = locate(name) + # an unresolved name is left to nnUNet, which only searches inside dynamic_network_architectures for it + if resolved is not None and not (isinstance(resolved, type) and issubclass(resolved, torch.nn.Module)): + raise ValueError(f"Refusing to load {what} {name!r} from nnUNet plans: not a torch.nn.Module subclass.") + + _check(class_name, "network class") + init_kwargs = init_kwargs or {} + if not isinstance(init_kwargs, dict): + raise ValueError( + f"Invalid network arguments in nnUNet plans: expected a dict, got {type(init_kwargs).__name__}." + ) + for key in kwargs_requiring_import or (): + value = init_kwargs.get(key) + if value is not None: + _check(value, f"network argument {key!r}") + def get_nnunet_trainer( dataset_name_or_id: str | int, @@ -152,6 +206,9 @@ class ModelnnUNetWrapper(torch.nn.Module): The folder path where the model and related files are stored. model_name : str, optional The name of the model file, by default "model.pt". + allowed_network_modules : Sequence[str], optional + Module prefixes that the plans may reference for the network class and imported network arguments, + by default ``DEFAULT_NNUNET_NETWORK_MODULES``. Attributes ---------- @@ -166,7 +223,13 @@ class ModelnnUNetWrapper(torch.nn.Module): restoring network architecture, and setting up the predictor for inference. """ - def __init__(self, predictor: object, model_folder: str | Path, model_name: str = "model.pt"): # type: ignore + def __init__( + self, + predictor: object, + model_folder: str | Path, + model_name: str = "model.pt", + allowed_network_modules: Sequence[str] | None = None, + ): # type: ignore super().__init__() self.predictor = predictor @@ -216,6 +279,12 @@ def __init__(self, predictor: object, model_folder: str | Path, model_name: str f"Unable to locate trainer class {trainer_name} in nnunetv2.training.nnUNetTrainer. " f"Please place it there (in any .py file)!" ) + _check_nnunet_network_spec( + configuration_manager.network_arch_class_name, + configuration_manager.network_arch_init_kwargs, + configuration_manager.network_arch_init_kwargs_req_import, + allowed_network_modules, + ) network = trainer_class.build_network_architecture( configuration_manager.network_arch_class_name, configuration_manager.network_arch_init_kwargs, @@ -294,7 +363,9 @@ def forward(self, x: MetaTensor) -> MetaTensor: return MetaTensor(out_tensor, meta=x.meta) -def get_nnunet_monai_predictor(model_folder: str | Path, model_name: str = "model.pt") -> ModelnnUNetWrapper: +def get_nnunet_monai_predictor( + model_folder: str | Path, model_name: str = "model.pt", allowed_network_modules: Sequence[str] | None = None +) -> ModelnnUNetWrapper: """ Initializes and returns a `nnUNetMONAIModelWrapper` containing the corresponding `nnUNetPredictor`. The model folder should contain the following files, created during training: @@ -325,6 +396,9 @@ def get_nnunet_monai_predictor(model_folder: str | Path, model_name: str = "mode The folder where the model is stored. model_name : str, optional The name of the model file, by default "model.pt". + allowed_network_modules : Sequence[str], optional + Module prefixes that ``plans.json`` may reference for the network class and imported network arguments, + by default ``DEFAULT_NNUNET_NETWORK_MODULES``. Extend it to load a trusted custom architecture. Returns ------- @@ -344,7 +418,7 @@ def get_nnunet_monai_predictor(model_folder: str | Path, model_name: str = "mode allow_tqdm=True, ) # initializes the network architecture, loads the checkpoint - wrapper = ModelnnUNetWrapper(predictor, model_folder, model_name) + wrapper = ModelnnUNetWrapper(predictor, model_folder, model_name, allowed_network_modules) return wrapper @@ -428,6 +502,7 @@ def get_network_from_nnunet_plans( configuration: str, model_ckpt: str | None = None, model_key_in_ckpt: str = "model", + allowed_network_modules: Sequence[str] | None = None, ) -> torch.nn.Module | Any: """ Load and initialize a nnUNet network based on nnUNet plans and configuration. @@ -444,6 +519,9 @@ def get_network_from_nnunet_plans( Path to the model checkpoint file. If None, the network is returned without loading weights (default is None). model_key_in_ckpt : str, optional The key in the checkpoint file that contains the model state dictionary (default is "model"). + allowed_network_modules : Sequence[str], optional + Module prefixes that ``plans_file`` may reference for the network class and imported network arguments, + by default ``DEFAULT_NNUNET_NETWORK_MODULES``. Extend it to load a trusted custom architecture. Returns ------- @@ -465,6 +543,12 @@ def get_network_from_nnunet_plans( enable_deep_supervision = True + _check_nnunet_network_spec( + configuration_manager.network_arch_class_name, + configuration_manager.network_arch_init_kwargs, + configuration_manager.network_arch_init_kwargs_req_import, + allowed_network_modules, + ) network = get_network_from_plans( configuration_manager.network_arch_class_name, configuration_manager.network_arch_init_kwargs, diff --git a/tests/apps/nnunet/test_nnunet_bundle_network_spec.py b/tests/apps/nnunet/test_nnunet_bundle_network_spec.py new file mode 100644 index 0000000000..de0cb2eab2 --- /dev/null +++ b/tests/apps/nnunet/test_nnunet_bundle_network_spec.py @@ -0,0 +1,266 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import json +import os +import tempfile +import unittest +from unittest import mock + +from monai.apps.nnunet.nnunet_bundle import _check_nnunet_network_spec, get_network_from_nnunet_plans +from monai.utils import optional_import + +_, has_nnunet = optional_import("nnunetv2") + +UNET = "dynamic_network_architectures.architectures.unet.PlainConvUNet" +IMPORTED_KWARGS = ["conv_op", "norm_op", "dropout_op", "nonlin"] + + +def _arch_kwargs(**overrides) -> dict: + kwargs = { + "n_stages": 2, + "features_per_stage": [4, 8], + "conv_op": "torch.nn.modules.conv.Conv2d", + "kernel_sizes": [[3, 3], [3, 3]], + "strides": [[1, 1], [2, 2]], + "n_conv_per_stage": [1, 1], + "n_conv_per_stage_decoder": [1], + "conv_bias": True, + "norm_op": "torch.nn.modules.instancenorm.InstanceNorm2d", + "norm_op_kwargs": {"eps": 1e-5, "affine": True}, + "dropout_op": None, + "dropout_op_kwargs": None, + "nonlin": "torch.nn.LeakyReLU", + "nonlin_kwargs": {"inplace": True}, + } + kwargs.update(overrides) + return kwargs + + +class TestNetworkSpecCheck(unittest.TestCase): + def test_default_spec_is_accepted(self): + with mock.patch("monai.apps.nnunet.nnunet_bundle.locate", return_value=None): + _check_nnunet_network_spec(UNET, _arch_kwargs(), IMPORTED_KWARGS) + + def test_torch_nn_modules_are_accepted(self): + _check_nnunet_network_spec("torch.nn.Identity", _arch_kwargs(), IMPORTED_KWARGS) + + def test_class_outside_allowed_modules_is_rejected(self): + for name in ( + "numpy.savez", + "os.system", + "builtins.eval", + "torch.nnx.Thing", + "dynamic_network_architecturesx.A", + ): + with self.subTest(name=name): + with mock.patch("monai.apps.nnunet.nnunet_bundle.locate") as locate: + with self.assertRaisesRegex(ValueError, "not in"): + _check_nnunet_network_spec(name, _arch_kwargs(), IMPORTED_KWARGS) + locate.assert_not_called() + + def test_imported_argument_outside_allowed_modules_is_rejected(self): + for key, value in (("nonlin", "subprocess.run"), ("conv_op", "os.system"), ("norm_op", 3)): + with self.subTest(key=key): + with self.assertRaisesRegex(ValueError, repr(key)): + _check_nnunet_network_spec("torch.nn.Identity", _arch_kwargs(**{key: value}), IMPORTED_KWARGS) + + def test_allowed_name_that_is_not_a_module_class_is_rejected(self): + for name in ("torch.nn.functional.relu", "torch.nn.init.zeros_", "torch.nn.utils"): + with self.subTest(name=name): + with self.assertRaisesRegex(ValueError, "torch.nn.Module"): + _check_nnunet_network_spec(name, {}, []) + + def test_non_string_class_name_is_rejected(self): + with self.assertRaises(ValueError): + _check_nnunet_network_spec(None, {}, []) + + def test_custom_modules_can_be_allowed(self): + with self.assertRaises(ValueError): + _check_nnunet_network_spec("monai.networks.nets.BasicUNet", {}, []) + _check_nnunet_network_spec("monai.networks.nets.BasicUNet", {}, [], allowed_modules=["monai.networks"]) + + +def _plans(class_name: str = UNET, **kwarg_overrides) -> dict: + resampling = "resample_data_or_seg_to_shape" + return { + "dataset_name": "Dataset001_Test", + "plans_name": "nnUNetPlans", + "original_median_spacing_after_transp": [1, 1], + "original_median_shape_after_transp": [8, 8], + "image_reader_writer": "SimpleITKIO", + "transpose_forward": [0, 1], + "transpose_backward": [0, 1], + "experiment_planner_used": "ExperimentPlanner", + "label_manager": "LabelManager", + "foreground_intensity_properties_per_channel": {}, + "configurations": { + "2d": { + "data_identifier": "nnUNetPlans_2d", + "preprocessor_name": "DefaultPreprocessor", + "batch_size": 1, + "patch_size": [8, 8], + "median_image_size_in_voxels": [8, 8], + "spacing": [1, 1], + "normalization_schemes": ["ZScoreNormalization"], + "use_mask_for_norm": [False], + "resampling_fn_data": resampling, + "resampling_fn_seg": resampling, + "resampling_fn_probabilities": resampling, + "resampling_fn_data_kwargs": {}, + "resampling_fn_seg_kwargs": {}, + "resampling_fn_probabilities_kwargs": {}, + "batch_dice": True, + "architecture": { + "network_class_name": class_name, + "arch_kwargs": _arch_kwargs(**kwarg_overrides), + "_kw_requires_import": IMPORTED_KWARGS, + }, + } + }, + } + + +@unittest.skipUnless(has_nnunet, "Requires nnunetv2") +class TestGetNetworkFromPlans(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.plans_file = os.path.join(self._tmp.name, "plans.json") + self.dataset_file = os.path.join(self._tmp.name, "dataset.json") + with open(self.dataset_file, "w") as f: + json.dump({"channel_names": {"0": "CT"}, "labels": {"background": 0, "a": 1}, "file_ending": ".nii.gz"}, f) + + def tearDown(self): + self._tmp.cleanup() + + def _load(self, plans: dict): + with open(self.plans_file, "w") as f: + json.dump(plans, f) + return get_network_from_nnunet_plans(self.plans_file, self.dataset_file, "2d") + + def test_default_plans_build_the_network(self): + self.assertEqual(type(self._load(_plans())).__name__, "PlainConvUNet") + + def test_network_class_outside_allowed_modules_is_not_called(self): + output = os.path.join(self._tmp.name, "out.npz") + with self.assertRaisesRegex(ValueError, "not in"): + self._load(_plans(class_name="numpy.savez", file=output)) + self.assertFalse(os.path.exists(output)) + + def test_imported_argument_outside_allowed_modules_is_not_called(self): + with mock.patch("os.system") as system: + with self.assertRaisesRegex(ValueError, "'nonlin'"): + self._load(_plans(nonlin="os.system")) + system.assert_not_called() + + +@unittest.skipUnless(has_nnunet, "Requires nnunetv2") +class TestModelWrapperNetworkSpec(unittest.TestCase): + """``ModelnnUNetWrapper`` and ``get_nnunet_monai_predictor`` apply the same check before building the network.""" + + def setUp(self): + """Create a model folder with ``dataset.json`` and a minimal ``nnunet_checkpoint.pth``.""" + import torch + + self._tmp = tempfile.TemporaryDirectory() + self.model_folder = os.path.join(self._tmp.name, "fold_0") + os.makedirs(self.model_folder) + with open(os.path.join(self._tmp.name, "dataset.json"), "w") as f: + json.dump({"channel_names": {"0": "CT"}, "labels": {"background": 0, "a": 1}, "file_ending": ".nii.gz"}, f) + torch.save( + {"trainer_name": "nnUNetTrainer", "init_args": {"configuration": "2d"}}, + os.path.join(self._tmp.name, "nnunet_checkpoint.pth"), + ) + + def tearDown(self): + """Remove the model folder.""" + self._tmp.cleanup() + + def _wrap(self, plans: dict, **kwargs): + """ + Write ``plans`` next to the model folder and build a ``ModelnnUNetWrapper`` from it. + + Args: + plans: the content of ``plans.json``. + kwargs: extra arguments for ``ModelnnUNetWrapper``. + + Returns: + the wrapper. + + Raises: + ValueError: if the plans reference a name outside the allowed modules. + """ + from monai.apps.nnunet.nnunet_bundle import ModelnnUNetWrapper + + with open(os.path.join(self._tmp.name, "plans.json"), "w") as f: + json.dump(plans, f) + return ModelnnUNetWrapper(mock.MagicMock(), self.model_folder, **kwargs) + + def test_default_plans_pass_the_check(self): + """Standard plans reach the trainer's network builder.""" + import torch + from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer + + # mocked because the trainer's build_network_architecture signature differs across nnunetv2 releases + with mock.patch.object(nnUNetTrainer, "build_network_architecture", return_value=torch.nn.Identity()) as build: + self._wrap(_plans()) + build.assert_called_once() + + def test_disallowed_network_is_not_built(self): + """A network class outside the allowed modules is rejected before the trainer builds anything.""" + from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer + + output = os.path.join(self._tmp.name, "out.npz") + with mock.patch.object(nnUNetTrainer, "build_network_architecture") as build: + with self.assertRaisesRegex(ValueError, "not in"): + self._wrap(_plans(class_name="numpy.savez", file=output)) + build.assert_not_called() + self.assertFalse(os.path.exists(output)) + + def test_disallowed_imported_argument_is_not_built(self): + """An imported network argument outside the allowed modules is rejected before the trainer builds anything.""" + from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer + + with mock.patch("os.system") as system, mock.patch.object(nnUNetTrainer, "build_network_architecture") as build: + with self.assertRaisesRegex(ValueError, "'nonlin'"): + self._wrap(_plans(nonlin="os.system")) + build.assert_not_called() + system.assert_not_called() + + def test_custom_modules_are_forwarded(self): + """A custom prefix is accepted only when passed in ``allowed_network_modules``.""" + import torch + from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer + + plans = _plans(class_name="monai.networks.nets.BasicUNet") + with mock.patch.object(nnUNetTrainer, "build_network_architecture", return_value=torch.nn.Identity()) as build: + with self.assertRaisesRegex(ValueError, "not in"): + self._wrap(plans) + build.assert_not_called() + self._wrap(plans, allowed_network_modules=["monai.networks", "torch.nn"]) + build.assert_called_once() + + def test_predictor_factory_forwards_allowed_modules(self): + """``get_nnunet_monai_predictor`` passes ``allowed_network_modules`` on to the wrapper.""" + from monai.apps.nnunet import nnunet_bundle + + with ( + mock.patch("nnunetv2.inference.predict_from_raw_data.nnUNetPredictor"), + mock.patch.object(nnunet_bundle, "ModelnnUNetWrapper") as wrapper, + ): + nnunet_bundle.get_nnunet_monai_predictor(self.model_folder, allowed_network_modules=["custom_nets"]) + self.assertEqual(wrapper.call_args.args[1:], (self.model_folder, "model.pt", ["custom_nets"])) + + +if __name__ == "__main__": + unittest.main()