Added name attribute to HeadType (#176)

luxonis · Oct 3, 2024 · 46e2348 · 46e2348
1 parent 9bda0c5
commit 46e2348
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 3 deletions.
diff --git a/luxonis_ml/nn_archive/config.py b/luxonis_ml/nn_archive/config.py
@@ -6,7 +6,7 @@
 
 from .model import Model
 
-CONFIG_VERSION = Literal["1.0"]
+CONFIG_VERSION = Literal["1.0", "1.1"]
 
 
 class Config(BaseModelExtraForbid):

diff --git a/luxonis_ml/nn_archive/config_building_blocks/base_models/head.py b/luxonis_ml/nn_archive/config_building_blocks/base_models/head.py
@@ -16,6 +16,8 @@
 class Head(BaseModel, ABC):
     """Represents head of a model.
 
+    @type name: str | None
+    @ivar name: Optional name of the head.
     @type parser: str
     @ivar parser: Name of the parser responsible for processing the models output.
     @type outputs: List[str] | None
@@ -27,6 +29,7 @@ class Head(BaseModel, ABC):
     @ivar metadata: Metadata of the parser.
     """
 
+    name: Optional[str] = Field(None, description="Optional name of the head.")
     parser: str = Field(
         description="Name of the parser responsible for processing the models output."
     )

diff --git a/requirements-dev.txt b/requirements-dev.txt
@@ -1,7 +1,8 @@
+onnx>=1.14.0,<=1.16.2
 pre-commit>=3.2.1
 pytest-cov>=4.1.0
 pytest-dependency>=0.6.0
 pytest-subtests>=0.12.1
 pytest-md>=0.2.0
 gdown>=4.7.1
-coverage-badge>=1.1.0
+coverage-badge>=1.1.0
diff --git a/tests/test_nn_archive/heads.py b/tests/test_nn_archive/heads.py
@@ -156,6 +156,7 @@
 
 classification_head = dict(
     Head(
+        name="ClassificationHead",
         parser="Classification",
         outputs=["output"],
         metadata=head_classification_metadata,
@@ -164,6 +165,7 @@
 
 ssd_object_detection_head = dict(
     Head(
+        name="ObjectDetectionSSDHead",
         parser="ObjectDetectionSSD",
         outputs=["boxes"],
         metadata=head_object_detection_ssd_metadata,
@@ -172,6 +174,7 @@
 
 yolo_object_detection_head = dict(
     Head(
+        name="YoloDetectionHead",
         parser="YOLO",
         outputs=["output"],
         metadata=head_yolo_obb_det_metadata,
@@ -180,6 +183,7 @@
 
 yolo_instance_segmentation_head = dict(
     Head(
+        name="YoloInstanceSegHead",
         parser="YOLO",
         outputs=["output"],
         metadata=head_yolo_instance_seg_metadata,
@@ -188,6 +192,7 @@
 
 yolo_keypoint_detection_head = dict(
     Head(
+        name="YoloKeypointDetectionHead",
         parser="YOLO",
         outputs=["output"],
         metadata=head_yolo_keypoint_det_metadata,
@@ -196,6 +201,7 @@
 
 yolo_obb_detection_head = dict(
     Head(
+        name="YoloOBBHead",
         parser="YOLO",
         outputs=["output"],
         metadata=head_yolo_obb_det_metadata,
@@ -204,6 +210,7 @@
 
 yolo_instance_seg_kpts_head = dict(
     Head(
+        name="YoloInstaceSegKptHead",
         parser="YOLO",
         outputs=["outputs"],
         metadata=head_yolo_instance_seg_kpts_metadata,
@@ -212,6 +219,7 @@
 
 custom_segmentation_head = dict(
     Head(
+        name="SegmentationHead",
         parser="Segmentation",
         outputs=["output"],
         metadata=head_segmentation_metadata,

diff --git a/tests/test_nn_archive/test_nn_archive.py b/tests/test_nn_archive/test_nn_archive.py
@@ -83,7 +83,7 @@ def test_archive_generator(
         archive_name=archive_name,
         save_path="tests/data/test_nn_archive",
         cfg_dict={
-            "config_version": "1.0",
+            "config_version": "1.1",
             "model": {
                 "metadata": {
                     "name": "test_model",