diff --git a/monai/bundle/__init__.py b/monai/bundle/__init__.py index f30ee9c40c..62e754e135 100644 --- a/monai/bundle/__init__.py +++ b/monai/bundle/__init__.py @@ -13,4 +13,4 @@ from .config_parser import ConfigParser from .reference_resolver import ReferenceResolver from .scripts import ckpt_export, download, load, run, verify_metadata, verify_net_in_out -from .utils import EXPR_KEY, ID_REF_KEY, ID_SEP_KEY, MACRO_KEY +from .utils import EXPR_KEY, ID_REF_KEY, ID_SEP_KEY, MACRO_KEY, load_bundle_config diff --git a/monai/bundle/__main__.py b/monai/bundle/__main__.py index 3e3534ef74..ace3701d19 100644 --- a/monai/bundle/__main__.py +++ b/monai/bundle/__main__.py @@ -10,7 +10,7 @@ # limitations under the License. -from monai.bundle.scripts import ckpt_export, download, run, verify_metadata, verify_net_in_out +from monai.bundle.scripts import ckpt_export, download, init_bundle, run, verify_metadata, verify_net_in_out if __name__ == "__main__": from monai.utils import optional_import diff --git a/monai/bundle/config_parser.py b/monai/bundle/config_parser.py index 448b5ecf20..edbe4d1923 100644 --- a/monai/bundle/config_parser.py +++ b/monai/bundle/config_parser.py @@ -183,6 +183,19 @@ def set(self, config: Any, id: str = ""): """ self[id] = config + def __contains__(self, id: Union[str, int]) -> bool: + """ + Returns True if `id` is stored in this configuration. + + Args: + id: id to specify the expected position. See also :py:meth:`__getitem__`. + """ + try: + _ = self[id] + return True + except KeyError: + return False + def parse(self, reset: bool = True): """ Recursively resolve `self.config` to replace the macro tokens with target content. @@ -230,7 +243,7 @@ def read_meta(self, f: Union[PathLike, Sequence[PathLike], Dict], **kwargs): Args: f: filepath of the metadata file, the content must be a dictionary, - if providing a list of files, wil merge the content of them. + if providing a list of files, will merge the content of them. if providing a dictionary directly, use it as metadata. kwargs: other arguments for ``json.load`` or ``yaml.safe_load``, depends on the file format. diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index 9ae0fae6c6..1b5118e0bb 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -17,6 +17,8 @@ import warnings from logging.config import fileConfig from pathlib import Path +from shutil import copyfile +from textwrap import dedent from typing import Dict, Mapping, Optional, Sequence, Tuple, Union import torch @@ -25,9 +27,10 @@ from monai.apps.utils import _basename, download_url, extractall, get_logger from monai.bundle.config_item import ConfigComponent from monai.bundle.config_parser import ConfigParser +from monai.bundle.utils import DEFAULT_INFERENCE, DEFAULT_METADATA from monai.config import IgniteInfo, PathLike from monai.data import load_net_with_metadata, save_net_with_metadata -from monai.networks import convert_to_torchscript, copy_model_state, get_state_dict +from monai.networks import convert_to_torchscript, copy_model_state, get_state_dict, save_state from monai.utils import check_parent_dir, get_equivalent_dtype, min_version, optional_import from monai.utils.misc import ensure_tuple @@ -621,3 +624,80 @@ def ckpt_export( more_extra_files=extra_files, ) logger.info(f"exported to TorchScript file: {filepath_}.") + + +def init_bundle( + bundle_dir: PathLike, + ckpt_file: Optional[PathLike] = None, + network: Optional[torch.nn.Module] = None, + metadata_str: Union[Dict, str] = DEFAULT_METADATA, + inference_str: Union[Dict, str] = DEFAULT_INFERENCE, +): + """ + Initialise a new bundle directory with some default configuration files and optionally network weights. + + Typical usage example: + + .. code-block:: bash + + python -m monai.bundle init_bundle /path/to/bundle_dir network_ckpt.pt + + Args: + bundle_dir: directory name to create, must not exist but parent direct must exist + ckpt_file: optional checkpoint file to copy into bundle + network: if given instead of ckpt_file this network's weights will be stored in bundle + """ + + bundle_dir = Path(bundle_dir).absolute() + + if bundle_dir.exists(): + raise ValueError(f"Specified bundle directory '{str(bundle_dir)}' already exists") + + if not bundle_dir.parent.is_dir(): + raise ValueError(f"Parent directory of specified bundle directory '{str(bundle_dir)}' does not exist") + + configs_dir = bundle_dir / "configs" + models_dir = bundle_dir / "models" + docs_dir = bundle_dir / "docs" + + bundle_dir.mkdir() + configs_dir.mkdir() + models_dir.mkdir() + docs_dir.mkdir() + + if isinstance(metadata_str, dict): + metadata_str = json.dumps(metadata_str, indent=4) + + if isinstance(inference_str, dict): + inference_str = json.dumps(inference_str, indent=4) + + with open(str(configs_dir / "metadata.json"), "w") as o: + o.write(metadata_str) + + with open(str(configs_dir / "inference.json"), "w") as o: + o.write(inference_str) + + with open(str(docs_dir / "README.md"), "w") as o: + readme = """ + # Your Model Name + + Describe your model here and how to run it, for example using `inference.json`: + + ``` + python -m monai.bundle run evaluating \ + --meta_file /path/to/bundle/configs/metadata.json \ + --config_file /path/to/bundle/configs/inference.json \ + --dataset_dir ./input \ + --bundle_root /path/to/bundle + ``` + """ + + o.write(dedent(readme)) + + with open(str(docs_dir / "license.txt"), "w") as o: + o.write("Select a license and place its terms here\n") + + if ckpt_file is not None: + copyfile(str(ckpt_file), str(models_dir / "model.pt")) + elif network is not None: + save_state(network, str(models_dir / "model.pt")) diff --git a/monai/bundle/utils.py b/monai/bundle/utils.py index ba5c2729e7..e3eb362941 100644 --- a/monai/bundle/utils.py +++ b/monai/bundle/utils.py @@ -9,6 +9,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json +import os +import zipfile +from typing import Any + +from monai.config.deviceconfig import get_config_values +from monai.utils import optional_import + +yaml, _ = optional_import("yaml") + __all__ = ["ID_REF_KEY", "ID_SEP_KEY", "EXPR_KEY", "MACRO_KEY"] @@ -16,3 +26,147 @@ ID_SEP_KEY = "#" # separator for the ID of a ConfigItem EXPR_KEY = "$" # start of a ConfigExpression MACRO_KEY = "%" # start of a macro of a config + + +_conf_values = get_config_values() + +DEFAULT_METADATA = { + "version": "0.0.1", + "changelog": {"0.0.1": "Initial version"}, + "monai_version": _conf_values["MONAI"], + "pytorch_version": _conf_values["Pytorch"], + "numpy_version": _conf_values["Numpy"], + "optional_packages_version": {}, + "task": "Describe what the network predicts", + "description": "A longer description of what the network does, use context, inputs, outputs, etc.", + "authors": "Your Name Here", + "copyright": "Copyright (c) Your Name Here", + "network_data_format": {"inputs": {}, "outputs": {}}, +} + +DEFAULT_INFERENCE = { + "imports": ["$import glob"], + "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", + "ckpt_path": "$@bundle_root + '/models/model.pt'", + "dataset_dir": "/workspace/data", + "datalist": "$list(sorted(glob.glob(@dataset_dir + '/*.jpeg')))", + "network_def": {"_target_": "???", "spatial_dims": 2}, + "network": "$@network_def.to(@device)", + "preprocessing": { + "_target_": "Compose", + "transforms": [ + {"_target_": "LoadImaged", "keys": "image"}, + {"_target_": "AddChanneld", "keys": "image"}, + {"_target_": "ScaleIntensityd", "keys": "image"}, + {"_target_": "EnsureTyped", "keys": "image", "device": "@device"}, + ], + }, + "dataset": {"_target_": "Dataset", "data": "$[{'image': i} for i in @datalist]", "transform": "@preprocessing"}, + "dataloader": { + "_target_": "DataLoader", + "dataset": "@dataset", + "batch_size": 1, + "shuffle": False, + "num_workers": 0, + }, + "inferer": {"_target_": "SimpleInferer"}, + "postprocessing": { + "_target_": "Compose", + "transforms": [ + {"_target_": "Activationsd", "keys": "pred", "softmax": True}, + {"_target_": "AsDiscreted", "keys": "pred", "argmax": True}, + ], + }, + "handlers": [ + { + "_target_": "CheckpointLoader", + "_disabled_": "$not os.path.exists(@ckpt_path)", + "load_path": "@ckpt_path", + "load_dict": {"model": "@network"}, + } + ], + "evaluator": { + "_target_": "SupervisedEvaluator", + "device": "@device", + "val_data_loader": "@dataloader", + "network": "@network", + "inferer": "@inferer", + "postprocessing": "@postprocessing", + "val_handlers": "@handlers", + }, + "evaluating": ["$@evaluator.run()"], +} + + +def load_bundle_config(bundle_path: str, *config_names, **load_kw_args) -> Any: + """ + Load the metadata and nominated configuration files from a MONAI bundle without loading the network itself. + + This function will load the information from the bundle, which can be a directory or a zip file containing a + directory or a Torchscript bundle, and return the parser object with the information. This saves having to load + the model if only the information is wanted, and can work on any sort of bundle format. + + Args: + bundle_path: path to the bundle directory or zip file + config_names: names of configuration files with extensions to load, should not be full paths but just name+ext + load_kw_args: keyword arguments to pass to the ConfigParser object when loading + + Returns: + ConfigParser object containing the parsed information + """ + + from monai.bundle.config_parser import ConfigParser # avoids circular import + + parser = ConfigParser() + + if not os.path.exists(bundle_path): + raise ValueError(f"Cannot find bundle file/directory '{bundle_path}'") + + # bundle is a directory, read files directly + if os.path.isdir(bundle_path): + conf_data = [] + parser.read_meta(f=os.path.join(bundle_path, "configs", "metadata.json"), **load_kw_args) + + for cname in config_names: + cpath = os.path.join(bundle_path, "configs", cname) + if not os.path.exists(cpath): + raise ValueError(f"Cannot find config file '{cpath}'") + + conf_data.append(cpath) + + parser.read_config(f=conf_data, **load_kw_args) + else: + # bundle is a zip file which is either a zipped directory or a Torchscript archive + + name, _ = os.path.splitext(os.path.basename(bundle_path)) + + archive = zipfile.ZipFile(bundle_path, "r") + + all_files = archive.namelist() + + zip_meta_name = f"{name}/configs/metadata.json" + + if zip_meta_name in all_files: + prefix = f"{name}/configs/" # zipped directory location for files + else: + zip_meta_name = f"{name}/extra/metadata.json" + prefix = f"{name}/extra/" # Torchscript location for files + + meta_json = json.loads(archive.read(zip_meta_name)) + parser.read_meta(f=meta_json) + + for cname in config_names: + full_cname = prefix + cname + if full_cname not in all_files: + raise ValueError(f"Cannot find config file '{full_cname}'") + + ardata = archive.read(full_cname) + + if full_cname.lower().endswith("json"): + cdata = json.loads(ardata, **load_kw_args) + elif full_cname.lower().endswith(("yaml", "yml")): + cdata = yaml.safe_load(ardata, **load_kw_args) + + parser.read_config(f=cdata) + + return parser diff --git a/tests/min_tests.py b/tests/min_tests.py index f17aaa85b0..f7e563dc79 100644 --- a/tests/min_tests.py +++ b/tests/min_tests.py @@ -168,6 +168,8 @@ def run_testsuit(): "test_bundle_verify_metadata", "test_bundle_verify_net", "test_bundle_ckpt_export", + "test_bundle_utils", + "test_bundle_init_bundle", ] assert sorted(exclude_cases) == sorted(set(exclude_cases)), f"Duplicated items in {exclude_cases}" diff --git a/tests/test_bundle_init_bundle.py b/tests/test_bundle_init_bundle.py new file mode 100644 index 0000000000..24fc425c31 --- /dev/null +++ b/tests/test_bundle_init_bundle.py @@ -0,0 +1,41 @@ +# 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. + +import os +import subprocess +import tempfile +import unittest + +import torch + +from monai.networks.nets import UNet +from tests.utils import skip_if_windows + + +@skip_if_windows +class TestBundleInit(unittest.TestCase): + def test_bundle(self): + with tempfile.TemporaryDirectory() as tempdir: + net = UNet(2, 1, 1, [4, 8], [2]) + torch.save(net.state_dict(), tempdir + "/test.pt") + + bundle_root = tempdir + "/test_bundle" + + cmd = ["coverage", "run", "-m", "monai.bundle", "init_bundle", bundle_root, tempdir + "/test.pt"] + subprocess.check_call(cmd) + + self.assertTrue(os.path.exists(bundle_root + "/configs/metadata.json")) + self.assertTrue(os.path.exists(bundle_root + "/configs/inference.json")) + self.assertTrue(os.path.exists(bundle_root + "/models/model.pt")) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_bundle_utils.py b/tests/test_bundle_utils.py new file mode 100644 index 0000000000..46b29651cd --- /dev/null +++ b/tests/test_bundle_utils.py @@ -0,0 +1,117 @@ +# 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. + +import os +import shutil +import subprocess +import tempfile +import unittest + +import torch + +from monai.bundle.utils import load_bundle_config +from monai.networks.nets import UNet +from tests.utils import skip_if_windows + +metadata = """ +{ + "test_value": 1, + "test_list": [2,3] +} +""" + +test_json = """ +{ + "test_dict": { + "a": 3, + "b": "c" + }, + "network_def": { + "_target_": "UNet", + "spatial_dims": 2, + "in_channels": 1, + "out_channels": 1, + "channels": [4,8], + "strides": [2] + } +} +""" + + +@skip_if_windows +class TestLoadBundleConfig(unittest.TestCase): + def setUp(self): + self.bundle_dir = tempfile.TemporaryDirectory() + self.dir_name = os.path.join(self.bundle_dir.name, "TestBundle") + self.configs_name = os.path.join(self.dir_name, "configs") + self.models_name = os.path.join(self.dir_name, "models") + self.metadata_name = os.path.join(self.configs_name, "metadata.json") + self.test_name = os.path.join(self.configs_name, "test.json") + self.modelpt_name = os.path.join(self.models_name, "model.pt") + + self.zip_file = os.path.join(self.bundle_dir.name, "TestBundle.zip") + self.ts_file = os.path.join(self.bundle_dir.name, "TestBundle.ts") + + # create the directories for the bundle + os.mkdir(self.dir_name) + os.mkdir(self.configs_name) + os.mkdir(self.models_name) + + # fill bundle configs + + with open(self.metadata_name, "w") as o: + o.write(metadata) + + with open(self.test_name, "w") as o: + o.write(test_json) + + # save network + net = UNet(2, 1, 1, [4, 8], [2]) + torch.save(net.state_dict(), self.modelpt_name) + + def tearDown(self): + self.bundle_dir.cleanup() + + def test_load_config_dir(self): + p = load_bundle_config(self.dir_name, "test.json") + + self.assertEqual(p["_meta_"]["test_value"], 1) + + self.assertEqual(p["test_dict"]["b"], "c") + + def test_load_config_zip(self): + # create a zip of the bundle + shutil.make_archive(self.zip_file[:-4], "zip", self.bundle_dir.name) + + p = load_bundle_config(self.zip_file, "test.json") + + self.assertEqual(p["_meta_"]["test_value"], 1) + + self.assertEqual(p["test_dict"]["b"], "c") + + def test_load_config_ts(self): + # create a Torchscript zip of the bundle + cmd = ["python", "-m", "monai.bundle", "ckpt_export", "network_def", "--filepath", self.ts_file] + cmd += ["--meta_file", self.metadata_name] + cmd += ["--config_file", self.test_name] + cmd += ["--ckpt_file", self.modelpt_name] + + subprocess.check_output(cmd, stderr=subprocess.STDOUT) + + p = load_bundle_config(self.ts_file, "test.json") + + self.assertEqual(p["_meta_"]["test_value"], 1) + + self.assertEqual(p["test_dict"]["b"], "c") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_config_parser.py b/tests/test_config_parser.py index da230e7794..ffb6ebc634 100644 --- a/tests/test_config_parser.py +++ b/tests/test_config_parser.py @@ -195,6 +195,22 @@ def test_list_expressions(self): parser.get_parsed_content("training", lazy=True, instantiate=True, eval_expr=True) np.testing.assert_allclose(parser.get_parsed_content("training#1", lazy=True), [0.7942, 1.5885], atol=1e-4) + def test_contains(self): + empty_parser = ConfigParser({}) + empty_parser.parse() + + parser = ConfigParser({"value": 1, "entry": "string content"}) + parser.parse() + + with self.subTest("Testing empty parser"): + self.assertFalse("something" in empty_parser) + + with self.subTest("Testing with keys"): + self.assertTrue("value" in parser) + self.assertFalse("value1" in parser) + self.assertTrue("entry" in parser) + self.assertFalse("entr" in parser) + if __name__ == "__main__": unittest.main()