# 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 ast

import json

import os

import re

import urllib

import warnings

from collections.abc import Mapping, Sequence

from functools import partial

from pathlib import Path

from pydoc import locate

from shutil import copyfile

from textwrap import dedent

from typing import Any, Callable

import torch

from torch.cuda import is_available

from monai._version import get_versions

from monai.apps.utils import _basename, _extract_zip, download_url, extractall, get_logger

from monai.bundle.config_parser import ConfigParser

from monai.bundle.utils import DEFAULT_INFERENCE, DEFAULT_METADATA, merge_kv

from monai.bundle.workflows import BundleWorkflow, ConfigWorkflow

from monai.config import PathLike

from monai.data import load_net_with_metadata, save_net_with_metadata

from monai.networks import (

convert_to_onnx,

convert_to_torchscript,

convert_to_trt,

copy_model_state,

get_state_dict,

save_state,

)

from monai.utils import (

IgniteInfo,

check_parent_dir,

ensure_tuple,

get_equivalent_dtype,

min_version,

optional_import,

pprint_edges,

)

validate, _ = optional_import("jsonschema", name="validate")

ValidationError, _ = optional_import("jsonschema.exceptions", name="ValidationError")

Checkpoint, has_ignite = optional_import("ignite.handlers", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Checkpoint")

requests, has_requests = optional_import("requests")

onnx, _ = optional_import("onnx")

huggingface_hub, _ = optional_import("huggingface_hub")

logger = get_logger(module_name=__name__)

# set BUNDLE_DOWNLOAD_SRC="ngc" to use NGC source in default for bundle download

# set BUNDLE_DOWNLOAD_SRC="github" to use github source in default for bundle download

DEFAULT_DOWNLOAD_SOURCE = os.environ.get("BUNDLE_DOWNLOAD_SRC", "monaihosting")

PPRINT_CONFIG_N = 5

MONAI_HOSTING_BASE_URL = "https://api.ngc.nvidia.com/v2/models/nvidia/monaihosting"

NGC_BASE_URL = "https://api.ngc.nvidia.com/v2/models/nvidia/monaitoolkit"

def update_kwargs(args: str | dict | None = None, ignore_none: bool = True, **kwargs: Any) -> dict:

"""

Update the `args` dictionary with the input `kwargs`.

For dict data, recursively update the content based on the keys.

Example::

from monai.bundle import update_kwargs

update_kwargs({'exist': 1}, exist=2, new_arg=3)

# return {'exist': 2, 'new_arg': 3}

Args:

args: source `args` dictionary (or a json/yaml filename to read as dictionary) to update.

ignore_none: whether to ignore input args with None value, default to `True`.

kwargs: key=value pairs to be merged into `args`.

"""

args_: dict = args if isinstance(args, dict) else {}

if isinstance(args, str):

# args are defined in a structured file

args_ = ConfigParser.load_config_file(args)

if isinstance(args, (tuple, list)) and all(isinstance(x, str) for x in args):

primary, overrides = args

args_ = update_kwargs(primary, ignore_none, **update_kwargs(overrides, ignore_none, **kwargs))

if not isinstance(args_, dict):

return args_

# recursively update the default args with new args

for k, v in kwargs.items():

if ignore_none and v is None:

continue

if isinstance(v, dict) and isinstance(args_.get(k), dict):

args_[k] = update_kwargs(args_[k], ignore_none, **v)

else:

merge_kv(args_, k, v)

return args_

_update_args = update_kwargs # backward compatibility

def _pop_args(src: dict, *args: Any, **kwargs: Any) -> tuple:

"""

Pop args from the `src` dictionary based on specified keys in `args` and (key, default value) pairs in `kwargs`.

"""

return tuple([src.pop(i) for i in args] + [src.pop(k, v) for k, v in kwargs.items()])

def _log_input_summary(tag: str, args: dict) -> None:

logger.info(f"--- input summary of monai.bundle.scripts.{tag} ---")

for name, val in args.items():

logger.info(f"> {name}: {pprint_edges(val, PPRINT_CONFIG_N)}")

logger.info("---\n\n")

def _get_var_names(expr: str) -> list[str]:

"""

Parse the expression and discover what variables are present in it based on ast module.

Args:

expr: source expression to parse.

"""

tree = ast.parse(expr)

return [m.id for m in ast.walk(tree) if isinstance(m, ast.Name)]

def _get_fake_spatial_shape(shape: Sequence[str | int], p: int = 1, n: int = 1, any: int = 1) -> tuple:

"""

Get spatial shape for fake data according to the specified shape pattern.

It supports `int` number and `string` with formats like: "32", "32 * n", "32 ** p", "32 ** p *n".

Args:

shape: specified pattern for the spatial shape.

p: power factor to generate fake data shape if dim of expected shape is "x**p", default to 1.

p: multiply factor to generate fake data shape if dim of expected shape is "x*n", default to 1.

any: specified size to generate fake data shape if dim of expected shape is "*", default to 1.

"""

ret = []

for i in shape:

if isinstance(i, int):

ret.append(i)

elif isinstance(i, str):

if i == "*":

ret.append(any)

else:

for c in _get_var_names(i):

if c not in ["p", "n"]:

raise ValueError(f"only support variables 'p' and 'n' so far, but got: {c}.")

ret.append(eval(i, {"p": p, "n": n}))

else:

raise ValueError(f"spatial shape items must be int or string, but got: {type(i)} {i}.")

return tuple(ret)

def _get_git_release_url(repo_owner: str, repo_name: str, tag_name: str, filename: str) -> str:

return f"https://github.com/{repo_owner}/{repo_name}/releases/download/{tag_name}/{filename}"

def _get_ngc_bundle_url(model_name: str, version: str) -> str:

return f"{NGC_BASE_URL}/{model_name.lower()}/versions/{version}/files"

def _get_ngc_private_base_url(repo: str) -> str:

return f"https://api.ngc.nvidia.com/v2/{repo}/models"

def _get_ngc_private_bundle_url(model_name: str, version: str, repo: str) -> str:

return f"{_get_ngc_private_base_url(repo)}/{model_name.lower()}/versions/{version}/zip"

def _get_monaihosting_bundle_url(model_name: str, version: str) -> str:

return f"{MONAI_HOSTING_BASE_URL}/{model_name.lower()}/versions/{version}/files/{model_name}_v{version}.zip"

def _download_from_github(repo: str, download_path: Path, filename: str, progress: bool = True) -> None:

repo_owner, repo_name, tag_name = repo.split("/")

if ".zip" not in filename:

filename += ".zip"

url = _get_git_release_url(repo_owner, repo_name, tag_name=tag_name, filename=filename)

filepath = download_path / f"{filename}"

download_url(url=url, filepath=filepath, hash_val=None, progress=progress)

extractall(filepath=filepath, output_dir=download_path, has_base=True)

def _download_from_monaihosting(download_path: Path, filename: str, version: str, progress: bool) -> None:

url = _get_monaihosting_bundle_url(model_name=filename, version=version)

filepath = download_path / f"{filename}_v{version}.zip"

download_url(url=url, filepath=filepath, hash_val=None, progress=progress)

extractall(filepath=filepath, output_dir=download_path, has_base=True)

def _download_from_bundle_info(download_path: Path, filename: str, version: str, progress: bool) -> None:

bundle_info = get_bundle_info(bundle_name=filename, version=version)

if not bundle_info:

raise ValueError(f"Bundle info not found for {filename} v{version}.")

url = bundle_info["browser_download_url"]

filepath = download_path / f"{filename}_v{version}.zip"

download_url(url=url, filepath=filepath, hash_val=None, progress=progress)

extractall(filepath=filepath, output_dir=download_path, has_base=True)

def _add_ngc_prefix(name: str, prefix: str = "monai_") -> str:

if name.startswith(prefix):

return name

return f"{prefix}{name}"

def _remove_ngc_prefix(name: str, prefix: str = "monai_") -> str:

if name.startswith(prefix):

return name[len(prefix) :]

return name

def _get_all_download_files(request_url: str, headers: dict | None = None) -> list[str]:

if not has_requests:

raise ValueError("requests package is required, please install it.")

headers = {} if headers is None else headers

response = requests.get(request_url, headers=headers)

response.raise_for_status()

model_info = json.loads(response.text)

if not isinstance(model_info, dict) or "modelFiles" not in model_info:

raise ValueError("The data is not a dictionary or it does not have the key 'modelFiles'.")

model_files = model_info["modelFiles"]

return [f["path"] for f in model_files]

def _download_from_ngc(

download_path: Path,

filename: str,

version: str,

prefix: str = "monai_",

remove_prefix: str | None = "monai_",

progress: bool = True,

) -> None:

# ensure prefix is contained

filename = _add_ngc_prefix(filename, prefix=prefix)

url = _get_ngc_bundle_url(model_name=filename, version=version)

if remove_prefix:

filename = _remove_ngc_prefix(filename, prefix=remove_prefix)

filepath = download_path / filename

filepath.mkdir(parents=True, exist_ok=True)

for file in _get_all_download_files(url):

download_url(url=f"{url}/{file}", filepath=f"{filepath}/{file}", hash_val=None, progress=progress)

def _download_from_ngc_private(

download_path: Path,

filename: str,

version: str,

repo: str,

prefix: str = "monai_",

remove_prefix: str | None = "monai_",

headers: dict | None = None,

) -> None:

# ensure prefix is contained

filename = _add_ngc_prefix(filename, prefix=prefix)

request_url = _get_ngc_private_bundle_url(model_name=filename, version=version, repo=repo)

if has_requests:

headers = {} if headers is None else headers

response = requests.get(request_url, headers=headers)

response.raise_for_status()

else:

raise ValueError("NGC API requires requests package. Please install it.")

os.makedirs(download_path, exist_ok=True)

zip_path = download_path / f"{filename}_v{version}.zip"

with open(zip_path, "wb") as f:

f.write(response.content)

logger.info(f"Downloading: {zip_path}.")

if remove_prefix:

filename = _remove_ngc_prefix(filename, prefix=remove_prefix)

extract_path = download_path / f"{filename}"

_extract_zip(zip_path, extract_path)

logger.info(f"Writing into directory: {extract_path}.")

def _get_ngc_token(api_key, retry=0):

"""Try to connect to NGC."""

url = "https://authn.nvidia.com/token?service=ngc"

headers = {"Accept": "application/json", "Authorization": "ApiKey " + api_key}

if has_requests:

response = requests.get(url, headers=headers)

if not response.ok:

# retry 3 times, if failed, raise an error.

if retry < 3:

logger.info(f"Retrying {retry} time(s) to GET {url}.")

return _get_ngc_token(url, retry + 1)

raise RuntimeError("NGC API response is not ok. Failed to get token.")

else:

token = response.json()["token"]

return token

def _examine_monai_version(monai_version: str) -> tuple[bool, str]:

"""Examine if the package version is compatible with the MONAI version in the metadata."""

version_dict = get_versions()

package_version = version_dict.get("version", "0+unknown")

if package_version == "0+unknown":

return False, "Package version is not available. Skipping version check."

if monai_version == "0+unknown":

return False, "MONAI version is not specified in the bundle. Skipping version check."

# treat rc versions as the same as the release version

package_version = re.sub(r"rc\d.*", "", package_version)

monai_version = re.sub(r"rc\d.*", "", monai_version)

if package_version < monai_version:

return (

False,

f"Your MONAI version is {package_version}, but the bundle is built on MONAI version {monai_version}.",

)

return True, ""

def _check_monai_version(bundle_dir: PathLike, name: str) -> None:

"""Get the `monai_version` from the metadata.json and compare if it is smaller than the installed `monai` package version"""

metadata_file = Path(bundle_dir) / name / "configs" / "metadata.json"

if not metadata_file.exists():

logger.warning(f"metadata file not found in {metadata_file}.")

return

with open(metadata_file) as f:

metadata = json.load(f)

is_compatible, msg = _examine_monai_version(metadata.get("monai_version", "0+unknown"))

if not is_compatible:

logger.warning(msg)

def _list_latest_versions(data: dict, max_versions: int = 3) -> list[str]:

"""

Extract the latest versions from the data dictionary.

Args:

data: the data dictionary.

max_versions: the maximum number of versions to return.

Returns:

versions of the latest models in the reverse order of creation date, e.g. ['1.0.0', '0.9.0', '0.8.0'].

"""

# Check if the data is a dictionary and it has the key 'modelVersions'

if not isinstance(data, dict) or "modelVersions" not in data:

raise ValueError("The data is not a dictionary or it does not have the key 'modelVersions'.")

# Extract the list of model versions

model_versions = data["modelVersions"]

if (

not isinstance(model_versions, list)

or len(model_versions) == 0

or "createdDate" not in model_versions[0]

or "versionId" not in model_versions[0]

):

raise ValueError(

"The model versions are not a list or it is empty or it does not have the keys 'createdDate' and 'versionId'."

)

# Sort the versions by the 'createdDate' in descending order

sorted_versions = sorted(model_versions, key=lambda x: x["createdDate"], reverse=True)

return [v["versionId"] for v in sorted_versions[:max_versions]]

def _get_latest_bundle_version_ngc(name: str, repo: str | None = None, headers: dict | None = None) -> str:

base_url = _get_ngc_private_base_url(repo) if repo else NGC_BASE_URL

version_endpoint = base_url + f"/{name.lower()}/versions/"

if not has_requests:

raise ValueError("requests package is required, please install it.")

version_header = {"Accept-Encoding": "gzip, deflate"} # Excluding 'zstd' to fit NGC requirements

if headers:

version_header.update(headers)

resp = requests.get(version_endpoint, headers=version_header)

resp.raise_for_status()

model_info = json.loads(resp.text)

latest_versions = _list_latest_versions(model_info)

for version in latest_versions:

file_endpoint = base_url + f"/{name.lower()}/versions/{version}/files/configs/metadata.json"

resp = requests.get(file_endpoint, headers=headers)

metadata = json.loads(resp.text)

resp.raise_for_status()

# if the package version is not available or the model is compatible with the package version

is_compatible, _ = _examine_monai_version(metadata["monai_version"])

if is_compatible:

if version != latest_versions[0]:

logger.info(f"Latest version is {latest_versions[0]}, but the compatible version is {version}.")

return version

# if no compatible version is found, return the latest version

return latest_versions[0]

def _get_latest_bundle_version(

source: str, name: str, repo: str, **kwargs: Any

) -> dict[str, list[str] | str] | Any | None:

if source == "ngc":

name = _add_ngc_prefix(name)

return _get_latest_bundle_version_ngc(name)

elif source == "monaihosting":

return get_bundle_versions(name, repo="Project-MONAI/model-zoo", tag="dev")["latest_version"]

elif source == "ngc_private":

headers = kwargs.pop("headers", {})

name = _add_ngc_prefix(name)

return _get_latest_bundle_version_ngc(name, repo=repo, headers=headers)

elif source == "github":

repo_owner, repo_name, tag_name = repo.split("/")

return get_bundle_versions(name, repo=f"{repo_owner}/{repo_name}", tag=tag_name)["latest_version"]

elif source == "huggingface_hub":

refs = huggingface_hub.list_repo_refs(repo_id=repo)

if len(refs.tags) > 0:

all_versions = [t.name for t in refs.tags] # git tags, not to be confused with `tag`

latest_version = ["latest_version" if "latest_version" in all_versions else all_versions[-1]][0]

else:

latest_version = [b.name for b in refs.branches][0] # use the branch that was last updated

return latest_version

else:

raise ValueError(

f"To get the latest bundle version, source should be 'github', 'monaihosting' or 'ngc', got {source}."

)

def _process_bundle_dir(bundle_dir: PathLike | None = None) -> Path:

if bundle_dir is None:

get_dir, has_home = optional_import("torch.hub", name="get_dir")

if has_home:

bundle_dir = Path(get_dir()) / "bundle"

else:

raise ValueError("bundle_dir=None, but no suitable default directory computed. Upgrade Pytorch to 1.6+ ?")

return Path(bundle_dir)

def download(

name: str | None = None,

version: str | None = None,

bundle_dir: PathLike | None = None,

source: str = DEFAULT_DOWNLOAD_SOURCE,

repo: str | None = None,

url: str | None = None,

remove_prefix: str | None = "monai_",

progress: bool = True,

args_file: str | None = None,

) -> None:

"""

download bundle from the specified source or url. The bundle should be a zip file and it

will be extracted after downloading.

This function refers to:

https://pytorch.org/docs/stable/_modules/torch/hub.html

Typical usage examples:

.. code-block:: bash

# Execute this module as a CLI entry, and download bundle from the model-zoo repo:

python -m monai.bundle download --name <bundle_name> --version "0.1.0" --bundle_dir "./"

# Execute this module as a CLI entry, and download bundle from specified github repo:

python -m monai.bundle download --name <bundle_name> --source "github" --repo "repo_owner/repo_name/release_tag"

# Execute this module as a CLI entry, and download bundle from ngc with latest version:

python -m monai.bundle download --name <bundle_name> --source "ngc" --bundle_dir "./"

# Execute this module as a CLI entry, and download bundle from monaihosting with latest version:

python -m monai.bundle download --name <bundle_name> --source "monaihosting" --bundle_dir "./"

# Execute this module as a CLI entry, and download bundle from Hugging Face Hub:

python -m monai.bundle download --name "bundle_name" --source "huggingface_hub" --repo "repo_owner/repo_name"

# Execute this module as a CLI entry, and download bundle via URL:

python -m monai.bundle download --name <bundle_name> --url <url>

# Execute this module as a CLI entry, and download bundle from ngc_private with latest version:

python -m monai.bundle download --name <bundle_name> --source "ngc_private" --bundle_dir "./" --repo "org/org_name"

# Set default args of `run` in a JSON / YAML file, help to record and simplify the command line.

# Other args still can override the default args at runtime.

# The content of the JSON / YAML file is a dictionary. For example:

# {"name": "spleen", "bundle_dir": "download", "source": ""}

# then do the following command for downloading:

python -m monai.bundle download --args_file "args.json" --source "github"

Args:

name: bundle name. If `None` and `url` is `None`, it must be provided in `args_file`.

for example:

"spleen_ct_segmentation", "prostate_mri_anatomy" in model-zoo:

https://github.com/Project-MONAI/model-zoo/releases/tag/hosting_storage_v1.

"monai_brats_mri_segmentation" in ngc:

https://catalog.ngc.nvidia.com/models?filters=&orderBy=scoreDESC&query=monai.

version: version name of the target bundle to download, like: "0.1.0". If `None`, will download

the latest version (or the last commit to the `main` branch in the case of Hugging Face Hub).

bundle_dir: target directory to store the downloaded data.

Default is `bundle` subfolder under `torch.hub.get_dir()`.

source: storage location name. This argument is used when `url` is `None`.

In default, the value is achieved from the environment variable BUNDLE_DOWNLOAD_SRC, and

it should be "ngc", "monaihosting", "github", "ngc_private", or "huggingface_hub".

If source is "ngc_private", you need specify the NGC_API_KEY in the environment variable.

repo: repo name. This argument is used when `url` is `None` and `source` is "github" or "huggingface_hub".

If `source` is "github", it should be in the form of "repo_owner/repo_name/release_tag".

If `source` is "huggingface_hub", it should be in the form of "repo_owner/repo_name". Please note that

bundles for "monaihosting" source are also hosted on Hugging Face Hub, but the "repo_id" is always in the form

of "MONAI/bundle_name", therefore, this argument is not required for "monaihosting" source.

If `source` is "ngc_private", it should be in the form of "org/org_name" or "org/org_name/team/team_name",

or you can specify the environment variable NGC_ORG and NGC_TEAM.

url: url to download the data. If not `None`, data will be downloaded directly

and `source` will not be checked.

If `name` is `None`, filename is determined by `monai.apps.utils._basename(url)`.

remove_prefix: This argument is used when `source` is "ngc" or "ngc_private". Currently, all ngc bundles

have the ``monai_`` prefix, which is not existing in their model zoo contrasts. In order to

maintain the consistency between these two sources, remove prefix is necessary.

Therefore, if specified, downloaded folder name will remove the prefix.

progress: whether to display a progress bar.

args_file: a JSON or YAML file to provide default values for all the args in this function.

so that the command line inputs can be simplified.

"""

_args = update_kwargs(

args=args_file,

name=name,

version=version,

bundle_dir=bundle_dir,

source=source,

repo=repo,

url=url,

remove_prefix=remove_prefix,

progress=progress,

)

_log_input_summary(tag="download", args=_args)

source_, progress_, remove_prefix_, repo_, name_, version_, bundle_dir_, url_ = _pop_args(

_args, "source", "progress", remove_prefix=None, repo=None, name=None, version=None, bundle_dir=None, url=None

)

bundle_dir_ = _process_bundle_dir(bundle_dir_)

if repo_ is None:

org_ = os.getenv("NGC_ORG", None)

team_ = os.getenv("NGC_TEAM", None)

if org_ is not None and source_ == "ngc_private":

repo_ = f"org/{org_}/team/{team_}" if team_ is not None else f"org/{org_}"

else:

repo_ = "Project-MONAI/model-zoo/hosting_storage_v1"

if len(repo_.split("/")) not in (2, 4) and source_ == "ngc_private":

raise ValueError(f"repo should be in the form of `org/org_name/team/team_name` or `org/org_name`, got {repo_}.")

if len(repo_.split("/")) != 3 and source_ == "github":

raise ValueError(f"repo should be in the form of `repo_owner/repo_name/release_tag`, got {repo_}.")

elif len(repo_.split("/")) != 2 and source_ == "huggingface_hub":

raise ValueError(f"Hugging Face Hub repo should be in the form of `repo_owner/repo_name`, got {repo_}.")

if url_ is not None:

if name_ is not None:

filepath = bundle_dir_ / f"{name_}.zip"

else:

filepath = bundle_dir_ / f"{_basename(url_)}"

download_url(url=url_, filepath=filepath, hash_val=None, progress=progress_)

extractall(filepath=filepath, output_dir=bundle_dir_, has_base=True)

else:

headers = {}

if name_ is None:

raise ValueError(f"To download from source: {source_}, `name` must be provided.")

if source == "ngc_private":

api_key = os.getenv("NGC_API_KEY", None)

if api_key is None:

raise ValueError("API key is required for ngc_private source.")

else:

token = _get_ngc_token(api_key)

headers = {"Authorization": f"Bearer {token}"}

if version_ is None:

version_ = _get_latest_bundle_version(source=source_, name=name_, repo=repo_, headers=headers)

if source_ == "github":

name_ver = "_v".join([name_, version_]) if version_ is not None else name_

_download_from_github(repo=repo_, download_path=bundle_dir_, filename=name_ver, progress=progress_)

elif source_ == "monaihosting":

try:

extract_path = os.path.join(bundle_dir_, name_)

huggingface_hub.snapshot_download(repo_id=f"MONAI/{name_}", revision=version_, local_dir=extract_path)

except (huggingface_hub.errors.RevisionNotFoundError, huggingface_hub.errors.RepositoryNotFoundError):

# if bundle or version not found from huggingface, download from ngc monaihosting

_download_from_monaihosting(

download_path=bundle_dir_, filename=name_, version=version_, progress=progress_

)

except urllib.error.HTTPError:

# if also cannot download from ngc monaihosting, download according to bundle_info

_download_from_bundle_info(

download_path=bundle_dir_, filename=name_, version=version_, progress=progress_

)

elif source_ == "ngc":

_download_from_ngc(

download_path=bundle_dir_,

filename=name_,

version=version_,

remove_prefix=remove_prefix_,

progress=progress_,

)

elif source_ == "ngc_private":

_download_from_ngc_private(

download_path=bundle_dir_,

filename=name_,

version=version_,

remove_prefix=remove_prefix_,

repo=repo_,

headers=headers,

)

elif source_ == "huggingface_hub":

extract_path = os.path.join(bundle_dir_, name_)

huggingface_hub.snapshot_download(repo_id=repo_, revision=version_, local_dir=extract_path)

else:

raise NotImplementedError(

"Currently only download from `url`, source 'github', 'monaihosting', 'huggingface_hub' or 'ngc' are implemented,"

f"got source: {source_}."

)

_check_monai_version(bundle_dir_, name_)

def load(

name: str,

model: torch.nn.Module | None = None,

version: str | None = None,

workflow_type: str = "train",

model_file: str | None = None,

load_ts_module: bool = False,

bundle_dir: PathLike | None = None,

source: str = DEFAULT_DOWNLOAD_SOURCE,

repo: str | None = None,

remove_prefix: str | None = "monai_",

progress: bool = True,

device: str | None = None,

key_in_ckpt: str | None = None,

config_files: Sequence[str] = (),

workflow_name: str | BundleWorkflow | None = None,

args_file: str | None = None,

copy_model_args: dict | None = None,

net_override: dict | None = None,

) -> object | tuple[torch.nn.Module, dict, dict] | Any:

"""

Load model weights or TorchScript module of a bundle.

Args:

name: bundle name. If `None` and `url` is `None`, it must be provided in `args_file`.

for example:

"spleen_ct_segmentation", "prostate_mri_anatomy" in model-zoo:

https://github.com/Project-MONAI/model-zoo/releases/tag/hosting_storage_v1.

"monai_brats_mri_segmentation" in ngc:

https://catalog.ngc.nvidia.com/models?filters=&orderBy=scoreDESC&query=monai.

"mednist_gan" in monaihosting:

https://api.ngc.nvidia.com/v2/models/nvidia/monaihosting/mednist_gan/versions/0.2.0/files/mednist_gan_v0.2.0.zip

model: a pytorch module to be updated. Default to None, using the "network_def" in the bundle.

version: version name of the target bundle to download, like: "0.1.0". If `None`, will download

the latest version. If `source` is "huggingface_hub", this argument is a Git revision id.

workflow_type: specifies the workflow type: "train" or "training" for a training workflow,

or "infer", "inference", "eval", "evaluation" for a inference workflow,

other unsupported string will raise a ValueError.

default to `train` for training workflow.

model_file: the relative path of the model weights or TorchScript module within bundle.

If `None`, "models/model.pt" or "models/model.ts" will be used.

load_ts_module: a flag to specify if loading the TorchScript module.

bundle_dir: directory the weights/TorchScript module will be loaded from.

Default is `bundle` subfolder under `torch.hub.get_dir()`.

source: storage location name. This argument is used when `model_file` is not existing locally and need to be

downloaded first.

In default, the value is achieved from the environment variable BUNDLE_DOWNLOAD_SRC, and

it should be "ngc", "monaihosting", "github", or "huggingface_hub".

repo: repo name. This argument is used when `url` is `None` and `source` is "github" or "huggingface_hub".

If `source` is "github", it should be in the form of "repo_owner/repo_name/release_tag".

If `source` is "huggingface_hub", it should be in the form of "repo_owner/repo_name".

remove_prefix: This argument is used when `source` is "ngc". Currently, all ngc bundles

have the ``monai_`` prefix, which is not existing in their model zoo contrasts. In order to

maintain the consistency between these three sources, remove prefix is necessary.

Therefore, if specified, downloaded folder name will remove the prefix.

progress: whether to display a progress bar when downloading.

device: target device of returned weights or module, if `None`, prefer to "cuda" if existing.

key_in_ckpt: for nested checkpoint like `{"model": XXX, "optimizer": XXX, ...}`, specify the key of model

weights. if not nested checkpoint, no need to set.

config_files: extra filenames would be loaded. The argument only works when loading a TorchScript module,

see `_extra_files` in `torch.jit.load` for more details.

workflow_name: specified bundle workflow name, should be a string or class, default to "ConfigWorkflow".

args_file: a JSON or YAML file to provide default values for all the args in "download" function.

copy_model_args: other arguments for the `monai.networks.copy_model_state` function.

net_override: id-value pairs to override the parameters in the network of the bundle, default to `None`.

Returns:

1. If `load_ts_module` is `False` and `model` is `None`,

return model weights if can't find "network_def" in the bundle,

else return an instantiated network that loaded the weights.

2. If `load_ts_module` is `False` and `model` is not `None`,

return an instantiated network that loaded the weights.

3. If `load_ts_module` is `True`, return a triple that include a TorchScript module,

the corresponding metadata dict, and extra files dict.

please check `monai.data.load_net_with_metadata` for more details.

"""

bundle_dir_ = _process_bundle_dir(bundle_dir)

net_override = {} if net_override is None else net_override

copy_model_args = {} if copy_model_args is None else copy_model_args

if device is None:

device = "cuda:0" if is_available() else "cpu"

if model_file is None:

model_file = os.path.join("models", "model.ts" if load_ts_module is True else "model.pt")

if source == "ngc":

name = _add_ngc_prefix(name)

if remove_prefix:

name = _remove_ngc_prefix(name, prefix=remove_prefix)

full_path = os.path.join(bundle_dir_, name, model_file)

if not os.path.exists(full_path):

download(

name=name,

version=version,

bundle_dir=bundle_dir_,

source=source,

repo=repo,

remove_prefix=remove_prefix,

progress=progress,

args_file=args_file,

)

# loading with `torch.jit.load`

if load_ts_module is True:

return load_net_with_metadata(full_path, map_location=torch.device(device), more_extra_files=config_files)

# loading with `torch.load`

model_dict = torch.load(full_path, map_location=torch.device(device), weights_only=True)

if not isinstance(model_dict, Mapping):

warnings.warn(f"the state dictionary from {full_path} should be a dictionary but got {type(model_dict)}.")

model_dict = get_state_dict(model_dict)

_workflow = None

if model is None:

bundle_config_file = bundle_dir_ / name / "configs" / f"{workflow_type}.json"

if bundle_config_file.is_file():

_net_override = {f"network_def#{key}": value for key, value in net_override.items()}

_workflow = create_workflow(

workflow_name=workflow_name,

args_file=args_file,

config_file=str(bundle_config_file),

workflow_type=workflow_type,

**_net_override,

)

else:

warnings.warn(f"Cannot find the config file: {bundle_config_file}, return state dict instead.")

return model_dict

if _workflow is not None:

if not hasattr(_workflow, "network_def"):

warnings.warn("No available network definition in the bundle, return state dict instead.")

return model_dict

else:

model = _workflow.network_def

model.to(device) # type: ignore

copy_model_state(

dst=model, src=model_dict if key_in_ckpt is None else model_dict[key_in_ckpt], **copy_model_args # type: ignore

)

return model

def _get_all_bundles_info(

repo: str = "Project-MONAI/model-zoo", tag: str = "dev", auth_token: str | None = None

) -> dict[str, dict[str, dict[str, Any]]]:

if has_requests:

if tag == "hosting_storage_v1":

request_url = f"https://api.github.com/repos/{repo}/releases"

else:

request_url = f"https://raw.githubusercontent.com/{repo}/{tag}/models/model_info.json"

if auth_token is not None:

headers = {"Authorization": f"Bearer {auth_token}"}

resp = requests.get(request_url, headers=headers)

else:

resp = requests.get(request_url)

resp.raise_for_status()

else:

raise ValueError("requests package is required, please install it.")

releases_list = json.loads(resp.text)

bundle_name_pattern = re.compile(r"_v\d*.")

bundles_info: dict[str, dict[str, dict[str, Any]]] = {}

if tag == "hosting_storage_v1":

for release in releases_list:

if release["tag_name"] == tag:

for asset in release["assets"]:

asset_name = bundle_name_pattern.split(asset["name"])[0]

if asset_name not in bundles_info:

bundles_info[asset_name] = {}

asset_version = asset["name"].split(f"{asset_name}_v")[-1].replace(".zip", "")

bundles_info[asset_name][asset_version] = dict(asset)

return bundles_info

else:

for asset in releases_list.keys():

asset_name = bundle_name_pattern.split(asset)[0]

if asset_name not in bundles_info:

bundles_info[asset_name] = {}

asset_version = asset.split(f"{asset_name}_v")[-1]

bundles_info[asset_name][asset_version] = {

"name": asset,

"browser_download_url": releases_list[asset]["source"],

}

return bundles_info

def get_all_bundles_list(

repo: str = "Project-MONAI/model-zoo", tag: str = "dev", auth_token: str | None = None

) -> list[tuple[str, str]]:

"""

Get all bundles names (and the latest versions) that are stored in the release of specified repository

with the provided tag. If tag is "dev", will get model information from

https://raw.githubusercontent.com/repo_owner/repo_name/dev/models/model_info.json.

The default values of arguments correspond to the release of MONAI model zoo. In order to increase the

rate limits of calling Github APIs, you can input your personal access token.

Please check the following link for more details about rate limiting:

https://docs.github.com/en/rest/overview/resources-in-the-rest-api#rate-limiting

The following link shows how to create your personal access token:

https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/creating-a-personal-access-token

Args:

repo: it should be in the form of "repo_owner/repo_name/".

tag: the tag name of the release.

auth_token: github personal access token.

Returns:

a list of tuple in the form of (bundle name, latest version).

"""

bundles_info = _get_all_bundles_info(repo=repo, tag=tag, auth_token=auth_token)

bundles_list = []

for bundle_name in bundles_info:

latest_version = sorted(bundles_info[bundle_name].keys())[-1]

bundles_list.append((bundle_name, latest_version))

return bundles_list

def get_bundle_versions(

bundle_name: str, repo: str = "Project-MONAI/model-zoo", tag: str = "dev", auth_token: str | None = None

) -> dict[str, list[str] | str]:

"""

Get the latest version, as well as all existing versions of a bundle that is stored in the release of specified

repository with the provided tag. If tag is "dev", will get model information from

https://raw.githubusercontent.com/repo_owner/repo_name/dev/models/model_info.json.

In order to increase the rate limits of calling Github APIs, you can input your personal access token.

Please check the following link for more details about rate limiting:

https://docs.github.com/en/rest/overview/resources-in-the-rest-api#rate-limiting

The following link shows how to create your personal access token:

https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/creating-a-personal-access-token

Args:

bundle_name: bundle name.

repo: it should be in the form of "repo_owner/repo_name/".

tag: the tag name of the release.

auth_token: github personal access token.

Returns:

a dictionary that contains the latest version and all versions of a bundle.

"""

bundles_info = _get_all_bundles_info(repo=repo, tag=tag, auth_token=auth_token)

if bundle_name not in bundles_info:

raise ValueError(f"bundle: {bundle_name} is not existing in repo: {repo}.")

bundle_info = bundles_info[bundle_name]

all_versions = sorted(bundle_info.keys())

return {"latest_version": all_versions[-1], "all_versions": all_versions}

def get_bundle_info(

bundle_name: str,

version: str | None = None,

repo: str = "Project-MONAI/model-zoo",

tag: str = "dev",

auth_token: str | None = None,

) -> dict[str, Any]:

"""

Get all information (include "name" and "browser_download_url") of a bundle

with the specified bundle name and version which is stored in the release of specified repository with the provided tag.

In order to increase the rate limits of calling Github APIs, you can input your personal access token.

Please check the following link for more details about rate limiting:

https://docs.github.com/en/rest/overview/resources-in-the-rest-api#rate-limiting

The following link shows how to create your personal access token:

https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/creating-a-personal-access-token

Args:

bundle_name: bundle name.

version: version name of the target bundle, if None, the latest version will be used.

repo: it should be in the form of "repo_owner/repo_name/".

tag: the tag name of the release.

auth_token: github personal access token.

Returns:

a dictionary that contains the bundle's information.

"""

bundles_info = _get_all_bundles_info(repo=repo, tag=tag, auth_token=auth_token)

if bundle_name not in bundles_info:

raise ValueError(f"bundle: {bundle_name} is not existing.")

bundle_info = bundles_info[bundle_name]

if version is None:

version = sorted(bundle_info.keys())[-1]

if version not in bundle_info:

raise ValueError(f"version: {version} of bundle: {bundle_name} is not existing.")

return bundle_info[version]

def run(

run_id: str | None = None,

init_id: str | None = None,

final_id: str | None = None,

meta_file: str | Sequence[str] | None = None,

config_file: str | Sequence[str] | None = None,

logging_file: str | None = None,

tracking: str | dict | None = None,

args_file: str | None = None,

**override: Any,

) -> None:

"""

Specify `config_file` to run monai bundle components and workflows.

Typical usage examples:

.. code-block:: bash

# Execute this module as a CLI entry:

python -m monai.bundle run --meta_file <meta path> --config_file <config path>

# Execute with specified `run_id=training`:

python -m monai.bundle run training --meta_file <meta path> --config_file <config path>

# Execute with all specified `run_id=runtest`, `init_id=inittest`, `final_id=finaltest`:

python -m monai.bundle run --run_id runtest --init_id inittest --final_id finaltest ...

# Override config values at runtime by specifying the component id and its new value:

python -m monai.bundle run --net#input_chns 1 ...

# Override config values with another config file `/path/to/another.json`:

python -m monai.bundle run --net %/path/to/another.json ...

# Override config values with part content of another config file:

python -m monai.bundle run --net %/data/other.json#net_arg ...

# Set default args of `run` in a JSON / YAML file, help to record and simplify the command line.

# Other args still can override the default args at runtime:

python -m monai.bundle run --args_file "/workspace/data/args.json" --config_file <config path>

Args:

run_id: ID name of the expected config expression to run, default to "run".

to run the config, the target config must contain this ID.

init_id: ID name of the expected config expression to initialize before running, default to "initialize".

it's optional for both configs and this `run` function.

final_id: ID name of the expected config expression to finalize after running, default to "finalize".

it's optional for both configs and this `run` function.

meta_file: filepath of the metadata file, if it is a list of file paths, the content of them will be merged.

Default to None.

config_file: filepath of the config file, if `None`, must be provided in `args_file`.

if it is a list of file paths, the content of them will be merged.

logging_file: config file for `logging` module in the program. for more details:

https://docs.python.org/3/library/logging.config.html#logging.config.fileConfig.

Default to None.

tracking: if not None, enable the experiment tracking at runtime with optionally configurable and extensible.

If "mlflow", will add `MLFlowHandler` to the parsed bundle with default tracking settings where a set of

common parameters shown below will be added and can be passed through the `override` parameter of this method.

- ``"output_dir"``: the path to save mlflow tracking outputs locally, default to "<bundle root>/eval".

- ``"tracking_uri"``: uri to save mlflow tracking outputs, default to "/output_dir/mlruns".

- ``"experiment_name"``: experiment name for this run, default to "monai_experiment".

- ``"run_name"``: the name of current run.

- ``"save_execute_config"``: whether to save the executed config files. It can be `False`, `/path/to/artifacts`

or `True`. If set to `True`, will save to the default path "<bundle_root>/eval". Default to `True`.

If other string, treat it as file path to load the tracking settings.

If `dict`, treat it as tracking settings.

Will patch the target config content with `tracking handlers` and the top-level items of `configs`.

for detailed usage examples, please check the tutorial:

https://github.com/Project-MONAI/tutorials/blob/main/experiment_management/bundle_integrate_mlflow.ipynb.

args_file: a JSON or YAML file to provide default values for `run_id`, `meta_file`,

`config_file`, `logging`, and override pairs. so that the command line inputs can be simplified.

override: id-value pairs to override or add the corresponding config content.

e.g. ``--net#input_chns 42``, ``--net %/data/other.json#net_arg``.

"""