From 9e1a1630e5aca3c3a1bce401c2337992297fa651 Mon Sep 17 00:00:00 2001 From: Petros Toupas Date: Mon, 11 Nov 2024 20:49:23 +0200 Subject: [PATCH] Added dai_type, updated encoding flags, and removed reverse_input_channels flag (#48) --- README.md | 71 +- modelconverter/__main__.py | 33 +- modelconverter/packages/hailo/exporter.py | 2 +- modelconverter/packages/rvc2/exporter.py | 42 +- modelconverter/utils/config.py | 133 +--- modelconverter/utils/nn_archive.py | 111 ++- modelconverter/utils/onnx_tools.py | 26 +- modelconverter/utils/types.py | 1 + requirements.txt | 2 +- shared_with_container/configs/defaults.yaml | 4 +- tests/test_utils/test_config.py | 775 +++++++++++++++++++- 11 files changed, 1028 insertions(+), 172 deletions(-) diff --git a/README.md b/README.md index 315ce36..b4324dd 100644 --- a/README.md +++ b/README.md @@ -21,13 +21,18 @@ Convert your **ONNX** models to a format compatible with any generation of Luxon ## Table of Contents -- [MLOps - Compilation Library](#mlops---compilation-library) +- [ModelConverter - Compilation Library](#modelconverter---compilation-library) + - [Status](#status) - [Table of Contents](#table-of-contents) - [Installation](#installation) + - [System Requirements](#system-requirements) - [Before You Begin](#before-you-begin) - [Instructions](#instructions) - [GPU Support](#gpu-support) - [Running ModelConverter](#running-modelconverter) + - [Encoding Configuration Flags](#encoding-configuration-flags) + - [YAML Configuration File](#yaml-configuration-file) + - [NN Archive Configuration File](#nn-archive-configuration-file) - [Sharing Files](#sharing-files) - [Usage](#usage) - [Examples](#examples) @@ -101,9 +106,58 @@ To enable GPU acceleration for `hailo` conversion, install the [Nvidia Container ## Running ModelConverter -Configuration for the conversion predominantly relies on a `yaml` config file. For reference, see [defaults.yaml](shared_with_container/configs/defaults.yaml) and other examples located in the [shared_with_container/configs](shared_with_container/configs) directory. +There are two main ways to execute configure the conversion process: -However, you have the flexibility to modify specific settings without altering the config file itself. This is done using command line arguments. You provide the arguments in the form of `key value` pairs. For better understanding, see [Examples](#examples). +1. **YAML Config File (Primary Method)**: + The primary way to configure the conversion is through a YAML configuration file. For reference, you can check [defaults.yaml](shared_with_container/configs/defaults.yaml) and other examples located in the [shared_with_container/configs](shared_with_container/configs) directory. +1. **NN Archive**: + Alternatively, you can use an [NN Archive](https://rvc4.docs.luxonis.com/software/ai-inference/nn-archive/#NN%20Archive) as input. An NN Archive includes a model in one of the supported formats—ONNX (.onnx), OpenVINO IR (.xml and .bin), or TensorFlow Lite (.tflite)—alongside a `config.json` file. The config.json file follows a specific configuration format as described in the [NN Archive Configuration Guide](https://rvc4.docs.luxonis.com/software/ai-inference/nn-archive/#NN%20Archive-Configuration). + +**Modifying Settings with Command-Line Arguments**: +In addition to these two configuration methods, you have the flexibility to override specific settings directly via command-line arguments. By supplying `key-value` pairs in the CLI, you can adjust particular settings without explicitly altering the config files (YAML or NN Archive). For further details, refer to the [Examples](#examples) section. + +### Encoding Configuration Flags + +In the conversion process, you have options to control the color encoding format in both the YAML configuration file and the NN Archive configuration. Here’s a breakdown of each available flag: + +#### YAML Configuration File + +The `encoding` flag in the YAML configuration file allows you to specify color encoding as follows: + +- **Single-Value `encoding`**: + Setting encoding to a single value, such as *"RGB"*, *"BGR"*, *"GRAY"*, or *"NONE"*, will automatically apply this setting to both `encoding.from` and `encoding.to`. For example, `encoding: RGB` sets both `encoding.from` and `encoding.to` to *"RGB"* internally. +- **Multi-Value `encoding.from` and `encoding.to`**: + Alternatively, you can explicitly set `encoding.from` and `encoding.to` to different values. For example: + ```yaml + encoding: + from: RGB + to: BGR + ``` + This configuration specifies that the input data is in RGB format and will be converted to BGR format during processing. + +> [!NOTE] +> If the encoding is not specified in the YAML configuration, the default values are set to `encoding.from=RGB` and `encoding.to=BGR`. + +> [!NOTE] +> Certain options can be set **globally**, applying to all inputs of the model, or **per input**. If specified per input, these settings will override the global configuration for that input alone. The options that support this flexibility include `scale_values`, `mean_values`, `encoding`, `data_type`, `shape`, and `layout`. + +#### NN Archive Configuration File + +In the NN Archive configuration, there are two flags related to color encoding control: + +- **`dai_type`**: + Provides a more comprehensive control over the input type compatible with the DAI backend. It is read by DepthAI to automatically configure the processing pipeline, including any necessary modifications to the input image format. +- **`reverse_channels` (Deprecated)**: + Determines the input color format of the model: when set to *True*, the input is considered to be *"RGB"*, and when set to *False*, it is treated as *"BGR"*. This flag is deprecated and will be replaced by the `dai_type` flag in future versions. + +> [!NOTE] +> If neither `dai_type` nor `reverse_channels` the input to the model is considered to be *"RGB"*. + +> [!NOTE] +> If both `dai_type` and `reverse_channels` are provided, the converter will give priority to `dai_type`. + +> [!IMPORTANT] +> Provide mean/scale values in the original color format used during model training (e.g., RGB or BGR). Any necessary channel permutation is handled internally—do not reorder values manually. ### Sharing Files @@ -158,6 +212,7 @@ You can run the built image either manually using the `docker run` command or us 1. Execute the conversion: - If using the `docker run` command: + ```bash docker run --rm -it \ -v $(pwd)/shared_with_container:/app/shared_with_container/ \ @@ -168,11 +223,15 @@ You can run the built image either manually using the `docker run` command or us convert \ --path [ config overrides ] ``` + - If using the `modelconverter` CLI: + ```bash modelconverter convert --path [ config overrides ] ``` + - If using `docker-compose`: + ```bash docker compose run convert ... ``` @@ -200,13 +259,17 @@ Specify all options via the command line without a config file: ```bash modelconverter convert rvc2 input_model models/yolov6n.onnx \ scale_values "[255,255,255]" \ - reverse_input_channels True \ + inputs.0.encoding.from RGB \ + inputs.0.encoding.to BGR \ shape "[1,3,256,256]" \ outputs.0.name out_0 \ outputs.1.name out_1 \ outputs.2.name out_2 ``` +> [!WARNING] +> If you modify the default stages names (`stages.stage_name`) in the configuration file (`config.yaml`), you need to provide the full path to each stage in the command-line arguments. For instance, if a stage name is changed to `stage1`, use `stages.stage1.inputs.0.name` instead of `inputs.0.name`. + ## Multi-Stage Conversion The converter supports multi-stage conversion. This means conversion of multiple diff --git a/modelconverter/__main__.py b/modelconverter/__main__.py index 84fd0a8..6c19919 100644 --- a/modelconverter/__main__.py +++ b/modelconverter/__main__.py @@ -37,7 +37,7 @@ MODELS_DIR, OUTPUTS_DIR, ) -from modelconverter.utils.types import Target +from modelconverter.utils.types import DataType, Encoding, Target logger = logging.getLogger(__name__) @@ -114,7 +114,7 @@ class Format(str, Enum): DevOption: TypeAlias = Annotated[ bool, typer.Option( - help="Builds a new iamge and uses the development docker-compose file." + help="Builds a new image and uses the development docker-compose file." ), ] @@ -212,18 +212,31 @@ def extract_preprocessing( for inp in stage_cfg.inputs: mean = inp.mean_values or [0, 0, 0] scale = inp.scale_values or [1, 1, 1] + encoding = inp.encoding + layout = inp.layout + + dai_type = encoding.to.value + if dai_type != "NONE": + if inp.data_type == DataType.FLOAT16: + type = "F16F16F16" + else: + type = "888" + dai_type += type + dai_type += "i" if layout == "NHWC" else "p" preproc_block = PreprocessingBlock( - reverse_channels=inp.reverse_input_channels, mean=mean, scale=scale, - interleaved_to_planar=False, + reverse_channels=encoding.to == Encoding.RGB, + interleaved_to_planar=layout == "NHWC", + dai_type=dai_type, ) preprocessing[inp.name] = preproc_block inp.mean_values = None inp.scale_values = None - inp.reverse_input_channels = False + inp.encoding.from_ = Encoding.NONE + inp.encoding.to = Encoding.NONE return cfg, preprocessing @@ -477,9 +490,13 @@ def convert( archive_cfg, preprocessing, main_stage, - exporter.inference_model_path - if isinstance(exporter, Exporter) - else exporter.exporters[main_stage].inference_model_path, + ( + exporter.inference_model_path + if isinstance(exporter, Exporter) + else exporter.exporters[ + main_stage + ].inference_model_path + ), ) generator = ArchiveGenerator( archive_name=f"{cfg.name}.{target.value.lower()}", diff --git a/modelconverter/packages/hailo/exporter.py b/modelconverter/packages/hailo/exporter.py index 2709f67..6d0decd 100644 --- a/modelconverter/packages/hailo/exporter.py +++ b/modelconverter/packages/hailo/exporter.py @@ -237,7 +237,7 @@ def _get_alls(self, runner: ClientRunner) -> str: f"{mean_values},{scale_values},{hn_name})" ) - if inp.reverse_input_channels: + if inp.encoding_mismatch: alls.append( f"bgr_to_rgb_{safe_name} = input_conversion(" f"{hn_name},bgr_to_rgb)" diff --git a/modelconverter/packages/rvc2/exporter.py b/modelconverter/packages/rvc2/exporter.py index 7a0817f..a8e73eb 100644 --- a/modelconverter/packages/rvc2/exporter.py +++ b/modelconverter/packages/rvc2/exporter.py @@ -109,25 +109,43 @@ def _export_openvino_ir(self) -> Path: reverse_only=True, ) for inp in self.inputs.values(): - if inp.mean_values is not None and inp.reverse_input_channels: + if inp.mean_values is not None and inp.encoding_mismatch: inp.mean_values = inp.mean_values[::-1] - if inp.scale_values is not None and inp.reverse_input_channels: + if inp.scale_values is not None and inp.encoding_mismatch: inp.scale_values = inp.scale_values[::-1] - inp.reverse_input_channels = False + inp.encoding.from_ = Encoding.BGR + inp.encoding.to = Encoding.BGR + mean_values_str = "" + scale_values_str = "" for name, inp in self.inputs.items(): + # Append mean values in a similar style if inp.mean_values is not None: - self._add_args( - args, - ["--mean_values", f"{name}{_lst_join(inp.mean_values)}"], + if mean_values_str: + mean_values_str += "," + mean_values_str += ( + f"{name}[{', '.join(str(v) for v in inp.mean_values)}]" ) + + # Append scale values in a similar style if inp.scale_values is not None: - self._add_args( - args, - ["--scale_values", f"{name}{_lst_join(inp.scale_values)}"], + if scale_values_str: + scale_values_str += "," + scale_values_str += ( + f"{name}[{', '.join(str(v) for v in inp.scale_values)}]" ) - if inp.reverse_input_channels: - self._add_args(args, ["--reverse_input_channels"]) + # Extend args with mean and scale values if they were collected + if mean_values_str: + args.extend(["--mean_values", mean_values_str]) + if scale_values_str: + args.extend(["--scale_values", scale_values_str]) + + # Append reverse_input_channels flag only once if needed + reverse_input_flag = any( + inp.encoding_mismatch for inp in self.inputs.values() + ) + if reverse_input_flag: + args.append("--reverse_input_channels") self._add_args(args, ["--input_model", self.input_model]) @@ -137,7 +155,7 @@ def _export_openvino_ir(self) -> Path: return self.input_model.with_suffix(".xml") def _check_reverse_channels(self): - reverses = [inp.reverse_input_channels for inp in self.inputs.values()] + reverses = [inp.encoding_mismatch for inp in self.inputs.values()] return all(reverses) or not any(reverses) @staticmethod diff --git a/modelconverter/utils/config.py b/modelconverter/utils/config.py index d5cdeec..08a1bcc 100644 --- a/modelconverter/utils/config.py +++ b/modelconverter/utils/config.py @@ -134,40 +134,12 @@ class InputConfig(OutputConfig): ] = RandomCalibrationConfig() scale_values: Optional[Annotated[List[float], Field(min_length=1)]] = None mean_values: Optional[Annotated[List[float], Field(min_length=1)]] = None - reverse_input_channels: bool = False frozen_value: Optional[Any] = None encoding: EncodingConfig = EncodingConfig() - @model_validator(mode="before") - @classmethod - def _validate_encoding(cls, data: Dict[str, Any]) -> Dict[str, Any]: - encoding = data.get("encoding") - if encoding is None: - return data - if isinstance(encoding, str): - data["encoding"] = {"from": encoding, "to": encoding} - return data - - @model_validator(mode="after") - def _validate_reverse_input_channels(self) -> Self: - if self.reverse_input_channels: - return self - - if self.encoding.from_ == Encoding.NONE: - self.reverse_input_channels = False - return self - - if ( - self.encoding.from_ == Encoding.GRAY - or self.encoding.to == Encoding.GRAY - ): - self.encoding.from_ = self.encoding.to = Encoding.GRAY - self.reverse_input_channels = False - - if self.encoding.from_ != self.encoding.to: - self.reverse_input_channels = True - - return self + @property + def encoding_mismatch(self) -> bool: + return self.encoding.from_ != self.encoding.to @model_validator(mode="after") def _validate_grayscale_inputs(self) -> Self: @@ -186,6 +158,25 @@ def _validate_grayscale_inputs(self) -> Self: return self + @model_validator(mode="before") + @classmethod + def _validate_encoding(cls, data: Dict[str, Any]) -> Dict[str, Any]: + encoding = data.get("encoding") + if encoding is None or encoding == {}: + data["encoding"] = {"from": "RGB", "to": "BGR"} + return data + if isinstance(encoding, str): + data["encoding"] = {"from": encoding, "to": encoding} + if isinstance(encoding, dict): + if ( + "from" in encoding + and encoding["from"] == "GRAY" + or "to" in encoding + and encoding["to"] == "GRAY" + ): + data["encoding"] = {"from": "GRAY", "to": "GRAY"} + return data + @model_validator(mode="before") @classmethod def _random_calibration(cls, data: Dict[str, Any]) -> Dict[str, Any]: @@ -219,63 +210,6 @@ def _parse_values( return [value, value, value] return value - @model_validator(mode="before") - @classmethod - def _validate_deprecated_reverse_channels( - cls, data: Dict[str, Any] - ) -> Dict[str, Any]: - if "reverse_input_channels" not in data: - return data - logger.warning( - "Field `reverse_input_channels` is deprecated. " - "Please use `encoding.from` and `encoding.to` instead." - ) - reverse = data["reverse_input_channels"] - calib = data.get("calibration", {}) or {} - if isinstance(calib, str): - calib_encoding = Encoding.BGR - else: - calib_encoding = calib.get("encoding", Encoding.BGR) - - if reverse: - if calib_encoding == Encoding.GRAY: - raise ValueError( - "Cannot reverse channels for grayscale images." - ) - else: - encoding = {"from": Encoding.RGB, "to": Encoding.BGR} - else: - if calib_encoding == Encoding.GRAY: - encoding = "GRAY" - else: - encoding = {"from": Encoding.BGR, "to": Encoding.BGR} - - data["encoding"] = encoding - return data - - @model_validator(mode="before") - @classmethod - def _validate_deprecated_reverse_calib_encoding( - cls, data: Dict[str, Any] - ) -> Dict[str, Any]: - calib = data.get("calibration", {}) or {} - if isinstance(calib, str): - return data - - calib_encoding = calib.pop("encoding", None) - if calib_encoding is None: - return data - - logger.warning( - "Field `calibration.encoding` is deprecated. Please use `encoding.to` instead." - ) - encoding = data.get("encoding", {}) - if isinstance(encoding, str): - encoding = {"from": encoding, "to": encoding} - encoding["to"] = calib_encoding - data["encoding"] = encoding - return data - class TargetConfig(CustomBaseModel): disable_calibration: bool = False @@ -363,7 +297,6 @@ def _validate_model(cls, data: Dict[str, Any]) -> Dict[str, Any]: data_type = data.pop("data_type", None) shape = data.pop("shape", None) layout = data.pop("layout", None) - reverse_input_channels = data.pop("reverse_input_channels", None) top_level_calibration = data.pop("calibration", {}) input_file_type = InputFileType.from_path(data["input_model"]) @@ -405,16 +338,16 @@ def _validate_model(cls, data: Dict[str, Any]) -> Dict[str, Any]: inp["layout"] = inp.get("layout") or layout inp["data_type"] = inp.get("data_type") or data_type or onnx_dtype inp["encoding"] = inp.get("encoding") or encoding - inp["mean_values"] = inp.get("mean_values") or mean_values - inp["scale_values"] = inp.get("scale_values") or scale_values - - if ( - inp.get("reverse_input_channels") is not None - or reverse_input_channels is not None - ): - inp["reverse_input_channels"] = inp.get( - "reverse_input_channels" - ) or (reverse_input_channels or False) + inp["mean_values"] = ( + inp.get("mean_values") + if inp.get("mean_values") is not None + else mean_values + ) + inp["scale_values"] = ( + inp.get("scale_values") + if inp.get("scale_values") is not None + else scale_values + ) inp_calibration: Dict[str, Any] = inp.get("calibration", {}) if not inp_calibration and not top_level_calibration: @@ -462,7 +395,6 @@ def _validate_model(cls, data: Dict[str, Any]) -> Dict[str, Any]: data[target] = {} data[target]["disable_calibration"] = disable_calibration - return data @model_validator(mode="before") @@ -522,7 +454,6 @@ def _validate_stages(cls, data: Dict[str, Any]) -> Dict[str, Any]: for key, value in extra.items(): if key not in stage: stage[key] = value - return data @model_validator(mode="after") diff --git a/modelconverter/utils/nn_archive.py b/modelconverter/utils/nn_archive.py index 87505ff..19812f3 100644 --- a/modelconverter/utils/nn_archive.py +++ b/modelconverter/utils/nn_archive.py @@ -1,4 +1,5 @@ import json +import logging import tarfile from pathlib import Path from typing import Any, Dict, Optional, Tuple @@ -17,6 +18,9 @@ from modelconverter.utils.constants import MISC_DIR from modelconverter.utils.layout import guess_new_layout, make_default_layout from modelconverter.utils.metadata import get_metadata +from modelconverter.utils.types import DataType, Encoding + +logger = logging.getLogger(__name__) def get_archive_input(cfg: NNArchiveConfig, name: str) -> NNArchiveInput: @@ -64,13 +68,67 @@ def process_nn_archive( for inp in archive_config.model.inputs: reverse = inp.preprocessing.reverse_channels + interleaved_to_planar = inp.preprocessing.interleaved_to_planar + dai_type = inp.preprocessing.dai_type + + layout = inp.layout + encoding = "NONE" if inp.input_type == InputType.IMAGE: - if reverse: - encoding = {"from": "RGB", "to": "BGR"} + if dai_type is not None: + if (reverse and dai_type.startswith("BGR")) or ( + reverse is False and dai_type.startswith("RGB") + ): + logger.warning( + "'reverse_channels' and 'dai_type' are conflicting, using dai_type" + ) + + if dai_type.startswith("RGB"): + encoding = {"from": "RGB", "to": "BGR"} + elif dai_type.startswith("BGR"): + encoding = "BGR" + elif dai_type.startswith("GRAY"): + encoding = "GRAY" + else: + logger.warning("unknown dai_type, using RGB888p") + encoding = {"from": "RGB", "to": "BGR"} + + if (interleaved_to_planar and dai_type.endswith("p")) or ( + interleaved_to_planar is False and dai_type.endswith("i") + ): + logger.warning( + "'interleaved_to_planar' and 'dai_type' are conflicting, using dai_type" + ) + if dai_type.endswith("i"): + layout = "NHWC" + elif dai_type.endswith("p"): + layout = "NCHW" else: - encoding = "BGR" - else: - encoding = "NONE" + if reverse is not None: + logger.warning( + "'reverse_channels' flag is deprecated and will be removed in the future, use 'dai_type' instead" + ) + if reverse: + encoding = {"from": "RGB", "to": "BGR"} + else: + encoding = "BGR" + else: + encoding = {"from": "RGB", "to": "BGR"} + + if interleaved_to_planar is not None: + logger.warning( + "'interleaved_to_planar' flag is deprecated and will be removed in the future, use 'dai_type' instead" + ) + if interleaved_to_planar: + layout = "NHWC" + else: + layout = "NCHW" + channels = ( + inp.shape[layout.index("C")] + if layout and "C" in layout + else None + ) + if channels and channels == 1: + encoding = "GRAY" mean = inp.preprocessing.mean or [0, 0, 0] scale = inp.preprocessing.scale or [1, 1, 1] @@ -79,7 +137,7 @@ def process_nn_archive( { "name": inp.name, "shape": inp.shape, - "layout": inp.layout, + "layout": layout, "data_type": inp.dtype.value, "mean_values": mean, "scale_values": scale, @@ -97,7 +155,7 @@ def process_nn_archive( } ) - main_stage_key = Path(archive_config.model.metadata.path).stem + main_stage_key = archive_config.model.metadata.name config = { "name": main_stage_key, "stages": { @@ -134,7 +192,7 @@ def modelconverter_config_to_nn( cfg = config.stages[main_stage_key] archive_cfg = { - "config_version": CONFIG_VERSION.__args__[-1], # type: ignore + "config_version": CONFIG_VERSION, "model": { "metadata": { "name": model_name.stem, @@ -148,37 +206,43 @@ def modelconverter_config_to_nn( for inp in cfg.inputs: new_shape = model_metadata.input_shapes[inp.name] - # new_dtype = model_metadata.input_dtypes[inp.name] if inp.shape is not None and not any(s == 0 for s in inp.shape): assert inp.layout is not None layout = guess_new_layout(inp.layout, inp.shape, new_shape) else: layout = make_default_layout(new_shape) + dai_type = inp.encoding.to.value + if inp.data_type == DataType.FLOAT16: + type = "F16F16F16" + else: + type = "888" + dai_type += type + dai_type += "i" if layout == "NHWC" else "p" archive_cfg["model"]["inputs"].append( { "name": inp.name, "shape": new_shape, "layout": layout, - # "dtype": new_dtype.value, "dtype": inp.data_type.value, - # "dtype": "float32", "input_type": "image", "preprocessing": { "mean": [0 for _ in inp.mean_values] if inp.mean_values else None, - "scale": [1 for _ in inp.scale_values] - if inp.scale_values - else None, - "reverse_channels": False, - "interleaved_to_planar": False, + "scale": ( + [1 for _ in inp.scale_values] + if inp.scale_values + else None + ), + "reverse_channels": inp.encoding.to == Encoding.RGB, + "interleaved_to_planar": layout == "NHWC", + "dai_type": dai_type, }, } ) for out in cfg.outputs: new_shape = model_metadata.output_shapes[out.name] - # new_dtype = model_metadata.output_dtypes[out.name] if out.shape is not None and not any(s == 0 for s in out.shape): assert out.layout is not None layout = guess_new_layout(out.layout, out.shape, new_shape) @@ -190,9 +254,7 @@ def modelconverter_config_to_nn( "name": out.name, "shape": new_shape, "layout": layout, - # "dtype": new_dtype.value, "dtype": out.data_type.value, - # "dtype": "float32", } ) @@ -225,7 +287,7 @@ def archive_from_model(model_path: Path) -> NNArchiveConfig: metadata = get_metadata(model_path) archive_cfg = { - "config_version": "1.0", + "config_version": CONFIG_VERSION, "model": { "metadata": { "name": model_path.stem, @@ -243,13 +305,14 @@ def archive_from_model(model_path: Path) -> NNArchiveConfig: "name": name, "shape": shape, "layout": make_default_layout(shape), - "dtype": "float32", + "dtype": metadata.input_dtypes[name].value, "input_type": "image", "preprocessing": { "mean": None, "scale": None, - "reverse_channels": False, - "interleaved_to_planar": False, + "reverse_channels": None, + "interleaved_to_planar": None, + "dai_type": None, }, } ) @@ -260,7 +323,7 @@ def archive_from_model(model_path: Path) -> NNArchiveConfig: "name": name, "shape": shape, "layout": make_default_layout(shape), - "dtype": "float32", + "dtype": metadata.output_dtypes[name].value, } ) diff --git a/modelconverter/utils/onnx_tools.py b/modelconverter/utils/onnx_tools.py index 42c6b53..98138a8 100644 --- a/modelconverter/utils/onnx_tools.py +++ b/modelconverter/utils/onnx_tools.py @@ -41,7 +41,7 @@ def onnx_attach_normalization_to_inputs( continue cfg = input_configs[input_name] if ( - not cfg.reverse_input_channels + cfg.encoding.from_ == cfg.encoding.to and cfg.mean_values is None and cfg.scale_values is None ): @@ -70,7 +70,7 @@ def onnx_attach_normalization_to_inputs( last_output = input_name # 1. Reverse channels if needed - if cfg.reverse_input_channels: + if cfg.encoding_mismatch: split_names = [f"split_{i}_{input_name}" for i in range(3)] split_node = helper.make_node( "Split", @@ -89,8 +89,15 @@ def onnx_attach_normalization_to_inputs( new_nodes.append(concat_node) last_output = f"normalized_{input_name}" - # 2. Subtract (mean) if mean_values is not None - if not reverse_only and cfg.mean_values is not None: + # 2. Subtract (mean) if mean_values is not None and not all 0 + if ( + not reverse_only + and cfg.mean_values is not None + and any(v != 0 for v in cfg.mean_values) + ): + if cfg.encoding_mismatch: + cfg.mean_values = cfg.mean_values[::-1] + sub_out = f"sub_out_{input_name}" sub_node = helper.make_node( "Sub", @@ -110,8 +117,15 @@ def onnx_attach_normalization_to_inputs( ) new_initializers.append(mean_tensor) - # 3. Divide (scale) if scale_values is not None - if not reverse_only and cfg.scale_values is not None: + # 3. Divide (scale) if scale_values is not None and not all 1 + if ( + not reverse_only + and cfg.scale_values is not None + and any(v != 1 for v in cfg.scale_values) + ): + if cfg.encoding_mismatch: + cfg.scale_values = cfg.scale_values[::-1] + div_out = f"div_out_{input_name}" div_node = helper.make_node( "Mul", diff --git a/modelconverter/utils/types.py b/modelconverter/utils/types.py index fee1d89..7b17e54 100644 --- a/modelconverter/utils/types.py +++ b/modelconverter/utils/types.py @@ -24,6 +24,7 @@ class DataType(Enum): FLOAT16 = "float16" FLOAT32 = "float32" FLOAT64 = "float64" + INT4 = "int4" INT8 = "int8" INT16 = "int16" INT32 = "int32" diff --git a/requirements.txt b/requirements.txt index 172b92d..40cf181 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,7 @@ Pillow PyYAML gcsfs -luxonis-ml[data,nn_archive] >= 0.4.0 +luxonis-ml[data,nn_archive] >= 0.5.0 onnx<1.17.0 onnxruntime onnxsim diff --git a/shared_with_container/configs/defaults.yaml b/shared_with_container/configs/defaults.yaml index b29c890..095bfbe 100644 --- a/shared_with_container/configs/defaults.yaml +++ b/shared_with_container/configs/defaults.yaml @@ -94,7 +94,9 @@ stages: # shape: [ 1, 3, 256, 256 ] # data_type: float32 (default) # scale_values: [ 255., 255., 255. ] - # reverse_input_channels: false + # encoding: + # from: RGB + # to: BGR # - name: is_training # freeze_value: false # - name: sequence_len diff --git a/tests/test_utils/test_config.py b/tests/test_utils/test_config.py index 8f6f013..156c0ff 100644 --- a/tests/test_utils/test_config.py +++ b/tests/test_utils/test_config.py @@ -1,14 +1,25 @@ +import json +import re import shutil +import tarfile import tempfile from pathlib import Path -from typing import List, Optional +from typing import Dict, List, Optional +import numpy as np import onnx import pytest +from luxonis_ml.nn_archive.config_building_blocks import PreprocessingBlock from onnx import checker, helper from onnx.onnx_pb import TensorProto -from modelconverter.utils.config import Config +from modelconverter.__main__ import extract_preprocessing +from modelconverter.utils.config import Config, EncodingConfig +from modelconverter.utils.nn_archive import ( + modelconverter_config_to_nn, + process_nn_archive, +) +from modelconverter.utils.onnx_tools import onnx_attach_normalization_to_inputs from modelconverter.utils.types import ( DataType, Encoding, @@ -128,6 +139,74 @@ def setup(): shutil.rmtree(DATA_DIR) +def set_nested_config_value( + config: Dict, keys: List[str], values: List[str] +) -> Dict: + for key, value in zip(keys, values): + keys = key.split(".") + current_level = config["model"] + + # Traverse through the keys except for the last one + for k in keys[:-1]: + # Handle integer keys for list indexing + if re.match(r"^\d+$", k): + k = int(k) + + # Move to the next level + current_level = current_level[k] + + # Set the final key to the value + final_key = keys[-1] + current_level[final_key] = value + + return config + + +def create_json( + keys: Optional[List[str]] = None, values: Optional[List[str]] = None +) -> str: + config = { + "config_version": "1.0", + "model": { + "metadata": { + "name": "dummy_model", + "path": "dummy_model.onnx", + "precision": "float32", + }, + "inputs": [ + { + "name": "input0", + "dtype": "float32", + "input_type": "image", + "shape": [1, 3, 64, 64], + "preprocessing": {}, + }, + { + "name": "input1", + "dtype": "float32", + "input_type": "image", + "shape": [1, 3, 128, 128], + "preprocessing": {}, + }, + ], + "outputs": [ + {"name": "output0", "dtype": "float32", "shape": [1, 10]}, + {"name": "output1", "dtype": "float32", "shape": [1, 5, 5, 5]}, + ], + "heads": [], + }, + } + + if keys and values: + config = set_nested_config_value(config, keys, values) + + with tempfile.NamedTemporaryFile(mode="w", delete=False) as f: + f.write(json.dumps(config)) + f.flush() + + return f.name + + def create_yaml(append: str = "") -> str: config = ( f""" @@ -213,7 +292,6 @@ def test_correct(): "layout": "NCHW", "scale_values": [255, 255, 255], "mean_values": [120, 0, 0], - "reverse_input_channels": False, "data_type": DataType.FLOAT32, "frozen_value": None, "encoding": { @@ -231,7 +309,6 @@ def test_correct(): "shape": [1, 3, 256, 256], "layout": "NCHW", "data_type": DataType.FLOAT32, - "reverse_input_channels": True, "mean_values": [256, 256], "scale_values": None, "frozen_value": None, @@ -302,7 +379,6 @@ def test_top_level(): "layout": "NCHW", "scale_values": [255, 255, 255], "mean_values": [123.675, 116.28, 103.53], - "reverse_input_channels": True, "data_type": DataType.FLOAT32, "frozen_value": None, "encoding": DEFAULT_ENCODINGS, @@ -318,7 +394,6 @@ def test_top_level(): "layout": "NCHW", "scale_values": [255, 255, 255], "mean_values": [123.675, 116.28, 103.53], - "reverse_input_channels": True, "data_type": DataType.FLOAT32, "frozen_value": None, "encoding": DEFAULT_ENCODINGS, @@ -399,7 +474,6 @@ def test_top_level_override(): "layout": "NCHW", "scale_values": [1.0, 2.0, 3.0], "mean_values": [4.0, 5.0, 6.0], - "reverse_input_channels": False, "data_type": DataType.FLOAT16, "frozen_value": [7, 8, 9], "encoding": { @@ -418,7 +492,6 @@ def test_top_level_override(): "layout": "NCHW", "scale_values": [255.0, 255.0, 255.0], "mean_values": [123.675, 116.28, 103.53], - "reverse_input_channels": True, "data_type": DataType.FLOAT32, "frozen_value": None, "encoding": DEFAULT_ENCODINGS, @@ -481,7 +554,6 @@ def test_no_top_level(): "layout": "NCHW", "scale_values": [1.0, 2.0, 3.0], "mean_values": [4.0, 5.0, 6.0], - "reverse_input_channels": False, "data_type": DataType.FLOAT16, "frozen_value": [7, 8, 9], "encoding": { @@ -532,7 +604,6 @@ def test_missing(key: str, value: str): @pytest.mark.parametrize( "key, value", [ - ("reverse_input_channels", "5"), ("inputs.0.encoding.from", "RGBA"), ("mean_values", "scale"), ("scale_values", "[1,2,dog]"), @@ -552,6 +623,686 @@ def test_incorrect_type(key: str, value: str): ) +@pytest.mark.parametrize( + "keys, values", + [ + ( + [], + [], + ), + ( + ["encoding"], + ["BGR"], + ), + ( + [ + "inputs.0.name", + "inputs.0.encoding", + "inputs.1.name", + "inputs.1.encoding.from", + "inputs.1.encoding.to", + ], + ["input0", "BGR", "input1", "RGB", "BGR"], + ), + ( + [ + "inputs.0.name", + "inputs.0.encoding.from", + "inputs.1.name", + ], + ["input0", "RGB", "input1"], + ), + ( + [ + "encoding", + "mean_values", + "scale_values", + "inputs.0.name", + "inputs.1.name", + "inputs.1.encoding.from", + "inputs.1.mean_values", + ], + ["BGR", 0, 255, "input0", "input1", "RGB", 127], + ), + ( + [ + "inputs.0.name", + "inputs.0.encoding", + "inputs.0.mean_values", + "inputs.0.scale_values", + "inputs.1.name", + "inputs.1.encoding.from", + "inputs.1.encoding.to", + "inputs.1.mean_values", + "inputs.1.scale_values", + ], + [ + "input0", + "BGR", + 0, + 1, + "input1", + "RGB", + "BGR", + [123.675, 116.28, 103.53], + [58.395, 57.12, 57.375], + ], + ), + ( + [ + "encoding", + "mean_values", + "scale_values", + "inputs.0.name", + "inputs.0.encoding", + "inputs.0.mean_values", + "inputs.0.scale_values", + "inputs.1.name", + "inputs.1.encoding.from", + "inputs.1.encoding.to", + "inputs.1.mean_values", + "inputs.1.scale_values", + ], + [ + "RGB", + 0, + 255, + "input0", + "BGR", + 0, + 1, + "input1", + "RGB", + "BGR", + [123.675, 116.28, 103.53], + [58.395, 57.12, 57.375], + ], + ), + ], +) +def test_modified_onnx(keys: List[str], values: List[str]): + overrides = {keys[i]: values[i] for i in range(len(keys))} + overrides["input_model"] = str(DATA_DIR / "dummy_model.onnx") + config = Config.get_config( + None, + overrides, + ) + inputs = next(iter(config.stages.values())).inputs + input_configs = {inp.name: inp for inp in inputs} + + modified_onnx_path = onnx_attach_normalization_to_inputs( + DATA_DIR / "dummy_model.onnx", + DATA_DIR / "dummy_model_modified.onnx", + input_configs, + reverse_only=False, + ) + modified_onnx = onnx.load(modified_onnx_path) + + reverse_inputs = {inp.name: False for inp in inputs} + normalize_inputs = {inp.name: [False, False] for inp in inputs} + for node in modified_onnx.graph.node: + if node.op_type == "Split": + reverse_inputs[node.input[0]] = True + elif node.op_type == "Sub": + inp_name = node.input[1].split("mean_")[1] + mean_tensor = onnx.numpy_helper.to_array( + next( + t + for t in modified_onnx.graph.initializer + if t.name == node.input[1] + ) + ) + assert np.allclose( + np.squeeze(mean_tensor), + np.array(input_configs[inp_name].mean_values), + ) + normalize_inputs[inp_name][0] = True + elif node.op_type == "Mul": + inp_name = node.input[1].split("scale_")[1] + scale_tensor = onnx.numpy_helper.to_array( + next( + t + for t in modified_onnx.graph.initializer + if t.name == node.input[1] + ) + ) + assert np.allclose( + np.squeeze(scale_tensor), + 1 / np.array(input_configs[inp_name].scale_values), + ) + normalize_inputs[inp_name][1] = True + + for inp, norm in reverse_inputs.items(): + if norm: + assert input_configs[inp].encoding_mismatch + else: + assert not input_configs[inp].encoding_mismatch + + for inp, norm in normalize_inputs.items(): + if norm[0] and norm[1]: + assert input_configs[inp].mean_values is not None + assert input_configs[inp].scale_values is not None + elif norm[0] and not norm[1]: + assert input_configs[inp].mean_values is not None + assert input_configs[inp].scale_values is None or [1, 1, 1] + elif not norm[0] and norm[1]: + assert input_configs[inp].mean_values is None or [0, 0, 0] + assert input_configs[inp].scale_values is not None + else: + assert input_configs[inp].mean_values is None or [0, 0, 0] + assert input_configs[inp].scale_values is None or [1, 1, 1] + + +@pytest.mark.parametrize( + "keys, values, nn_preprocess, expected", + [ + ( + [], + [], + False, + [ + { + "mean": None, + "scale": None, + "reverse_channels": False, + "interleaved_to_planar": False, + "dai_type": "BGR888p", + }, + { + "mean": None, + "scale": None, + "reverse_channels": False, + "interleaved_to_planar": False, + "dai_type": "BGR888p", + }, + ], + ), + ( + [ + "inputs.0.name", + "inputs.0.encoding", + "inputs.0.mean_values", + "inputs.0.scale_values", + "inputs.1.name", + "inputs.1.encoding.to", + ], + ["input0", "BGR", 127, 255, "input1", "RGB"], + False, + [ + { + "mean": [0, 0, 0], + "scale": [1, 1, 1], + "reverse_channels": False, + "interleaved_to_planar": False, + "dai_type": "BGR888p", + }, + { + "mean": None, + "scale": None, + "reverse_channels": True, + "interleaved_to_planar": False, + "dai_type": "RGB888p", + }, + ], + ), + ( + [], + [], + True, + [ + { + "mean": [0, 0, 0], + "scale": [1, 1, 1], + "reverse_channels": False, + "interleaved_to_planar": False, + "dai_type": "BGR888p", + }, + { + "mean": [0, 0, 0], + "scale": [1, 1, 1], + "reverse_channels": False, + "interleaved_to_planar": False, + "dai_type": "BGR888p", + }, + ], + ), + ( + [ + "inputs.0.name", + "inputs.0.encoding", + "inputs.0.mean_values", + "inputs.0.scale_values", + "inputs.1.name", + "inputs.1.encoding.to", + ], + ["input0", "BGR", 127, 255, "input1", "RGB"], + True, + [ + { + "mean": [127, 127, 127], + "scale": [255, 255, 255], + "reverse_channels": False, + "interleaved_to_planar": False, + "dai_type": "BGR888p", + }, + { + "mean": [0, 0, 0], + "scale": [1, 1, 1], + "reverse_channels": True, + "interleaved_to_planar": False, + "dai_type": "RGB888p", + }, + ], + ), + ], +) +def test_output_nn_config_from_yaml( + keys: List[str], + values: List[str], + nn_preprocess: bool, + expected: List[Dict], +): + overrides = {keys[i]: values[i] for i in range(len(keys))} + overrides["input_model"] = str(DATA_DIR / "dummy_model.onnx") + config = Config.get_config( + None, + overrides, + ) + preprocessing = {} + if nn_preprocess: + config, preprocessing = extract_preprocessing(config) + nn_config = modelconverter_config_to_nn( + config, + DATA_DIR / "dummy_model.onnx", + None, + preprocessing, + "dummy_model", + DATA_DIR / "dummy_model.onnx", + ) + + input_0_preprocessing = nn_config.model.inputs[0].preprocessing + input_1_preprocessing = nn_config.model.inputs[1].preprocessing + + expected_0_preprocessing = PreprocessingBlock( + mean=expected[0]["mean"], + scale=expected[0]["scale"], + reverse_channels=expected[0]["reverse_channels"], + interleaved_to_planar=expected[0]["interleaved_to_planar"], + dai_type=expected[0]["dai_type"], + ) + expected_1_preprocessing = PreprocessingBlock( + mean=expected[1]["mean"], + scale=expected[1]["scale"], + reverse_channels=expected[1]["reverse_channels"], + interleaved_to_planar=expected[1]["interleaved_to_planar"], + dai_type=expected[1]["dai_type"], + ) + + assert input_0_preprocessing == expected_0_preprocessing + assert input_1_preprocessing == expected_1_preprocessing + + +@pytest.mark.parametrize( + "keys, values, nn_preprocess, expected", + [ + ( + [], + [], + False, + [ + { + "mean": [0, 0, 0], + "scale": [1, 1, 1], + "reverse_channels": False, + "interleaved_to_planar": False, + "dai_type": "BGR888p", + }, + { + "mean": [0, 0, 0], + "scale": [1, 1, 1], + "reverse_channels": False, + "interleaved_to_planar": False, + "dai_type": "BGR888p", + }, + ], + ), + ( + [ + "inputs.0.preprocessing.reverse_channels", + "inputs.1.preprocessing.reverse_channels", + ], + ["False", "True"], + False, + [ + { + "mean": [0, 0, 0], + "scale": [1, 1, 1], + "reverse_channels": False, + "interleaved_to_planar": False, + "dai_type": "BGR888p", + }, + { + "mean": [0, 0, 0], + "scale": [1, 1, 1], + "reverse_channels": False, + "interleaved_to_planar": False, + "dai_type": "BGR888p", + }, + ], + ), + ( + [ + "inputs.0.preprocessing.dai_type", + "inputs.1.preprocessing.dai_type", + ], + ["BGR888p", "RGB888p"], + False, + [ + { + "mean": [0, 0, 0], + "scale": [1, 1, 1], + "reverse_channels": False, + "interleaved_to_planar": False, + "dai_type": "BGR888p", + }, + { + "mean": [0, 0, 0], + "scale": [1, 1, 1], + "reverse_channels": False, + "interleaved_to_planar": False, + "dai_type": "BGR888p", + }, + ], + ), + ( + [ + "inputs.0.preprocessing.dai_type", + "inputs.0.preprocessing.reverse_channels", + "inputs.0.preprocessing.mean", + "inputs.0.preprocessing.scale", + "inputs.1.preprocessing.dai_type", + "inputs.1.preprocessing.reverse_channels", + "inputs.1.preprocessing.mean", + "inputs.1.preprocessing.scale", + ], + [ + "BGR888p", + "True", + [0, 0, 0], + [255, 255, 255], + "RGB888p", + "False", + [127, 127, 127], + [255, 255, 255], + ], + False, + [ + { + "mean": [0, 0, 0], + "scale": [1, 1, 1], + "reverse_channels": False, + "interleaved_to_planar": False, + "dai_type": "BGR888p", + }, + { + "mean": [0, 0, 0], + "scale": [1, 1, 1], + "reverse_channels": False, + "interleaved_to_planar": False, + "dai_type": "BGR888p", + }, + ], + ), + ( + [ + "inputs.0.preprocessing.dai_type", + "inputs.0.preprocessing.reverse_channels", + "inputs.0.preprocessing.mean", + "inputs.0.preprocessing.scale", + "inputs.1.preprocessing.dai_type", + "inputs.1.preprocessing.reverse_channels", + "inputs.1.preprocessing.mean", + "inputs.1.preprocessing.scale", + ], + [ + "BGR888p", + "True", + [0, 0, 0], + [255, 255, 255], + "RGB888p", + "False", + [127, 127, 127], + [255, 255, 255], + ], + True, + [ + { + "mean": [0, 0, 0], + "scale": [255, 255, 255], + "reverse_channels": False, + "interleaved_to_planar": False, + "dai_type": "BGR888p", + }, + { + "mean": [127, 127, 127], + "scale": [255, 255, 255], + "reverse_channels": False, + "interleaved_to_planar": False, + "dai_type": "BGR888p", + }, + ], + ), + ], +) +def test_output_nn_config_from_nn_archive( + keys: List[str], + values: List[str], + nn_preprocess: bool, + expected: List[Dict], +): + nn_archive_path = DATA_DIR / "dummy_model.tar.xz" + with tempfile.TemporaryDirectory() as tmpdirname: + tar_path = Path(tmpdirname) / "dummy_model.tar.xz" + with tarfile.open(tar_path, "w") as tar: + tar.add( + str(DATA_DIR / "dummy_model.onnx"), arcname="dummy_model.onnx" + ) + tar.add( + create_json(keys=keys, values=values), arcname="config.json" + ) + shutil.copy(tar_path, nn_archive_path) + config, archive_cfg, main_stage = process_nn_archive( + nn_archive_path, overrides=None + ) + preprocessing = {} + if nn_preprocess: + config, preprocessing = extract_preprocessing(config) + nn_config = modelconverter_config_to_nn( + config, + DATA_DIR / "dummy_model.onnx", + archive_cfg, + preprocessing, + main_stage, + DATA_DIR / "dummy_model.onnx", + ) + + input_0_preprocessing = nn_config.model.inputs[0].preprocessing + input_1_preprocessing = nn_config.model.inputs[1].preprocessing + + expected_0_preprocessing = PreprocessingBlock( + mean=expected[0]["mean"], + scale=expected[0]["scale"], + reverse_channels=expected[0]["reverse_channels"], + interleaved_to_planar=expected[0]["interleaved_to_planar"], + dai_type=expected[0]["dai_type"], + ) + expected_1_preprocessing = PreprocessingBlock( + mean=expected[1]["mean"], + scale=expected[1]["scale"], + reverse_channels=expected[1]["reverse_channels"], + interleaved_to_planar=expected[1]["interleaved_to_planar"], + dai_type=expected[1]["dai_type"], + ) + + assert input_0_preprocessing == expected_0_preprocessing + assert input_1_preprocessing == expected_1_preprocessing + + +@pytest.mark.parametrize( + "key, value, expected", + [ + ( + "", + "", + EncodingConfig(**{"from": Encoding.RGB, "to": Encoding.BGR}), + ), + ( + "encoding", + "NONE", + EncodingConfig(**{"from": Encoding.NONE, "to": Encoding.NONE}), + ), + ( + "encoding", + "RGB", + EncodingConfig(**{"from": Encoding.RGB, "to": Encoding.RGB}), + ), + ( + "encoding", + "BGR", + EncodingConfig(**{"from": Encoding.BGR, "to": Encoding.BGR}), + ), + ( + "encoding", + "GRAY", + EncodingConfig(**{"from": Encoding.GRAY, "to": Encoding.GRAY}), + ), + ( + "encoding.from", + "BGR", + EncodingConfig(**{"from": Encoding.BGR, "to": Encoding.BGR}), + ), + ( + "encoding.to", + "RGB", + EncodingConfig(**{"from": Encoding.RGB, "to": Encoding.RGB}), + ), + ], +) +def test_encoding(key: str, value: str, expected: EncodingConfig): + if key and value: + config = Config.get_config( + None, + { + key: value, + "input_model": str(DATA_DIR / "dummy_model.onnx"), + }, + ) + else: + config = Config.get_config( + None, + { + "input_model": str(DATA_DIR / "dummy_model.onnx"), + }, + ) + assert config.get("stages.dummy_model.inputs.0.encoding") == expected + assert config.get("stages.dummy_model.inputs.1.encoding") == expected + + +@pytest.mark.parametrize( + "keys, values, expected", + [ + ( + [], + [], + EncodingConfig(**{"from": Encoding.RGB, "to": Encoding.BGR}), + ), + ( + [ + "inputs.0.preprocessing.reverse_channels", + ], + ["False"], + EncodingConfig(**{"from": Encoding.BGR, "to": Encoding.BGR}), + ), + ( + [ + "inputs.0.preprocessing.reverse_channels", + ], + ["True"], + EncodingConfig(**{"from": Encoding.RGB, "to": Encoding.BGR}), + ), + ( + [ + "inputs.0.preprocessing.dai_type", + ], + ["BGR888p"], + EncodingConfig(**{"from": Encoding.BGR, "to": Encoding.BGR}), + ), + ( + [ + "inputs.0.preprocessing.dai_type", + ], + ["RGB888p"], + EncodingConfig(**{"from": Encoding.RGB, "to": Encoding.BGR}), + ), + ( + [ + "inputs.0.preprocessing.dai_type", + "inputs.0.preprocessing.reverse_channels", + ], + ["BGR888p", "False"], + EncodingConfig(**{"from": Encoding.BGR, "to": Encoding.BGR}), + ), + ( + [ + "inputs.0.preprocessing.dai_type", + "inputs.0.preprocessing.reverse_channels", + ], + ["BGR888p", "True"], + EncodingConfig(**{"from": Encoding.BGR, "to": Encoding.BGR}), + ), + ( + [ + "inputs.0.preprocessing.dai_type", + "inputs.0.preprocessing.reverse_channels", + ], + ["RGB888p", "False"], + EncodingConfig(**{"from": Encoding.RGB, "to": Encoding.BGR}), + ), + ( + [ + "inputs.0.preprocessing.dai_type", + "inputs.0.preprocessing.reverse_channels", + ], + ["RGB888p", "True"], + EncodingConfig(**{"from": Encoding.RGB, "to": Encoding.BGR}), + ), + ( + [ + "inputs.0.preprocessing.dai_type", + ], + ["GRAY8"], + EncodingConfig(**{"from": Encoding.GRAY, "to": Encoding.GRAY}), + ), + ], +) +def test_encoding_nn_archive( + keys: List[str], values: List[str], expected: EncodingConfig +): + nn_archive_path = DATA_DIR / "dummy_model.tar.xz" + with tempfile.TemporaryDirectory() as tmpdirname: + tar_path = Path(tmpdirname) / "dummy_model.tar.xz" + with tarfile.open(tar_path, "w") as tar: + tar.add( + str(DATA_DIR / "dummy_model.onnx"), arcname="dummy_model.onnx" + ) + tar.add( + create_json(keys=keys, values=values), arcname="config.json" + ) + shutil.copy(tar_path, nn_archive_path) + config, _, _ = process_nn_archive(nn_archive_path, overrides=None) + assert config.get("stages.dummy_model.inputs.0.encoding") == expected + + def test_onnx_load(): load_and_compare( None, @@ -570,7 +1321,6 @@ def test_onnx_load(): "layout": "NCHW", "scale_values": None, "mean_values": None, - "reverse_input_channels": True, "data_type": DataType.FLOAT32, "frozen_value": None, "calibration": DEFAULT_CALIBRATION_CONFIG, @@ -582,7 +1332,6 @@ def test_onnx_load(): "layout": "NCHW", "scale_values": None, "mean_values": None, - "reverse_input_channels": True, "data_type": DataType.FLOAT32, "frozen_value": None, "calibration": DEFAULT_CALIBRATION_CONFIG, @@ -624,7 +1373,6 @@ def test_explicit_nones(): "layout": "NCHW", "scale_values": None, "mean_values": None, - "reverse_input_channels": True, "data_type": DataType.FLOAT32, "frozen_value": None, "encoding": DEFAULT_ENCODINGS, @@ -636,7 +1384,6 @@ def test_explicit_nones(): "layout": "NCHW", "scale_values": None, "mean_values": None, - "reverse_input_channels": True, "data_type": DataType.FLOAT32, "frozen_value": None, "encoding": DEFAULT_ENCODINGS,