kornia.onnx#

Run and compose ONNX models with ONNX Runtime: ONNXSequential chains several ONNX graphs (local files, onnx.ModelProto objects, or pre-exported Kornia operators and models from the kornia/ONNX_models Hugging Face repository) into one model with a PyTorch-like call interface. The ONNX support guide walks through installation, chaining, export and running on CUDA.

class kornia.onnx.module.ONNXModule(op, providers=None, session_options=None, cache_dir=None, target_ir_version=None, target_opset_version=None)[source]#

ONNXModule to wrap an ONNX operator.

Parameters:

arg – A variable number of ONNX models (either ONNX ModelProto objects or file paths). For Hugging Face-hosted models, use the format ‘hf://model_name’. Valid model_name can be found on https://huggingface.co/kornia/ONNX_models. Or a URL to the ONNX model.
providers (optional) – A list of execution providers for ONNXRuntime (e.g., [‘CUDAExecutionProvider’, ‘CPUExecutionProvider’]). Default: None
session_options (optional) – Optional ONNXRuntime session options for optimizing the session. Default: None
cache_dir (optional) – The directory where downloaded ONNX models (hf:// names and URLs) are cached locally. Defaults to None, which will use a default kornia.config.hub_onnx_dir directory. Default: None
target_ir_version (optional) – The target IR version to convert to. Default: None
target_opset_version (optional) – The target OPSET version to convert to. Default: None

add_metadata(additional_metadata=None)[source]#

Attach metadata entries to the wrapped ONNX model graph.

Parameters:: additional_metadata (Optional[list[tuple[str, str]]], optional) – Optional list of (key, value) string pairs to store in the ONNX model metadata properties. Default: None
Return type:: ModelProto
Returns:: The wrapped ONNX model graph after metadata has been added.

create_session(providers=None, session_options=None)[source]#

Create an ONNX Runtime session for the wrapped model graph.

Parameters:

providers (optional) – Optional ordered list of ONNX Runtime execution providers, such as "CUDAExecutionProvider" followed by "CPUExecutionProvider". The first available provider is used according to ONNX Runtime’s provider selection rules. Default: None
session_options (optional) – Optional ONNX Runtime session options controlling graph optimization, threading, logging, and other runtime execution settings. Default: None

Returns:

ONNX Runtime inference session bound to self.op, the wrapped ONNX model graph.

export(file_path, **kwargs)[source]#

Serialize the wrapped ONNX model graph to disk.

Parameters:

file_path (str) – Destination path for the exported .onnx file.
kwargs (Any) – Additional keyword arguments forwarded to the underlying ONNX export helper.

Return type:

None

class kornia.onnx.sequential.ONNXSequential(*args, providers=None, session_options=None, io_maps=None, cache_dir=None, auto_ir_version_conversion=False, target_ir_version=None, target_opset_version=None)[source]#

ONNXSequential to chain multiple ONNX operators together.

Parameters:

*args – A variable number of ONNX models (either ONNX ModelProto objects or file paths). For Hugging Face-hosted models, use the format ‘hf://model_name’. Valid model_name can be found on https://huggingface.co/kornia/ONNX_models. Or a URL to the ONNX model.
providers (optional) – A list of execution providers for ONNXRuntime (e.g., [‘CUDAExecutionProvider’, ‘CPUExecutionProvider’]). Default: None
session_options (optional) – Optional ONNXRuntime session options for optimizing the session. Default: None
io_maps (optional) – An optional list of list of tuples specifying input-output mappings for combining models. If None, we assume the default input name and output name are “input” and “output” accordingly, and only one input and output node for each graph. If not None, io_maps[0] shall represent the io_map for combining the first and second ONNX models. Default: None
cache_dir (optional) – The directory where downloaded ONNX models (hf:// names and URLs) are cached locally. Defaults to None, which will use a default kornia.config.hub_onnx_dir directory. Default: None
auto_ir_version_conversion (optional) – If True, convert the models to one IR version and one OPSET version before combining them: IR version 9, and the highest OPSET version among the models, to which the others are converted up. Other versions may be pointed to by target_ir_version and target_opset_version. Default: False
target_ir_version (optional) – The target IR version to convert to. Default: None
target_opset_version (optional) – The target OPSET version to convert to. Default: None

add_metadata(additional_metadata=None)[source]#

Attach metadata entries to the combined ONNX graph.

Parameters:: additional_metadata (Optional[list[tuple[str, str]]], optional) – Optional list of (key, value) pairs to add to the ONNX model metadata properties. Default: None
Return type:: ModelProto
Returns:: Combined ONNX model graph after metadata has been attached.

combine(io_maps=None)[source]#

Combine the stored ONNX operators into one executable graph.

Parameters:: io_maps (list[tuple[str, str]] | None, optional) – Optional list of (source_output, target_input) name pairs describing how adjacent graphs should be connected. When omitted, the helper assumes the default single input/output names used by the ONNX composition utility. Default: None
Return type:: ModelProto
Returns:: Combined ONNX model graph containing all operators in self.operators connected in sequence.

create_session(providers=None, session_options=None)[source]#

Create an ONNX Runtime session for the combined graph.

Parameters:

providers (optional) – Optional ordered list of execution providers used by ONNX Runtime, for example GPU first and CPU as fallback. Default: None
session_options (optional) – Optional ONNX Runtime session options controlling optimization, threading, and runtime behavior. Default: None

Returns:

ONNX Runtime inference session for self._combined_op.

export(file_path, **kwargs)[source]#

Serialize the combined ONNX graph to an .onnx file.

Parameters:

file_path (str) – Destination path where the composed ONNX graph is saved.
kwargs (Any) – Additional keyword arguments forwarded to the export helper.

Return type:

None

class kornia.onnx.utils.ONNXLoader[source]#

Manages ONNX models, handling local caching, downloading from Hugging Face, and loading models.

from kornia.onnx.utils import ONNXLoader

# Load a Hugging Face operator
ONNXLoader.load_model("hf://operators/kornia.color.gray.GrayscaleToRgb")
# Load a local converted/downloaded operator
ONNXLoader.load_model("operators/kornia.color.gray.GrayscaleToRgb")

classmethod list_models()[source]#

List all available ONNX models in the ‘models’ folder of the Hugging Face repository.

Return type:: None

classmethod list_operators()[source]#

List all available ONNX operators in the ‘operators’ folder of the Hugging Face repository.

Return type:: None

classmethod load_config(url, download=True, **kwargs)[source]#

Load JSON config from the specified URL.

Parameters:

url (str) – The URL of the preprocessor config to load.
download (bool, optional) – If True, the config will be downloaded if it’s not already in the local cache. Default: True
kwargs (Any) – Keyword arguments of CachedDownloader.download_to_cache. Only cache_dir applies: the root the config is stored under, at <scheme>/<host>/<directories of the URL's path>/<file name>, with a %3F<hash> directory before the file name when the URL has a query, as for load_model. None or "" means kornia_config.hub_cache_dir. Any other keyword raises TypeError, whatever url is.

Returns:

The loaded preprocessor config.

Return type:

dict[str, Any]

classmethod load_model(model_name, download=True, with_data=False, **kwargs)[source]#

Load an ONNX model from the local cache or downloads it from Hugging Face if necessary.

Parameters:

model_name (str) – The name of the ONNX model or operator. For Hugging Face-hosted models, use the format ‘hf://model_name’. Valid model_name can be found on https://huggingface.co/kornia/ONNX_models. Or a URL to the ONNX model.
download (bool, optional) – If True, the model will be downloaded from Hugging Face if it’s not already in the local cache. Default: True
cache_dir – The root of the cache. An hf:// model is stored under it at the path of its name after the folder (<cache_dir>/a/b.onnx for hf://models/a/b): the kornia/ONNX_models names are fully qualified (kornia.color.gray.RgbToGrayscale), so the folders never share a name. A URL is stored at <cache_dir>/<scheme>/<host>/<directories of the URL's path>/<file name>, with a %3F<hash> directory before the file name when the URL has a query, so that two URLs share a file only when they differ in nothing but credentials, the fragment, ., .. or empty directories, or an .onnx extension that one of them lacks. Defaults to None ("" too): an hf:// model goes to <kornia_config.hub_onnx_dir>/<folder>/, and a URL under kornia_config.hub_onnx_dir. Limitations: on Windows, a URL whose file name is not a valid file name there, or whose cache path is longer than MAX_PATH (260 characters) while long paths are disabled, cannot be cached, and neither can a directory with a reserved name (con); on a case-insensitive file system, URLs that differ only in letter case share a file.
with_data (bool, optional) – If True, the model will be loaded with its .onnx_data weights, stored beside the model. Default: False
**kwargs – Only cache_dir; any other keyword raises TypeError, whatever model_name is.

Returns:

The loaded ONNX model.

Return type:

onnx.ModelProto