Function load_net_with_metadata

monai/data/torchscript_utils.py:103–134

Load the module object from the given Torchscript filename or stream, and convert the stored JSON metadata back to a dict object. This will produce an empty dict if the metadata file is not present. Args: filename_prefix_or_stream: filename or file-like stream object. m

(
    filename_prefix_or_stream: str | IO[Any],
    map_location: torch.device | None = None,
    more_extra_files: Sequence[str] = (),
)

def load_net_with_metadata(
    filename_prefix_or_stream: str | IO[Any],
    map_location: torch.device | None = None,
    more_extra_files: Sequence[str] = (),
) -> tuple[torch.nn.Module, dict, dict]:
    """
    Load the module object from the given Torchscript filename or stream, and convert the stored JSON metadata
    back to a dict object. This will produce an empty dict if the metadata file is not present.

    Args:
        filename_prefix_or_stream: filename or file-like stream object.
        map_location: network map location as in `torch.jit.load`.
        more_extra_files: other extra file data names to load from bundle, see `_extra_files` of `torch.jit.load`.
    Returns:
        Triple containing loaded object, metadata dict, and extra files dict containing other file data if present
    """
    extra_files = dict.fromkeys(more_extra_files, "")
    extra_files[METADATA_FILENAME] = ""

    jit_obj = torch.jit.load(filename_prefix_or_stream, map_location, extra_files)

    extra_files = dict(extra_files.items())  # compatibility with ExtraFilesMap

    if METADATA_FILENAME in extra_files:
        json_data = extra_files[METADATA_FILENAME]
        del extra_files[METADATA_FILENAME]
    else:
        json_data = "{}"

    json_data_dict = json.loads(json_data)

    return jit_obj, json_data_dict, extra_files