From 895499f1d228496f49974fbb86b6a8e562262457 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Kozlovsk=C3=BD?= <martin.kozlovsky@luxonis.com>
Date: Thu, 25 Jul 2024 19:29:05 +0200
Subject: [PATCH] Layout Support (#13)

---
 README.md                                     |   2 +
 modelconverter/__main__.py                    |  14 +-
 modelconverter/packages/base_exporter.py      |   7 +-
 modelconverter/packages/base_inferer.py       |  13 +-
 modelconverter/packages/hailo/exporter.py     |   8 +-
 modelconverter/packages/rvc2/exporter.py      |   6 +-
 modelconverter/packages/rvc3/exporter.py      |   1 +
 modelconverter/utils/__init__.py              |   6 +
 modelconverter/utils/config.py                | 226 +++++-------------
 modelconverter/utils/layout.py                |  81 +++++++
 modelconverter/utils/metadata.py              | 216 +++++++++++++++++
 modelconverter/utils/nn_archive.py            | 121 ++++++----
 modelconverter/utils/types.py                 |  33 +++
 requirements.txt                              |   3 +-
 shared_with_container/configs/defaults.yaml   |   7 +-
 shared_with_container/configs/yolov6n.yaml    |   4 +
 tests/conftest.py                             |   2 +-
 .../test_packages/test_cross_format_export.py |   2 +-
 tests/test_utils/test_config.py               |  24 +-
 tests/test_utils/test_layout.py               |  28 +++
 20 files changed, 578 insertions(+), 226 deletions(-)
 create mode 100644 modelconverter/utils/layout.py
 create mode 100644 modelconverter/utils/metadata.py
 create mode 100644 tests/test_utils/test_layout.py

diff --git a/README.md b/README.md
index 27e0d22..808a158 100644
--- a/README.md
+++ b/README.md
@@ -73,11 +73,13 @@ Requires `hailo_ai_sw_suite_2024-04:1` docker image to be present on the system.
    ```bash
    docker build -f docker/<package>/Dockerfile.public -t luxonis/modelconverter-<package>:latest .
    ```
+
 1. For easier use, you can install the ModelConverter CLI. You can install it from PyPI using the following command:
 
    ```bash
    pip install modelconv
    ```
+
    For usage instructions, see `modelconverter --help`.
 
 ### GPU Support
diff --git a/modelconverter/__main__.py b/modelconverter/__main__.py
index 6b5743e..3b945bd 100644
--- a/modelconverter/__main__.py
+++ b/modelconverter/__main__.py
@@ -221,7 +221,7 @@ def infer(
         ),
     ],
     path: PathOption,
-    output_dir: OutputDirOption = None,
+    output_dir: OutputDirOption,
     stage: Annotated[
         Optional[str],
         typer.Option(
@@ -251,11 +251,10 @@ def infer(
         try:
             mult_cfg, _, _ = get_configs(path, opts)
             cfg = mult_cfg.get_stage_config(stage)
-            output_path = get_output_dir_name(
-                target, mult_cfg.name, output_dir
-            )
             Inferer = get_inferer(target)
-            Inferer.from_config(model_path, input_path, output_path, cfg).run()
+            Inferer.from_config(
+                model_path, input_path, Path(output_dir), cfg
+            ).run()
         except Exception:
             logger.exception("Encountered an unexpected error!")
             exit(2)
@@ -430,6 +429,8 @@ def convert(
             if not isinstance(out_models, list):
                 out_models = [out_models]
             if to == Format.NN_ARCHIVE:
+                from modelconverter.packages.base_exporter import Exporter
+
                 logger.info("Converting to NN archive")
                 assert main_stage is not None
                 if len(out_models) > 1:
@@ -442,6 +443,9 @@ def convert(
                     archive_cfg,
                     preprocessing,
                     main_stage,
+                    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/base_exporter.py b/modelconverter/packages/base_exporter.py
index 653c661..5b9e6ee 100644
--- a/modelconverter/packages/base_exporter.py
+++ b/modelconverter/packages/base_exporter.py
@@ -4,7 +4,7 @@
 from importlib.metadata import version
 from logging import getLogger
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union, cast
+from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
 import onnx
@@ -85,10 +85,9 @@ def __init__(
             logger.warning(
                 f"Random calibration is being used for input '{name}'."
             )
-            shape = cast(List[int], inp.shape)
             dest = self.intermediate_outputs_dir / "random" / name
             dest.mkdir(parents=True)
-            if shape is None or not all(isinstance(dim, int) for dim in shape):
+            if inp.shape is None:
                 exit_with(
                     ValueError(
                         f"Random calibration requires shape to be specified for input '{name}'."
@@ -96,7 +95,7 @@ def __init__(
                 )
 
             for i in range(calib.max_images):
-                arr = np.random.normal(calib.mean, calib.std, shape)
+                arr = np.random.normal(calib.mean, calib.std, inp.shape)
                 arr = np.clip(arr, calib.min_value, calib.max_value)
 
                 arr = arr.astype(calib.data_type.as_numpy_dtype())
diff --git a/modelconverter/packages/base_inferer.py b/modelconverter/packages/base_inferer.py
index 9d92c06..b76e974 100644
--- a/modelconverter/packages/base_inferer.py
+++ b/modelconverter/packages/base_inferer.py
@@ -42,13 +42,22 @@ def __post_init__(self):
     def from_config(
         cls, model_path: str, src: Path, dest: Path, config: SingleStageConfig
     ):
+        for container, typ_name in zip(
+            [config.inputs, config.outputs], ["input", "output"]
+        ):
+            for node in container:
+                if node.shape is None:
+                    raise ValueError(
+                        f"Shape for {typ_name} '{node.name}' must be provided."
+                    )
+
         return cls(
             model_path=resolve_path(model_path, Path.cwd()),
             src=src,
             dest=dest,
-            in_shapes={inp.name: inp.shape for inp in config.inputs},
+            in_shapes={inp.name: inp.shape for inp in config.inputs},  # type: ignore
             in_dtypes={inp.name: inp.data_type for inp in config.inputs},
-            out_shapes={out.name: out.shape for out in config.outputs},
+            out_shapes={out.name: out.shape for out in config.outputs},  # type: ignore
             out_dtypes={out.name: out.data_type for out in config.outputs},
             resize_method={
                 inp.name: inp.calibration.resize_method
diff --git a/modelconverter/packages/hailo/exporter.py b/modelconverter/packages/hailo/exporter.py
index 1d3c978..2709f67 100644
--- a/modelconverter/packages/hailo/exporter.py
+++ b/modelconverter/packages/hailo/exporter.py
@@ -63,7 +63,8 @@ def __init__(self, config: SingleStageConfig, output_dir: Path):
         self.optimization_level = config.hailo.optimization_level
         self.compression_level = config.hailo.compression_level
         self.batch_size = config.hailo.batch_size
-        self.early_stop = config.hailo.early_stop
+        self.disable_compilation = config.hailo.disable_compilation
+        self._alls: List[str] = []
         self.hw_arch = config.hailo.hw_arch
         if not tf.config.list_physical_devices("GPU"):
             logger.error(
@@ -113,12 +114,13 @@ def export(self) -> Path:
         har_path = self.input_model.with_suffix(".har")
         runner.save_har(har_path)
         if self._disable_calibration:
+            self._inference_model_path = har_path
             return har_path
 
         quantized_har_path = self._calibrate(har_path)
         self._inference_model_path = Path(quantized_har_path)
-        if self.early_stop:
-            logger.info("Early stop enabled. Skipping compilation.")
+        if self.disable_compilation:
+            logger.warning("Compilation disabled, skipping compilation.")
             copy_path = Path(quantized_har_path).parent / (
                 Path(quantized_har_path).stem + "_copy.har"
             )
diff --git a/modelconverter/packages/rvc2/exporter.py b/modelconverter/packages/rvc2/exporter.py
index 6e9c3cc..60abdd6 100644
--- a/modelconverter/packages/rvc2/exporter.py
+++ b/modelconverter/packages/rvc2/exporter.py
@@ -21,9 +21,9 @@
 
 logger = getLogger(__name__)
 
-COMPILE_TOOL: Final[str] = (
-    f'{env["INTEL_OPENVINO_DIR"]}/tools/compile_tool/compile_tool'
-)
+COMPILE_TOOL: Final[
+    str
+] = f'{env["INTEL_OPENVINO_DIR"]}/tools/compile_tool/compile_tool'
 
 DEFAULT_SUPER_SHAVES: Final[int] = 8
 
diff --git a/modelconverter/packages/rvc3/exporter.py b/modelconverter/packages/rvc3/exporter.py
index ccdd59f..a28fdc3 100644
--- a/modelconverter/packages/rvc3/exporter.py
+++ b/modelconverter/packages/rvc3/exporter.py
@@ -62,6 +62,7 @@ def export(self) -> Path:
                 xml_path,
                 output_dir=str(output_dir),
             )
+            self._inference_model_path = calibrated_xml_path
             output_path = (
                 self.output_dir
                 / f"{self.model_name}-{self.target.name.lower()}-int8"
diff --git a/modelconverter/utils/__init__.py b/modelconverter/utils/__init__.py
index 21f91fe..549616f 100644
--- a/modelconverter/utils/__init__.py
+++ b/modelconverter/utils/__init__.py
@@ -19,6 +19,8 @@
     upload_file_to_remote,
 )
 from .image import read_calib_dir, read_image
+from .layout import guess_new_layout, make_default_layout
+from .metadata import Metadata, get_metadata
 from .nn_archive import (
     get_archive_input,
     modelconverter_config_to_nn,
@@ -49,4 +51,8 @@
     "get_docker_image",
     "docker_exec",
     "in_docker",
+    "guess_new_layout",
+    "make_default_layout",
+    "Metadata",
+    "get_metadata",
 ]
diff --git a/modelconverter/utils/config.py b/modelconverter/utils/config.py
index 7ab3bff..91716c3 100644
--- a/modelconverter/utils/config.py
+++ b/modelconverter/utils/config.py
@@ -12,11 +12,13 @@
     field_validator,
     model_validator,
 )
-from typing_extensions import Annotated, Self, TypeAlias
+from typing_extensions import Annotated, Self
 
 from modelconverter.utils.calibration_data import download_calibration_data
 from modelconverter.utils.constants import MODELS_DIR
 from modelconverter.utils.filesystem_utils import resolve_path
+from modelconverter.utils.layout import make_default_layout
+from modelconverter.utils.metadata import get_metadata
 from modelconverter.utils.types import (
     DataType,
     Encoding,
@@ -27,9 +29,6 @@
 
 logger = logging.getLogger(__name__)
 
-FileInfoType: TypeAlias = Dict[
-    str, Tuple[Optional[List[Optional[int]]], Optional[DataType]]
-]
 
 NAMED_VALUES = {
     "imagenet": {
@@ -91,16 +90,35 @@ class RandomCalibrationConfig(CustomBaseModel):
 
 class OutputConfig(CustomBaseModel):
     name: str
-    shape: Optional[List[Optional[int]]] = None
+    shape: Optional[List[int]] = None
+    layout: Optional[str] = None
     data_type: DataType = DataType.FLOAT32
 
-    @field_validator("data_type", mode="before")
-    @staticmethod
-    def _default_data_type(value: Any) -> DataType:
-        """Parses the data_type from the config."""
-        if value is None:
-            return DataType.FLOAT32
-        return DataType(value)
+    @model_validator(mode="before")
+    @classmethod
+    def _make_default_layout(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        shape = data.get("shape")
+        layout = data.get("layout")
+        if shape is None and layout is not None:
+            raise ValueError("`layout` cannot be provided without `shape`.")
+        elif shape is None:
+            return data
+        if layout is None:
+            layout = make_default_layout(shape)
+        data["layout"] = layout.upper()
+        return data
+
+    @model_validator(mode="after")
+    def validate_layout(self) -> Self:
+        if self.shape is None:
+            return self
+        assert self.layout is not None
+        if len(self.layout) != len(self.shape):
+            raise ValueError(
+                f"Length of `layout` ({len(self.layout)}) must match "
+                f"length of `shape` ({len(self.shape)})."
+            )
+        return self
 
 
 class EncodingConfig(CustomBaseModel):
@@ -250,7 +268,7 @@ class HailoConfig(TargetConfig):
     optimization_level: Literal[-100, 0, 1, 2, 3, 4] = 2
     compression_level: Literal[0, 1, 2, 3, 4, 5] = 2
     batch_size: int = 8
-    early_stop: bool = False
+    disable_compilation: bool = False
     alls: List[str] = []
     hw_arch: Literal[
         "hailo8", "hailo8l", "hailo8r", "hailo10h", "hailo15h", "hailo15m"
@@ -325,29 +343,23 @@ def _validate_model(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         encoding = data.pop("encoding", {})
         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 = _detect_input_file_type(data["input_model"])
         data["input_file_type"] = input_file_type
-        file_inputs: FileInfoType = {}
-        file_outputs: FileInfoType = {}
-        if input_file_type == InputFileType.ONNX:
-            file_inputs, file_outputs = _get_onnx_info(data["input_model"])
-        if input_file_type == InputFileType.TFLITE:
-            file_inputs, file_outputs = _get_tflite_info(data["input_model"])
-        elif input_file_type == InputFileType.IR:
-            file_inputs, file_outputs = _get_ir_info(
-                data["input_bin"], data["input_model"]
-            )
+        metadata = get_metadata(Path(data["input_model"]))
 
         inputs = data.get("inputs")
         if not inputs:
-            inputs = [{"name": cast(Any, name)} for name in file_inputs.keys()]
+            inputs = [
+                {"name": cast(Any, name)} for name in metadata.input_shapes
+            ]
         outputs = data.get("outputs")
         if not outputs:
             outputs = [
-                {"name": cast(Any, name)} for name in file_outputs.keys()
+                {"name": cast(Any, name)} for name in metadata.output_shapes
             ]
 
         for inp in inputs:
@@ -356,17 +368,22 @@ def _validate_model(cls, data: Dict[str, Any]) -> Dict[str, Any]:
                     f"Unable to determine name for input: `{inp}`."
                 )
             inp_name = str(inp["name"])
-            if inp_name not in file_inputs:
+            if inp_name not in metadata.input_shapes:
                 tensor_shape, tensor_dtype = _get_onnx_inter_info(
                     data["input_model"], inp_name
                 )
-                file_inputs[inp_name] = tensor_shape, tensor_dtype
+                metadata.input_shapes[inp_name] = tensor_shape  # type: ignore
+                metadata.input_dtypes[inp_name] = tensor_dtype  # type: ignore
                 logger.warning(
                     f"Input `{inp_name}` is not present in inputs of the ONNX model. "
                     f"Assuming it is an intermediate node."
                 )
-            onnx_shape, onnx_dtype = file_inputs[inp_name]
+            onnx_shape, onnx_dtype = (
+                metadata.input_shapes[inp_name],
+                metadata.input_dtypes[inp_name],
+            )
             inp["shape"] = inp.get("shape") or shape or onnx_shape
+            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
@@ -394,7 +411,7 @@ def _validate_model(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         for out in outputs:
             out_name = str(out["name"])
             if (
-                out_name not in file_outputs
+                out_name not in metadata.output_shapes
                 and out.get("data_type") is None
                 and out.get("shape") is None
             ):
@@ -402,8 +419,11 @@ def _validate_model(cls, data: Dict[str, Any]) -> Dict[str, Any]:
                     data["input_model"], out_name
                 )
                 onnx_shape, onnx_dtype = tensor_shape, tensor_dtype
-            elif out_name in file_outputs:
-                onnx_shape, onnx_dtype = file_outputs[out_name]
+            elif out_name in metadata.output_shapes:
+                onnx_shape, onnx_dtype = (
+                    metadata.output_shapes[out_name],
+                    metadata.output_dtypes[out_name],
+                )
             else:
                 onnx_shape, onnx_dtype = None, None
             out["shape"] = out.get("shape") or onnx_shape
@@ -560,129 +580,9 @@ def _extract_bin_xml_from_ir(ir_path: Any) -> Tuple[Path, Path]:
     return bin_path, xml_path
 
 
-def _get_tflite_info(
-    tflite_path: Path,
-) -> Tuple[FileInfoType, FileInfoType]:
-    """Reads names, shapes, and data types for all inputs and outputs of the provided
-    TFLite model.
-
-    Args:
-        tflite_path (Path): Path to the TFLite model.
-
-    Returns:
-        Tuple[FileInfoType, FileInfoType]: (input_info, output_info) where the keys
-        are the input names and the values are tuples of (shape, DataType).
-    """
-
-    import tflite
-
-    with open(tflite_path, "rb") as f:
-        data = f.read()
-
-    # Load the model
-    model = tflite.Model.GetRootAsModel(data, 0)
-
-    # Get the subgraph (model usually contains only one subgraph)
-    subgraph = model.Subgraphs(0)
-    if subgraph is None:
-        raise ValueError("Failed to load TFLite model.")
-
-    input_info = {}
-    output_info = {}
-
-    for i in range(subgraph.InputsLength()):
-        tensor = subgraph.Tensors(subgraph.Inputs(i))
-        input_info[tensor.Name().decode("utf-8")] = (  # type: ignore
-            tensor.ShapeAsNumpy().tolist(),  # type: ignore
-            DataType.from_tensorflow_dtype(tensor.Type()),  # type: ignore
-        )
-
-    for i in range(subgraph.OutputsLength()):
-        tensor = subgraph.Tensors(subgraph.Outputs(i))
-        output_info[tensor.Name().decode("utf-8")] = (  # type: ignore
-            tensor.ShapeAsNumpy().tolist(),  # type: ignore
-            DataType.from_tensorflow_dtype(tensor.Type()),  # type: ignore
-        )
-
-    return input_info, output_info
-
-
-def _get_onnx_info(onnx_path: Path) -> Tuple[FileInfoType, FileInfoType]:
-    """Reads names, shapes and data types for all inputs and outputs of the provided
-    ONNX model.
-
-    Args:
-        onnx_path (Path): Path to the ONNX model.
-
-    Returns:
-        Tuple[FileInfoType, FileInfoType]: (input_info, output_info) where the keys
-        are the input names and the values are tuples of (shape, DataType).
-    """
-
-    try:
-        model = onnx.load(str(onnx_path))
-    except Exception as e:
-        raise ValueError(f"Failed to load ONNX model: `{onnx_path}`") from e
-
-    input_info = {}
-    output_info = {}
-
-    for inp in model.graph.input:
-        shape = [dim.dim_value for dim in inp.type.tensor_type.shape.dim]
-        dtype = DataType.from_onnx_dtype(inp.type.tensor_type.elem_type)
-        input_info[inp.name] = (shape, dtype)
-
-    for output in model.graph.output:
-        shape = [dim.dim_value for dim in output.type.tensor_type.shape.dim]
-        dtype = DataType.from_onnx_dtype(output.type.tensor_type.elem_type)
-        output_info[output.name] = (shape, dtype)
-
-    return input_info, output_info
-
-
-def _get_ir_info(
-    bin_path: Path, xml_path: Path
-) -> Tuple[FileInfoType, FileInfoType]:
-    """Reads names, shapes and data types for all inputs and outputs of the provided IR
-    model (bin and xml).
-
-    Args:
-        bin_path (Path): Path to the OpenVINO binary weights of the model.
-        xml_path (Path): Path to the OpenVINO XML definition of the model.
-
-    Returns:
-        Tuple[FileInfoType, FileInfoType]: (input_info, output_info) where the keys
-        are the input names and the values are tuples of (shape, DataType).
-    """
-
-    from openvino.runtime import Core
-
-    ie = Core()
-    try:
-        model = ie.read_model(model=str(xml_path), weights=str(bin_path))
-    except Exception as e:
-        raise ValueError(
-            f"Failed to load IR model: `{bin_path}` and `{xml_path}`"
-        ) from e
-
-    input_info = {}
-    output_info = {}
-
-    for inp in model.inputs:
-        name = list(inp.names)[0]
-        dtype = DataType.from_ir_dtype(inp.element_type.get_type_name())
-        input_info[name] = (inp.shape, dtype)
-    for output in model.outputs:
-        name = list(output.names)[0]
-        dtype = DataType.from_ir_dtype(output.element_type.get_type_name())
-        output_info[name] = (output.shape, dtype)
-
-    return input_info, output_info
-
-
 def _get_onnx_node_info(
     model_path: Path, node_name: str
-) -> Tuple[List[Optional[int]], DataType]:
+) -> Tuple[List[int], DataType]:
     onnx_model = onnx.load(str(model_path))
     graph = onnx_model.graph
 
@@ -701,9 +601,13 @@ def _get_onnx_node_info(
         )
 
     shape = [
-        dim.dim_value if dim.dim_value > 0 else None
-        for dim in output_value_info.type.tensor_type.shape.dim
+        dim.dim_value for dim in output_value_info.type.tensor_type.shape.dim
     ]
+    if any(dim == 0 for dim in shape):
+        raise ValueError(
+            "Dynamic shapes are not supported. "
+            f"Shape of node '{node_name}' is {shape}."
+        )
     data_type = output_value_info.type.tensor_type.elem_type
 
     return shape, DataType.from_onnx_dtype(data_type)
@@ -711,14 +615,16 @@ def _get_onnx_node_info(
 
 def _get_onnx_tensor_info(
     model_path: Union[Path, str], tensor_name: str
-) -> Tuple[List[Optional[int]], DataType]:
+) -> Tuple[List[int], DataType]:
     model = onnx.load(str(model_path))
 
     def extract_tensor_info(tensor_type):
-        shape = [
-            dim.dim_value if dim.dim_value > 0 else None
-            for dim in tensor_type.shape.dim
-        ]
+        shape = [dim.dim_value for dim in tensor_type.shape.dim]
+        if any(dim == 0 for dim in shape):
+            raise ValueError(
+                "Dynamic shapes are not supported. "
+                f"Shape of tensor '{tensor_name}' is {shape}."
+            )
         return shape, DataType.from_onnx_dtype(tensor_type.elem_type)
 
     for tensor in chain(model.graph.input, model.graph.output):
@@ -740,7 +646,7 @@ def extract_tensor_info(tensor_type):
 
 def _get_onnx_inter_info(
     model_path: Path, name: str
-) -> Tuple[Optional[List[Optional[int]]], Optional[DataType]]:
+) -> Tuple[Optional[List[int]], Optional[DataType]]:
     try:
         logger.info(
             f"Attempting to find shape and data type for tensor '{name}'."
diff --git a/modelconverter/utils/layout.py b/modelconverter/utils/layout.py
new file mode 100644
index 0000000..b802b23
--- /dev/null
+++ b/modelconverter/utils/layout.py
@@ -0,0 +1,81 @@
+from typing import List
+
+
+def make_default_layout(shape: List[int]) -> str:
+    """Creates a default layout for the given shape.
+
+    Tries to guess most common layouts for the given shape pattern.
+    Otherwise, uses the first free letter of the alphabet for each dimension.
+
+    Example::
+        >>> make_default_layout([1, 3, 256, 256])
+        >>> "NCHW"
+        >>> make_default_layout([1, 19, 7, 8])
+        >>> "NABC"
+    """
+    layout = []
+    i = 0
+    if shape[0] == 1:
+        layout.append("N")
+        i += 1
+    if len(shape) - i == 3:
+        if shape[i] < shape[i + 1] and shape[i] < shape[i + 2]:
+            return "".join(layout + ["C", "H", "W"])
+        elif shape[-1] < shape[-2] and shape[-1] < shape[-3]:
+            return "".join(layout + ["H", "W", "C"])
+    i = 0
+    while len(layout) < len(shape):
+        # Starting with "C" for more sensible defaults
+        letter = chr(ord("A") + (i + 2) % 26)
+        if letter not in layout:
+            layout.append(letter)
+        i += 1
+    return "".join(layout)
+
+
+def guess_new_layout(
+    old_layout: str, old_shape: List[int], new_shape: List[int]
+) -> str:
+    """Tries to guess the layout of the new shape.
+
+    The new shape must contain the same elements as the old one.
+    If two values are the same, the order of their labels will be preserved.
+
+    Example::
+
+        >>> old_shape = [1, 3, 256, 256]
+        >>> old_layout = "NCHW"
+        >>> guess_new_layout(old_layout, old_shape, [1, 256, 256, 3])
+        >>> "NHWC"
+
+    @type old_layout: str
+    @param old_layout: Old layout
+
+    @type old_shape: List[int]
+    @param old_shape: Old shape
+
+    @type other: List[int]
+    @param other: New shape to guess the layout of
+
+    @rtype: str
+    @return: Lettercode representation of the new layout
+    """
+    if len(new_shape) != len(old_layout):
+        raise ValueError(
+            "The length of the new shape must be the same as the old one."
+        )
+    if sorted(old_shape) != sorted(new_shape):
+        raise ValueError(
+            "The new shape must contain the same elements as the old one."
+        )
+    old_shape_tuples = list(zip(old_layout, old_shape))
+
+    new_layout = []
+    for dim in new_shape:
+        for i, (old_label, old_dim) in enumerate(old_shape_tuples):
+            if old_dim == dim:
+                new_layout.append(old_label)
+                old_shape_tuples.pop(i)
+                break
+
+    return "".join(new_layout)
diff --git a/modelconverter/utils/metadata.py b/modelconverter/utils/metadata.py
new file mode 100644
index 0000000..57aa07f
--- /dev/null
+++ b/modelconverter/utils/metadata.py
@@ -0,0 +1,216 @@
+import io
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List
+
+import onnx
+import pandas as pd
+
+from modelconverter.utils.subprocess import subprocess_run
+from modelconverter.utils.types import DataType
+
+
+@dataclass
+class Metadata:
+    input_shapes: Dict[str, List[int]]
+    input_dtypes: Dict[str, DataType]
+    output_shapes: Dict[str, List[int]]
+    output_dtypes: Dict[str, DataType]
+
+
+def get_metadata(model_path: Path) -> Metadata:
+    suffix = model_path.suffix
+    if suffix == ".dlc":
+        return _get_metadata_dlc(model_path)
+    if suffix == ".onnx":
+        return _get_metadata_onnx(model_path)
+    if suffix in [".xml", ".bin"]:
+        if suffix == ".xml":
+            xml_path = model_path
+            bin_path = model_path.with_suffix(".bin")
+        else:
+            bin_path = model_path
+            xml_path = model_path.with_suffix(".xml")
+        return _get_metadata_ir(bin_path, xml_path)
+    if suffix in [".hef", ".har"]:
+        return _get_metadata_hailo(model_path)
+    if suffix == ".tflite":
+        return _get_metadata_tflite(model_path)
+    raise ValueError(f"Unsupported model format: {suffix}")
+
+
+def _get_metadata_dlc(model_path: Path) -> Metadata:
+    csv_path = Path("info.csv")
+    subprocess_run(
+        ["snpe-dlc-info", "-i", model_path, "-s", csv_path], silent=True
+    )
+    content = csv_path.read_text()
+
+    metadata = {}
+
+    for typ in ["input", "output"]:
+        start_marker = f"{typ.capitalize()} Name,Dimensions,Type,Encoding Info"
+        if typ == "input":
+            end_marker = "Output Name,Dimensions,Type,Encoding Info"
+        else:
+            end_marker = "Total parameters:"
+        start_index = content.find(start_marker)
+        end_index = content.find(end_marker, start_index)
+
+        relevant_csv_part = content[start_index:end_index].strip()
+        df = pd.read_csv(io.StringIO(relevant_csv_part))
+        metadata[f"{typ}_shapes"] = {
+            str(row[f"{typ.capitalize()} Name"]): list(
+                map(int, str(row["Dimensions"]).split(","))
+            )
+            for _, row in df.iterrows()
+        }
+        metadata[f"{typ}_dtypes"] = {
+            str(row[f"{typ.capitalize()} Name"]): DataType.from_dlc_dtype(
+                row["Type"]  # type: ignore
+            )
+            for _, row in df.iterrows()
+        }
+
+    return Metadata(**metadata)
+
+
+def _get_metadata_ir(bin_path: Path, xml_path: Path) -> Metadata:
+    from openvino.runtime import Core
+
+    ie = Core()
+    try:
+        model = ie.read_model(model=str(xml_path), weights=str(bin_path))
+    except Exception as e:
+        raise ValueError(
+            f"Failed to load IR model: `{bin_path}` and `{xml_path}`"
+        ) from e
+
+    input_shapes = {}
+    input_dtypes = {}
+    output_shapes = {}
+    output_dtypes = {}
+
+    for inp in model.inputs:
+        name = list(inp.names)[0]
+        input_shapes[name] = inp.shape
+        input_dtypes[name] = DataType.from_ir_dtype(
+            inp.element_type.get_type_name()
+        )
+    for output in model.outputs:
+        name = list(output.names)[0]
+        output_shapes[name] = output.shape
+        output_dtypes[name] = DataType.from_ir_dtype(
+            output.element_type.get_type_name()
+        )
+
+    return Metadata(
+        input_shapes=input_shapes,
+        input_dtypes=input_dtypes,
+        output_shapes=output_shapes,
+        output_dtypes=output_dtypes,
+    )
+
+
+def _get_metadata_onnx(onnx_path: Path) -> Metadata:
+    try:
+        model = onnx.load(str(onnx_path))
+    except Exception as e:
+        raise ValueError(f"Failed to load ONNX model: `{onnx_path}`") from e
+
+    input_shapes = {}
+    input_dtypes = {}
+    output_shapes = {}
+    output_dtypes = {}
+
+    for inp in model.graph.input:
+        shape = [dim.dim_value for dim in inp.type.tensor_type.shape.dim]
+        input_shapes[inp.name] = shape
+        input_dtypes[inp.name] = DataType.from_onnx_dtype(
+            inp.type.tensor_type.elem_type
+        )
+
+    for output in model.graph.output:
+        shape = [dim.dim_value for dim in output.type.tensor_type.shape.dim]
+        output_shapes[output.name] = shape
+        output_dtypes[output.name] = DataType.from_onnx_dtype(
+            output.type.tensor_type.elem_type
+        )
+
+    return Metadata(
+        input_shapes=input_shapes,
+        input_dtypes=input_dtypes,
+        output_shapes=output_shapes,
+        output_dtypes=output_dtypes,
+    )
+
+
+def _get_metadata_tflite(model_path: Path) -> Metadata:
+    import tflite
+
+    with open(model_path, "rb") as f:
+        data = f.read()
+
+    subgraph = tflite.Model.GetRootAsModel(data, 0).Subgraphs(0)
+
+    if subgraph is None:
+        raise ValueError("Failed to load TFLite model.")
+
+    input_shapes = {}
+    input_dtypes = {}
+    output_shapes = {}
+    output_dtypes = {}
+
+    for i in range(subgraph.InputsLength()):
+        tensor = subgraph.Tensors(subgraph.Inputs(i))
+        input_shapes[tensor.Name().decode("utf-8")] = (  # type: ignore
+            tensor.ShapeAsNumpy().tolist()  # type: ignore
+        )
+        input_dtypes[tensor.Name().decode("utf-8")] = (  # type: ignore
+            DataType.from_tensorflow_dtype(tensor.Type())  # type: ignore
+        )
+
+    for i in range(subgraph.OutputsLength()):
+        tensor = subgraph.Tensors(subgraph.Outputs(i))
+        output_shapes[tensor.Name().decode("utf-8")] = (  # type: ignore
+            tensor.ShapeAsNumpy().tolist()  # type: ignore
+        )
+        output_dtypes[tensor.Name().decode("utf-8")] = (  # type: ignore
+            DataType.from_tensorflow_dtype(tensor.Type())  # type: ignore
+        )
+
+    return Metadata(
+        input_shapes=input_shapes,
+        input_dtypes=input_dtypes,
+        output_shapes=output_shapes,
+        output_dtypes=output_dtypes,
+    )
+
+
+def _get_metadata_hailo(model_path: Path) -> Metadata:
+    from modelconverter.packages.hailo.exporter import ClientRunner
+
+    input_shapes = {}
+    input_dtypes = {}
+    output_shapes = {}
+    output_dtypes = {}
+    runner = ClientRunner(hw_arch="hailo8", har=str(model_path))
+    for params in runner.get_hn_dict()["layers"].values():
+        if params["type"] in ["input_layer", "output_layer"]:
+            name = params["original_names"][0]
+            shape = params["input_shapes"][0]
+            if shape[0] == -1:
+                shape[0] = 1
+            if params["type"] == "input_layer":
+                input_shapes[name] = shape
+                input_dtypes[name] = None
+            else:
+                output_shapes[name] = shape
+                output_dtypes[name] = None
+
+    return Metadata(
+        input_shapes=input_shapes,
+        input_dtypes=input_dtypes,
+        output_shapes=output_shapes,
+        output_dtypes=output_dtypes,
+    )
diff --git a/modelconverter/utils/nn_archive.py b/modelconverter/utils/nn_archive.py
index 70fad36..42e995d 100644
--- a/modelconverter/utils/nn_archive.py
+++ b/modelconverter/utils/nn_archive.py
@@ -14,6 +14,8 @@
 
 from modelconverter.utils.config import Config
 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
 
 
 def get_archive_input(cfg: NNArchiveConfig, name: str) -> NNArchiveInput:
@@ -76,6 +78,7 @@ def process_nn_archive(
             {
                 "name": inp.name,
                 "shape": inp.shape,
+                "layout": inp.layout,
                 "data_type": inp.dtype.value,
                 "mean_values": mean,
                 "scale_values": scale,
@@ -85,7 +88,12 @@ def process_nn_archive(
 
     for out in archive_config.model.outputs:
         main_stage_config["outputs"].append(
-            {"name": out.name, "data_type": out.dtype.value}
+            {
+                "name": out.name,
+                "shape": out.shape,
+                "layout": out.layout,
+                "data_type": out.dtype.value,
+            }
         )
 
     main_stage_key = Path(archive_config.model.metadata.path).stem
@@ -117,51 +125,80 @@ def modelconverter_config_to_nn(
     orig_nn: Optional[NNArchiveConfig],
     preprocessing: Dict[str, PreprocessingBlock],
     main_stage_key: str,
+    model_path: Path,
 ) -> NNArchiveConfig:
     is_multistage = len(config.stages) > 1
+    model_metadata = get_metadata(model_path)
 
     cfg = config.stages[main_stage_key]
-    archive_cfg = NNArchiveConfig(
-        **{
-            "config_version": "1.0",
-            "model": {
-                "metadata": {
-                    "name": model_name.stem,
-                    "path": str(model_name),
-                },
-                "inputs": [
-                    {
-                        "name": inp.name,
-                        "shape": inp.shape,
-                        "dtype": inp.data_type.value,
-                        "input_type": "image",
-                        "preprocessing": {
-                            "mean": [0] * len(inp.mean_values)
-                            if inp.mean_values
-                            else None,
-                            "scale": [1] * len(inp.scale_values)
-                            if inp.scale_values
-                            else None,
-                            "reverse_channels": False,
-                            "interleaved_to_planar": False,
-                        },
-                    }
-                    for inp in cfg.inputs
-                ],
-                "outputs": [
-                    {
-                        "name": out.name,
-                        "dtype": out.data_type.value,
-                    }
-                    for out in cfg.outputs
-                ],
-                "heads": orig_nn.model.heads if orig_nn else [],
+
+    archive_cfg = {
+        "config_version": "1.0",
+        "model": {
+            "metadata": {
+                "name": model_name.stem,
+                "path": str(model_name),
             },
-        }
-    )
+            "inputs": [],
+            "outputs": [],
+            "heads": orig_nn.model.heads if orig_nn else [],
+        },
+    }
+
+    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:
+            assert inp.layout is not None
+            layout = guess_new_layout(inp.layout, inp.shape, new_shape)
+        else:
+            layout = make_default_layout(new_shape)
+
+        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,
+                },
+            }
+        )
+    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:
+            assert out.layout is not None
+            layout = guess_new_layout(out.layout, out.shape, new_shape)
+        else:
+            layout = make_default_layout(new_shape)
+
+        archive_cfg["model"]["outputs"].append(
+            {
+                "name": out.name,
+                "shape": new_shape,
+                "layout": layout,
+                # "dtype": new_dtype.value,
+                "dtype": out.data_type.value,
+                # "dtype": "float32",
+            }
+        )
+
+    archive = NNArchiveConfig(**archive_cfg)
 
     for name, block in preprocessing.items():
-        nn_inp = get_archive_input(archive_cfg, name)
+        nn_inp = get_archive_input(archive, name)
         nn_inp.preprocessing = block
 
     if is_multistage:
@@ -172,12 +209,12 @@ def modelconverter_config_to_nn(
         post_stage_key = [
             key for key in config.stages if key != main_stage_key
         ][0]
-        if not archive_cfg.model.heads:
+        if not archive.model.heads:
             raise ValueError(
                 "Multistage NN Archives must sxpecify 1 head in the archive config"
             )
-        head = archive_cfg.model.heads[0]
+        head = archive.model.heads[0]
         head.metadata.postprocessor_path = (
             f"{post_stage_key}{model_name.suffix}"
         )
-    return archive_cfg
+    return archive
diff --git a/modelconverter/utils/types.py b/modelconverter/utils/types.py
index 5b3fcf3..7e53c78 100644
--- a/modelconverter/utils/types.py
+++ b/modelconverter/utils/types.py
@@ -32,6 +32,14 @@ class DataType(Enum):
     UINT64 = "uint64"
     BOOLEAN = "boolean"
     STRING = "string"
+    UFXP8 = "ufxp8"
+    UFXP16 = "ufxp16"
+    UFXP32 = "ufxp32"
+    UFXP64 = "ufxp64"
+    FXP8 = "fxp8"
+    FXP16 = "fxp16"
+    FXP32 = "fxp32"
+    FXP64 = "fxp64"
 
     @classmethod
     def from_tensorflow_dtype(cls, dtype: int) -> "DataType":
@@ -55,6 +63,31 @@ def from_tensorflow_dtype(cls, dtype: int) -> "DataType":
             raise ValueError(f"Unsupported TensorFlow data type: `{dtype}`")
         return cls(tensor_types[dtype])
 
+    @classmethod
+    def from_dlc_dtype(cls, dtype: str) -> "DataType":
+        dtype_map = {
+            "Float_16": "float16",
+            "Float_32": "float32",
+            "Float_64": "float64",
+            "Int_8": "int8",
+            "Int_16": "int16",
+            "Int_32": "int32",
+            "Int_64": "int64",
+            "uInt_8": "uint8",
+            "uInt_16": "uint16",
+            "uInt_32": "uint32",
+            "uInt_64": "uint64",
+            "uFxp_8": "ufxp8",
+            "uFxp_16": "ufxp16",
+            "uFxp_32": "ufxp32",
+            "Fxp_8": "fxp8",
+            "Fxp_16": "fxp16",
+            "Fxp_32": "fxp32",
+        }
+        if dtype not in dtype_map:
+            raise ValueError(f"Unsupported DLC data type: `{dtype}`")
+        return cls(dtype_map[dtype])
+
     @classmethod
     def from_onnx_dtype(cls, dtype: int) -> "DataType":
         dtype_map = {
diff --git a/requirements.txt b/requirements.txt
index bd6350d..5947eb3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,7 +1,8 @@
 Pillow
 PyYAML
 gcsfs
-luxonis-ml[data,nn_archive] >= 0.2.3
+# luxonis-ml[data,nn_archive] >= 0.2.3
+luxonis-ml[data,nn_archive] @ git+https://github.com/luxonis/luxonis-ml.git@dev
 onnx
 onnxruntime
 onnxsim
diff --git a/shared_with_container/configs/defaults.yaml b/shared_with_container/configs/defaults.yaml
index 828902c..7ebf38c 100644
--- a/shared_with_container/configs/defaults.yaml
+++ b/shared_with_container/configs/defaults.yaml
@@ -69,10 +69,13 @@ stages:
     # Same options as for scale_values. "imagenet" equivalent to [ 58.395, 57.12, 57.375 ].
     mean_values: ~
 
-    # The input shape of the network in NCHW format. If not provided,
+    # The input shape of the network. If not provided,
     # it will be inferred from the model if possible.
     shape: ~
 
+    # Lettercode representation of the input layout. e.g. NCHW
+    layout: ~
+
     # Remote path where to upload the compiled model. Optional.
     output_remote_url: ~
 
@@ -120,7 +123,7 @@ stages:
       batch_size: 8
 
       # Stop after quantization. Used in tests.
-      early_stop: false
+      disable_compilation: false
 
       # List of additional arguments to pass to the model optimizer.
       alls: []
diff --git a/shared_with_container/configs/yolov6n.yaml b/shared_with_container/configs/yolov6n.yaml
index 1b3e357..e8b4c26 100644
--- a/shared_with_container/configs/yolov6n.yaml
+++ b/shared_with_container/configs/yolov6n.yaml
@@ -11,8 +11,12 @@ calibration:
 inputs:
   - name: images
     shape: [ 1, 3, 416, 416 ]
+    layout: NCHW
 
 outputs:
   - name: output1_yolov6r2
+    layout: NCHW
   - name: output2_yolov6r2
+    layout: NCHW
   - name: output3_yolov6r2
+    layout: NCHW
diff --git a/tests/conftest.py b/tests/conftest.py
index 424c40a..0d1552b 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -199,7 +199,7 @@ def prepare(
             f"input_model {file_url} "
             "hailo.compression_level 0 "
             "hailo.optimization_level 0 "
-            "hailo.early_stop True "
+            "hailo.disable_compilation True "
             f"rvc2.superblob {'false' if 'superblob' not in service else 'true'} "
             "calibration.max_images 30 "
             + ("--to nn_archive" if model_type == "archive" else "")
diff --git a/tests/test_packages/test_cross_format_export.py b/tests/test_packages/test_cross_format_export.py
index 632b92d..774dec8 100644
--- a/tests/test_packages/test_cross_format_export.py
+++ b/tests/test_packages/test_cross_format_export.py
@@ -5,7 +5,7 @@
 
 from modelconverter.utils import subprocess_run
 
-URL_PREFIX: Final[str] = "shared_with_container/configs/"
+URL_PREFIX: Final[str] = "gs://luxonis-test-bucket/modelconverter/"
 
 
 @pytest.mark.parametrize(
diff --git a/tests/test_utils/test_config.py b/tests/test_utils/test_config.py
index 7a97420..93b72c3 100644
--- a/tests/test_utils/test_config.py
+++ b/tests/test_utils/test_config.py
@@ -56,7 +56,7 @@
         "optimization_level": 2,
         "compression_level": 2,
         "batch_size": 8,
-        "early_stop": False,
+        "disable_compilation": False,
         "alls": [],
         "disable_calibration": False,
         "hw_arch": "hailo8",
@@ -86,11 +86,13 @@
         "name": "output0",
         "data_type": DataType.FLOAT32,
         "shape": [1, 10],
+        "layout": "NC",
     },
     {
         "name": "output1",
         "data_type": DataType.FLOAT32,
         "shape": [1, 5, 5, 5],
+        "layout": "NCDE",
     },
 ]
 
@@ -145,6 +147,7 @@ def create_yaml(append: str = "") -> str:
   - name: "input0"
     scale_values: [255,255,255]
     shape: [1, 3, 64, 64]
+    layout: "NCHW"
     calibration:
       max_images: 100
 """
@@ -206,6 +209,7 @@ def test_correct():
                 {
                     "name": "input0",
                     "shape": [1, 3, 64, 64],
+                    "layout": "NCHW",
                     "scale_values": [255, 255, 255],
                     "mean_values": [120, 0, 0],
                     "reverse_input_channels": False,
@@ -224,6 +228,7 @@ def test_correct():
                 {
                     "name": "input1",
                     "shape": [1, 3, 256, 256],
+                    "layout": "NCHW",
                     "data_type": DataType.FLOAT32,
                     "reverse_input_channels": True,
                     "mean_values": [256, 256],
@@ -242,6 +247,7 @@ def test_correct():
                     "name": "output0",
                     "data_type": DataType.FLOAT32,
                     "shape": [1, 10],
+                    "layout": "NC",
                 },
             ],
             **DEFAULT_GENERAL_CONFIG,
@@ -259,7 +265,7 @@ def test_correct():
                 "optimization_level": 3,
                 "compression_level": 3,
                 "batch_size": 4,
-                "early_stop": False,
+                "disable_compilation": False,
                 "alls": [],
                 "hw_arch": "hailo8",
             },
@@ -292,6 +298,7 @@ def test_top_level():
                 {
                     "name": "input0",
                     "shape": [1, 3, 64, 64],
+                    "layout": "NCHW",
                     "scale_values": [255, 255, 255],
                     "mean_values": [123.675, 116.28, 103.53],
                     "reverse_input_channels": True,
@@ -307,6 +314,7 @@ def test_top_level():
                 {
                     "name": "input1",
                     "shape": [1, 3, 64, 64],
+                    "layout": "NCHW",
                     "scale_values": [255, 255, 255],
                     "mean_values": [123.675, 116.28, 103.53],
                     "reverse_input_channels": True,
@@ -325,6 +333,7 @@ def test_top_level():
                     "name": "output1",
                     "data_type": DataType.FLOAT32,
                     "shape": [1, 5, 5, 5],
+                    "layout": "NCDE",
                 },
             ],
             **DEFAULT_GENERAL_CONFIG,
@@ -355,6 +364,8 @@ def test_top_level_override():
             "[1,10]",
             "outputs.0.data_type",
             "float16",
+            "outputs.0.layout",
+            "NA",
             "inputs.0.name",
             "input0",
             "inputs.0.shape",
@@ -384,6 +395,7 @@ def test_top_level_override():
                 {
                     "name": "input0",
                     "shape": [1, 3, 256, 256],
+                    "layout": "NCHW",
                     "scale_values": [1.0, 2.0, 3.0],
                     "mean_values": [4.0, 5.0, 6.0],
                     "reverse_input_channels": False,
@@ -402,6 +414,7 @@ def test_top_level_override():
                 {
                     "name": "input1",
                     "shape": [1, 3, 64, 64],
+                    "layout": "NCHW",
                     "scale_values": [255.0, 255.0, 255.0],
                     "mean_values": [123.675, 116.28, 103.53],
                     "reverse_input_channels": True,
@@ -420,6 +433,7 @@ def test_top_level_override():
                     "name": "883",
                     "data_type": DataType.FLOAT16,
                     "shape": [1, 10],
+                    "layout": "NA",
                 },
             ],
             **DEFAULT_GENERAL_CONFIG,
@@ -463,6 +477,7 @@ def test_no_top_level():
                 {
                     "name": "input0",
                     "shape": [1, 3, 256, 256],
+                    "layout": "NCHW",
                     "scale_values": [1.0, 2.0, 3.0],
                     "mean_values": [4.0, 5.0, 6.0],
                     "reverse_input_channels": False,
@@ -484,6 +499,7 @@ def test_no_top_level():
                     "name": "output0",
                     "data_type": DataType.FLOAT32,
                     "shape": [1, 10],
+                    "layout": "NC",
                 },
             ],
             **DEFAULT_GENERAL_CONFIG,
@@ -552,6 +568,7 @@ def test_onnx_load():
                 {
                     "name": "input0",
                     "shape": [1, 3, 64, 64],
+                    "layout": "NCHW",
                     "scale_values": None,
                     "mean_values": None,
                     "reverse_input_channels": True,
@@ -563,6 +580,7 @@ def test_onnx_load():
                 {
                     "name": "input1",
                     "shape": [1, 3, 128, 128],
+                    "layout": "NCHW",
                     "scale_values": None,
                     "mean_values": None,
                     "reverse_input_channels": True,
@@ -604,6 +622,7 @@ def test_explicit_nones():
                 {
                     "name": "input0",
                     "shape": [1, 3, 64, 64],
+                    "layout": "NCHW",
                     "scale_values": None,
                     "mean_values": None,
                     "reverse_input_channels": True,
@@ -615,6 +634,7 @@ def test_explicit_nones():
                 {
                     "name": "input1",
                     "shape": [1, 3, 128, 128],
+                    "layout": "NCHW",
                     "scale_values": None,
                     "mean_values": None,
                     "reverse_input_channels": True,
diff --git a/tests/test_utils/test_layout.py b/tests/test_utils/test_layout.py
new file mode 100644
index 0000000..778afc3
--- /dev/null
+++ b/tests/test_utils/test_layout.py
@@ -0,0 +1,28 @@
+from modelconverter.utils.layout import guess_new_layout, make_default_layout
+
+
+def test_shape():
+    old_shape = [1, 3, 256, 256]
+    old_layout = "NCHW"
+    assert guess_new_layout(old_layout, old_shape, [3, 256, 256, 1]) == "CHWN"
+    assert guess_new_layout(old_layout, old_shape, [1, 256, 256, 3]) == "NHWC"
+    assert guess_new_layout(old_layout, old_shape, [1, 3, 256, 256]) == "NCHW"
+
+
+def test_shape_complex():
+    old_shape = [1, 2, 3, 4, 5, 5, 5, 6]
+    old_layout = "NABCWHDE"
+    assert (
+        guess_new_layout(old_layout, old_shape, [1, 2, 3, 4, 5, 5, 5, 6])
+        == "NABCWHDE"
+    )
+    assert (
+        guess_new_layout(old_layout, old_shape, [6, 5, 5, 1, 2, 3, 5, 4])
+        == "EWHNABDC"
+    )
+
+
+def test_make_default_layout():
+    assert make_default_layout([1, 3, 256, 256]) == "NCHW"
+    assert make_default_layout([1, 19, 7, 8]) == "NCDE"
+    assert make_default_layout([256, 256, 3]) == "HWC"