Added dai_type to NNArchive

luxonis_ml/nn_archive/config.py

@@ -1,14 +1,12 @@
-from typing import Literal
+import re
 
-from pydantic import Field
+from pydantic import Field, field_validator
 
 from luxonis_ml.utils import BaseModelExtraForbid
 
 from .model import Model
 
-# NOTE: Change in CONFIG_VERSION is a breaking change and should be communicated
-# to all other packages that relly on it before pushed to release
-CONFIG_VERSION = Literal["1.0"]
+CONFIG_VERSION = "1.0"
 
 
 class Config(BaseModelExtraForbid):
@@ -17,17 +15,26 @@ class Config(BaseModelExtraForbid):
     models).
 
     @type config_version: str
-    @ivar config_version: Static variable representing the version of
-        the config scheme.
+    @ivar config_version: String representing config schema version in
+        format 'x.y' where x is major version and y is minor version
     @type model: Model
     @ivar model: A Model object representing the neural network used in
         the archive.
     """
 
-    config_version: CONFIG_VERSION = Field(
-        ...,
-        description="Static variable representing the version of the config scheme.",
+    config_version: str = Field(
+        CONFIG_VERSION,
+        description="String representing config schema version in format 'x.y' where x is major version and y is minor version.",
     )
     model: Model = Field(
         description="A Model object representing the neural network used in the archive."
     )
+
+    @field_validator("config_version")
+    def validate_config_version_format(cls, v):
+        # Regular expression to match 'x.y' where x and y are integers
+        if not re.match(r"^\d+\.\d+$", v):
+            raise ValueError(
+                "'config_version' must be in format 'x.y' where x and y are integers"
+            )
+        return v

luxonis_ml/nn_archive/config_building_blocks/__init__.py

@@ -9,8 +9,6 @@
 from .enums import (
     DataType,
     InputType,
-    ObjectDetectionSubtypeSSD,
-    ObjectDetectionSubtypeYOLO,
 )
 
 __all__ = [
@@ -22,6 +20,4 @@
     "Metadata",
     "DataType",
     "InputType",
-    "ObjectDetectionSubtypeSSD",
-    "ObjectDetectionSubtypeYOLO",
 ]

luxonis_ml/nn_archive/config_building_blocks/base_models/head_metadata.py

@@ -2,8 +2,6 @@
 
 from pydantic import BaseModel, ConfigDict, Field, model_validator
 
-from ..enums import ObjectDetectionSubtypeYOLO
-
 
 class HeadMetadata(BaseModel):
     """Metadata for the basic head. It allows you to specify additional
@@ -138,7 +136,7 @@ class HeadYOLOMetadata(HeadObjectDetectionMetadata, HeadSegmentationMetadata):
     @ivar keypoints_outputs: A list of output names for the keypoints.
     @type angles_outputs: list | None
     @ivar angles_outputs: A list of output names for the angles.
-    @type subtype: C{ObjectDetectionSubtypeYOLO}
+    @type subtype: str
     @ivar subtype: YOLO family decoding subtype (e.g. yolov5, yolov6,
         yolov7 etc.)
     @type n_prototypes: int | None
@@ -174,7 +172,7 @@ class HeadYOLOMetadata(HeadObjectDetectionMetadata, HeadSegmentationMetadata):
         None, description="A list of output names for the angles."
     )
 
-    subtype: ObjectDetectionSubtypeYOLO = Field(
+    subtype: str = Field(
         description="YOLO family decoding subtype (e.g. yolov5, yolov6, yolov7 etc.)."
     )
     n_prototypes: Optional[int] = Field(
@@ -320,19 +318,3 @@ def validate_task_specific_fields(
                 )
 
         return values
-
-    @model_validator(mode="before")
-    def validate_anchors(cls, values):
-        if (
-            "anchors" in values
-            and values["anchors"] is not None
-            and (
-                values["subtype"] == ObjectDetectionSubtypeYOLO.YOLOv6
-                or values["subtype"] == ObjectDetectionSubtypeYOLO.YOLOv6r2
-                or values["subtype"] == ObjectDetectionSubtypeYOLO.YOLOv8
-            )
-        ):
-            raise ValueError(
-                "YOLOv6, YOLOv6r2, and YOLOv8 do not support anchors."
-            )
-        return values

luxonis_ml/nn_archive/config_building_blocks/base_models/input.py

@@ -13,35 +13,43 @@
 class PreprocessingBlock(BaseModelExtraForbid):
     """Represents preprocessing operations applied to the input data.
 
-    @type mean: list
-    @ivar mean: Mean values in channel order. Typically, this is BGR
-        order.
-    @type scale: list
-    @ivar scale: Standardization values in channel order. Typically,
-        this is BGR order.
+    @type mean: list | None
+    @ivar mean: Mean values in channel order. Order depends on the order
+        in which the model was trained on.
+    @type scale: list | None
+    @ivar scale: Standardization values in channel order. Order depends
+        on the order in which the model was trained on.
     @type reverse_channels: bool | None
-    @ivar reverse_channels: If True, color channels are reversed (e.g.
-        BGR to RGB or vice versa).
+    @ivar reverse_channels: If True input to the model is RGB else BGR.
     @type interleaved_to_planar: bool | None
-    @ivar interleaved_to_planar: If True, format is changed from
-        interleaved to planar.
+    @ivar interleaved_to_planar: If True input to the model is
+        interleaved (NHWC) else planar (NCHW).
+    @type dai_type: str | None
+    @ivar dai_type: DepthAI input type which is read by DepthAI to
+        automatically setup the pipeline.
     """
 
     mean: Optional[List[float]] = Field(
         None,
-        description="Mean values in channel order. Typically, this is BGR order.",
+        description="Mean values in channel order. Order depends on the order in which the model was trained on.",
     )
     scale: Optional[List[float]] = Field(
         None,
-        description="Standardization values in channel order. Typically, this is BGR order.",
+        description="Standardization values in channel order. Order depends on the order in which the model was trained on.",
     )
     reverse_channels: Optional[bool] = Field(
         None,
-        description="If True, color channels are reversed (e.g. BGR to RGB or vice versa).",
+        deprecated="Deprecated, use `dai_type` instead.",
+        description="If True input to the model is RGB else BGR.",
     )
     interleaved_to_planar: Optional[bool] = Field(
         None,
-        description="If True, format is changed from interleaved to planar.",
+        deprecated="Deprecated, use `dai_type` instead.",
+        description="If True input to the model is interleaved (NHWC) else planar (NCHW).",
+    )
+    dai_type: Optional[str] = Field(
+        None,
+        description="DepthAI input type which is read by DepthAI to automatically setup the pipeline.",
     )
 
 

luxonis_ml/nn_archive/config_building_blocks/enums/__init__.py

@@ -1,10 +1,7 @@
 from .data_type import DataType
-from .decoding import ObjectDetectionSubtypeSSD, ObjectDetectionSubtypeYOLO
 from .input_type import InputType
 
 __all__ = [
     "DataType",
     "InputType",
-    "ObjectDetectionSubtypeSSD",
-    "ObjectDetectionSubtypeYOLO",
 ]

luxonis_ml/nn_archive/config_building_blocks/enums/data_type.py

@@ -5,8 +5,16 @@ class DataType(Enum):
     """Represents all existing data types used in i/o streams of the
     model."""
 
+    INT4 = "int4"
     INT8 = "int8"
     INT32 = "int32"
+    INT64 = "int64"
     UINT8 = "uint8"
+    UINT16 = "uint16"
+    UINT32 = "uint32"
+    UINT64 = "uint64"
     FLOAT32 = "float32"
     FLOAT16 = "float16"
+    FLOAT64 = "float64"
+    BOOLEAN = "boolean"
+    STRING = "string"

tests/test_nn_archive/heads.py

@@ -7,9 +7,6 @@
     HeadSegmentationMetadata,
     HeadYOLOMetadata,
 )
-from luxonis_ml.nn_archive.config_building_blocks.enums import (
-    ObjectDetectionSubtypeYOLO,
-)
 
 head_metadata = HeadMetadata(
     postprocessor_path="postprocessor.onnx",
@@ -61,7 +58,7 @@
     n_prototypes=None,
     n_keypoints=None,
     is_softmax=None,
-    subtype=ObjectDetectionSubtypeYOLO.YOLOv6,
+    subtype="yolov6",
     postprocessor_path=None,
     yolo_outputs=["feats"],
     mask_outputs=None,
@@ -80,7 +77,7 @@
     n_prototypes=10,
     n_keypoints=None,
     is_softmax=True,
-    subtype=ObjectDetectionSubtypeYOLO.YOLOv6,
+    subtype="yolov6",
     postprocessor_path="postprocessor.onnx",
     yolo_outputs=["feats"],
     mask_outputs=["mask"],
@@ -99,7 +96,7 @@
     n_prototypes=None,
     n_keypoints=21,
     is_softmax=None,
-    subtype=ObjectDetectionSubtypeYOLO.YOLOv6,
+    subtype="yolov6",
     postprocessor_path=None,
     yolo_outputs=["feats"],
     mask_outputs=None,
@@ -118,7 +115,7 @@
     n_prototypes=None,
     n_keypoints=None,
     is_softmax=None,
-    subtype=ObjectDetectionSubtypeYOLO.YOLOv6,
+    subtype="yolov6",
     postprocessor_path="postprocessor.onnx",
     yolo_outputs=["feats"],
     mask_outputs=None,
@@ -137,7 +134,7 @@
     n_prototypes=10,
     n_keypoints=21,
     is_softmax=False,
-    subtype=ObjectDetectionSubtypeYOLO.YOLOv6,
+    subtype="yolov6",
     postprocessor_path="postprocessor.onnx",
     yolo_outputs=["feats"],
     mask_outputs=["mask"],

tests/test_nn_archive/test_nn_archive.py

@@ -132,6 +132,51 @@ def test_archive_generator(
         assert "config.json" in tar.getnames()
 
 
+def test_config_version():
+    from luxonis_ml.nn_archive import Config
+
+    cfg_dict = {
+        "config_version": "1.0",
+        "model": {
+            "metadata": {
+                "name": "test_model",
+                "path": "test_model.onnx",
+            },
+            "inputs": [
+                {
+                    "name": "input",
+                    "shape": [1, 3, 224, 224],
+                    "input_type": "image",
+                    "layout": "nchw",
+                    "dtype": "float32",
+                    "preprocessing": {
+                        "mean": [0.485, 0.456, 0.406],
+                        "scale": [0.229, 0.224, 0.225],
+                        "reverse_channels": False,
+                        "interleaved_to_planar": False,
+                    },
+                }
+            ],
+            "outputs": [
+                {
+                    "name": "output",
+                    "dtype": "float32",
+                }
+            ],
+            "heads": [],
+        },
+    }
+    Config(**cfg_dict)
+    cfg_dict["config_version"] = "1.2"
+    Config(**cfg_dict)
+    cfg_dict["config_version"] = "1.2.2"
+    with pytest.raises(ValidationError):
+        Config(**cfg_dict)
+    cfg_dict["config_version"] = "1.a"
+    with pytest.raises(ValidationError):
+        Config(**cfg_dict)
+
+
 def test_optional_head_name():
     from luxonis_ml.nn_archive.config_building_blocks.base_models.head_metadata import (
         HeadMetadata,