From 92eb83e3ca130780560e69d12bc6246fd16bc895 Mon Sep 17 00:00:00 2001
From: Nikita Sokovnin <nikitoos67@gmail.com>
Date: Mon, 25 Nov 2024 16:01:23 +0100
Subject: [PATCH 01/23] feat: add Archivable interface

---
 luxonis_train/nodes/heads/bisenet_head.py             | 3 ++-
 luxonis_train/nodes/heads/classification_head.py      | 3 ++-
 luxonis_train/nodes/heads/ddrnet_segmentation_head.py | 3 ++-
 luxonis_train/nodes/heads/efficient_bbox_head.py      | 4 +++-
 luxonis_train/nodes/heads/segmentation_head.py        | 3 ++-
 luxonis_train/nodes/interfaces/__init__.py            | 5 +++++
 luxonis_train/nodes/interfaces/archivable.py          | 6 ++++++
 7 files changed, 22 insertions(+), 5 deletions(-)
 create mode 100644 luxonis_train/nodes/interfaces/__init__.py
 create mode 100644 luxonis_train/nodes/interfaces/archivable.py

diff --git a/luxonis_train/nodes/heads/bisenet_head.py b/luxonis_train/nodes/heads/bisenet_head.py
index b1fa57b3..772e390a 100644
--- a/luxonis_train/nodes/heads/bisenet_head.py
+++ b/luxonis_train/nodes/heads/bisenet_head.py
@@ -5,10 +5,11 @@
 from luxonis_train.enums import TaskType
 from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import ConvModule
+from luxonis_train.nodes.interfaces import Archivable
 from luxonis_train.utils import infer_upscale_factor
 
 
-class BiSeNetHead(BaseNode[Tensor, Tensor]):
+class BiSeNetHead(BaseNode[Tensor, Tensor], Archivable):
     in_height: int
     in_width: int
     in_channels: int
diff --git a/luxonis_train/nodes/heads/classification_head.py b/luxonis_train/nodes/heads/classification_head.py
index 74033c2b..ac09b36a 100644
--- a/luxonis_train/nodes/heads/classification_head.py
+++ b/luxonis_train/nodes/heads/classification_head.py
@@ -4,9 +4,10 @@
 
 from luxonis_train.enums import TaskType
 from luxonis_train.nodes.base_node import BaseNode
+from luxonis_train.nodes.interfaces import Archivable
 
 
-class ClassificationHead(BaseNode[Tensor, Tensor]):
+class ClassificationHead(BaseNode[Tensor, Tensor], Archivable):
     in_channels: int
     tasks: list[TaskType] = [TaskType.CLASSIFICATION]
 
diff --git a/luxonis_train/nodes/heads/ddrnet_segmentation_head.py b/luxonis_train/nodes/heads/ddrnet_segmentation_head.py
index 53912d7c..f92c93fe 100644
--- a/luxonis_train/nodes/heads/ddrnet_segmentation_head.py
+++ b/luxonis_train/nodes/heads/ddrnet_segmentation_head.py
@@ -7,12 +7,13 @@
 
 from luxonis_train.enums import TaskType
 from luxonis_train.nodes.base_node import BaseNode
+from luxonis_train.nodes.interfaces import Archivable
 from luxonis_train.utils.general import infer_upscale_factor
 
 logger = logging.getLogger(__name__)
 
 
-class DDRNetSegmentationHead(BaseNode[Tensor, Tensor]):
+class DDRNetSegmentationHead(BaseNode[Tensor, Tensor], Archivable):
     attach_index: int = -1
     in_height: int
     in_width: int
diff --git a/luxonis_train/nodes/heads/efficient_bbox_head.py b/luxonis_train/nodes/heads/efficient_bbox_head.py
index 531a294f..0e6d344d 100644
--- a/luxonis_train/nodes/heads/efficient_bbox_head.py
+++ b/luxonis_train/nodes/heads/efficient_bbox_head.py
@@ -7,6 +7,7 @@
 from luxonis_train.enums import TaskType
 from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import EfficientDecoupledBlock
+from luxonis_train.nodes.interfaces import Archivable
 from luxonis_train.utils import (
     Packet,
     anchors_for_fpn_features,
@@ -18,7 +19,8 @@
 
 
 class EfficientBBoxHead(
-    BaseNode[list[Tensor], tuple[list[Tensor], list[Tensor], list[Tensor]]]
+    BaseNode[list[Tensor], tuple[list[Tensor], list[Tensor], list[Tensor]]],
+    Archivable,
 ):
     in_channels: list[int]
     tasks: list[TaskType] = [TaskType.BOUNDINGBOX]
diff --git a/luxonis_train/nodes/heads/segmentation_head.py b/luxonis_train/nodes/heads/segmentation_head.py
index 31c79ab6..21c40498 100644
--- a/luxonis_train/nodes/heads/segmentation_head.py
+++ b/luxonis_train/nodes/heads/segmentation_head.py
@@ -5,10 +5,11 @@
 from luxonis_train.enums import TaskType
 from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import UpBlock
+from luxonis_train.nodes.interfaces import Archivable
 from luxonis_train.utils import infer_upscale_factor
 
 
-class SegmentationHead(BaseNode[Tensor, Tensor]):
+class SegmentationHead(BaseNode[Tensor, Tensor], Archivable):
     in_height: int
     in_width: int
     in_channels: int
diff --git a/luxonis_train/nodes/interfaces/__init__.py b/luxonis_train/nodes/interfaces/__init__.py
new file mode 100644
index 00000000..3f8bb45a
--- /dev/null
+++ b/luxonis_train/nodes/interfaces/__init__.py
@@ -0,0 +1,5 @@
+from archivable import Archivable
+
+__all__ = [
+    'Archivable',
+]
diff --git a/luxonis_train/nodes/interfaces/archivable.py b/luxonis_train/nodes/interfaces/archivable.py
new file mode 100644
index 00000000..ceb48f9c
--- /dev/null
+++ b/luxonis_train/nodes/interfaces/archivable.py
@@ -0,0 +1,6 @@
+from abc import ABC, abstractmethod
+
+class Archivable(ABC):
+    @abstractmethod
+    def get_metadata(self) -> dict:
+        ...
\ No newline at end of file

From a43a091ef11cf36f3e309aaa21a5641fcc86f780 Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Mon, 25 Nov 2024 15:04:45 +0000
Subject: [PATCH 02/23] fix: archivable export

---
 luxonis_train/nodes/interfaces/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/luxonis_train/nodes/interfaces/__init__.py b/luxonis_train/nodes/interfaces/__init__.py
index 3f8bb45a..5a5c4300 100644
--- a/luxonis_train/nodes/interfaces/__init__.py
+++ b/luxonis_train/nodes/interfaces/__init__.py
@@ -1,4 +1,4 @@
-from archivable import Archivable
+from .archivable import Archivable
 
 __all__ = [
     'Archivable',

From cea5650e32e4f854eeda5b240f5926c3e55cedf9 Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Tue, 26 Nov 2024 15:08:34 +0000
Subject: [PATCH 03/23] feat: add BaseHead class

---
 luxonis_train/core/core.py                |  12 +-
 luxonis_train/core/utils/archive_utils.py | 156 +++++-----------------
 luxonis_train/nodes/heads/__init__.py     |   2 +
 luxonis_train/nodes/heads/base_head.py    |  69 ++++++++++
 luxonis_train/utils/__init__.py           |   2 +
 luxonis_train/utils/general.py            |  12 ++
 6 files changed, 124 insertions(+), 129 deletions(-)
 create mode 100644 luxonis_train/nodes/heads/base_head.py

diff --git a/luxonis_train/core/core.py b/luxonis_train/core/core.py
index 2a3f3678..cc1141a3 100644
--- a/luxonis_train/core/core.py
+++ b/luxonis_train/core/core.py
@@ -704,7 +704,11 @@ def archive(
             return self._archive(path)
 
     def _archive(self, path: str | Path | None = None) -> Path:
-        from .utils.archive_utils import get_heads, get_inputs, get_outputs
+        from .utils.archive_utils import (
+            get_head_configs,
+            get_inputs,
+            get_outputs,
+        )
 
         archive_name = self.cfg.archiver.name or self.cfg.model.name
         archive_save_directory = Path(self.run_save_dir, "archive")
@@ -761,11 +765,9 @@ def _mult(lst: list[float | int]) -> list[float]:
                 }
             )
 
-        heads = get_heads(
-            self.cfg,
+        heads = get_head_configs(
+            self.lightning_module,
             outputs,
-            self.loaders["train"].get_classes(),
-            self.lightning_module.nodes,  # type: ignore
         )
 
         model = {
diff --git a/luxonis_train/core/utils/archive_utils.py b/luxonis_train/core/utils/archive_utils.py
index 6d973785..32faf279 100644
--- a/luxonis_train/core/utils/archive_utils.py
+++ b/luxonis_train/core/utils/archive_utils.py
@@ -9,12 +9,8 @@
 )
 from onnx.onnx_pb import TensorProto
 
-from luxonis_train.config import Config
-from luxonis_train.nodes.base_node import BaseNode
-from luxonis_train.nodes.enums.head_categorization import (
-    ImplementedHeads,
-    ImplementedHeadsIsSoxtmaxed,
-)
+from luxonis_train.models import LuxonisLightningModule
+from luxonis_train.nodes.heads import BaseHead
 
 logger = logging.getLogger(__name__)
 
@@ -104,71 +100,9 @@ def _get_onnx_inputs(onnx_path: Path) -> dict[str, MetadataDict]:
     return inputs
 
 
-def _get_classes(
-    node_name: str, node_task: str | None, classes: dict[str, list[str]]
-) -> list[str]:
-    if not node_task:
-        match node_name:
-            case "ClassificationHead":
-                node_task = "classification"
-            case "EfficientBBoxHead":
-                node_task = "boundingbox"
-            case "SegmentationHead" | "BiSeNetHead" | "DDRNetSegmentationHead":
-                node_task = "segmentation"
-            case "EfficientKeypointBBoxHead":
-                node_task = "keypoints"
-            case _:  # pragma: no cover
-                raise ValueError("Node does not map to a default task.")
-
-    return classes.get(node_task, [])
-
-
-def _get_head_specific_parameters(
-    nodes: dict[str, BaseNode], head_name: str, head_alias: str
-) -> dict:
-    """Get parameters specific to head.
-
-    @type nodes: dict[str, BaseNode]
-    @param nodes: Dictionary of nodes.
-    @type head_name: str
-    @param head_name: Name of the head (e.g. 'EfficientBBoxHead').
-    @type head_alias: str
-    @param head_alias: Alias of the head (e.g. 'detection_head').
-    """
-
-    parameters = {}
-    if head_name == "ClassificationHead":
-        parameters["is_softmax"] = getattr(
-            ImplementedHeadsIsSoxtmaxed, head_name
-        ).value
-    elif head_name == "EfficientBBoxHead":
-        parameters["subtype"] = "yolov6"
-        head_node = nodes[head_alias]
-        parameters["iou_threshold"] = head_node.iou_thres
-        parameters["conf_threshold"] = head_node.conf_thres
-        parameters["max_det"] = head_node.max_det
-    elif head_name in [
-        "SegmentationHead",
-        "BiSeNetHead",
-        "DDRNetSegmentationHead",
-    ]:
-        parameters["is_softmax"] = getattr(
-            ImplementedHeadsIsSoxtmaxed, head_name
-        ).value
-    elif head_name == "EfficientKeypointBBoxHead":
-        # or appropriate subtype
-        head_node = nodes[head_alias]
-        parameters["iou_threshold"] = head_node.iou_thres
-        parameters["conf_threshold"] = head_node.conf_thres
-        parameters["max_det"] = head_node.max_det
-        parameters["n_keypoints"] = head_node.n_keypoints
-    else:  # pragma: no cover
-        raise ValueError("Unknown head name")
-    return parameters
-
-
 def _get_head_outputs(
-    outputs: list[dict], head_name: str, head_type: str
+    outputs: list[dict],
+    head_name: str,
 ) -> list[str]:
     """Get model outputs in a head-specific format.
 
@@ -177,8 +111,6 @@ def _get_head_outputs(
     @type head_name: str
     @param head_name: Type of the head (e.g. 'EfficientBBoxHead') or its
         custom alias.
-    @type head_type: str
-    @param head_name: Type of the head (e.g. 'EfficientBBoxHead').
     @rtype: list[str]
     @return: List of output names.
     """
@@ -192,63 +124,39 @@ def _get_head_outputs(
     return output_names
 
 
-def get_heads(
-    cfg: Config,
+def get_head_configs(
+    lightning_module: LuxonisLightningModule,
     outputs: list[dict],
-    class_dict: dict[str, list[str]],
-    nodes: dict[str, BaseNode],
 ) -> list[dict]:
     """Get model heads.
 
-    @type cfg: Config
-    @param cfg: Configuration object.
+    @type lightning_module: LuxonisLightningModule
+    @param lightning_module: Lightning module.
     @type outputs: list[dict]
-    @param outputs: List of model outputs.
-    @type class_dict: dict[str, list[str]]
-    @param class_dict: Dictionary of classes.
-    @type nodes: dict[str, BaseNode]
-    @param nodes: Dictionary of nodes.
+    @param outputs: List of NN Archive outputs.
+    @rtype: list[dict]
+    @return: List of head configurations.
     """
-    heads = []
+    head_configs = []
     head_names = set()
-    for node in cfg.model.nodes:
-        node_name = node.name
-        node_alias = node.alias or node_name
-        if "aux-segmentation" in node_alias:
+
+    for node_name, node in lightning_module.nodes.items():
+        if not isinstance(node, BaseHead) or node.remove_on_export:
             continue
-        if node_alias in cfg.model.outputs:
-            if node_name in ImplementedHeads.__members__:
-                parser = getattr(ImplementedHeads, node_name).value
-                task = node.task
-                if isinstance(task, dict):
-                    task = str(next(iter(task.values())))
-
-                classes = _get_classes(node_name, task, class_dict)
-
-                export_output_names = nodes[node_alias].export_output_names
-                if export_output_names is not None:
-                    head_outputs = export_output_names
-                else:
-                    head_outputs = _get_head_outputs(
-                        outputs, node_alias, node_name
-                    )
-
-                if node_alias in head_names:
-                    curr_head_name = f"{node_alias}_{len(head_names)}"  # add suffix if name is already present
-                else:
-                    curr_head_name = node_alias
-                head_names.add(curr_head_name)
-                head_dict = {
-                    "name": curr_head_name,
-                    "parser": parser,
-                    "metadata": {
-                        "classes": classes,
-                        "n_classes": len(classes),
-                    },
-                    "outputs": head_outputs,
-                }
-                head_dict["metadata"].update(
-                    _get_head_specific_parameters(nodes, node_name, node_alias)
-                )
-                heads.append(head_dict)
-    return heads
+
+        head_config = node.get_head_config()
+        head_name = (
+            node_name
+            if node_name not in head_names
+            else f"{node_name}_{len(head_names)}"
+        )
+        head_names.add(head_name)
+
+        head_outputs = node.export_output_names or _get_head_outputs(
+            outputs, node_name
+        )
+        head_config.update({"name": head_name, "outputs": head_outputs})
+
+        head_configs.append(head_config)
+
+    return head_configs
diff --git a/luxonis_train/nodes/heads/__init__.py b/luxonis_train/nodes/heads/__init__.py
index af18074d..6e06c7b2 100644
--- a/luxonis_train/nodes/heads/__init__.py
+++ b/luxonis_train/nodes/heads/__init__.py
@@ -1,3 +1,4 @@
+from .base_head import BaseHead
 from .bisenet_head import BiSeNetHead
 from .classification_head import ClassificationHead
 from .ddrnet_segmentation_head import DDRNetSegmentationHead
@@ -7,6 +8,7 @@
 from .segmentation_head import SegmentationHead
 
 __all__ = [
+    "BaseHead",
     "BiSeNetHead",
     "ClassificationHead",
     "EfficientBBoxHead",
diff --git a/luxonis_train/nodes/heads/base_head.py b/luxonis_train/nodes/heads/base_head.py
new file mode 100644
index 00000000..0c4de5dc
--- /dev/null
+++ b/luxonis_train/nodes/heads/base_head.py
@@ -0,0 +1,69 @@
+
+from abc import abstractmethod
+from typing import Generic, Literal
+
+from luxonis_train.nodes.base_node import BaseNode, ForwardInputT, ForwardOutputT
+from luxonis_train.utils import deep_merge_dicts
+
+
+class BaseHead(BaseNode[ForwardInputT, ForwardOutputT], Generic[ForwardInputT, ForwardOutputT]):
+    """Base class for all heads in the model.
+
+    @type parser: str | None
+    @ivar parser: Parser to use for the head.
+
+    @type is_softmaxed: bool | None
+    @ivar is_softmaxed: Whether the head uses softmax or not.
+    """
+
+    parser: Literal[
+        "ClassificationParser",
+        "YOLO",
+        "YoloDetectionNetwork",
+        "SegmentationParser"
+    ] | None = None
+    is_softmaxed: bool | None = None
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def __init_subclass__(cls, **kwargs):
+        """Check if head has all required attributes."""
+        super().__init_subclass__()
+        if cls.parser is None:
+            raise ValueError(f"Head `{cls.__name__}` must define a parser.")
+
+    def get_head_config(self) -> dict:
+        """Get head configuration.
+
+        @rtype: dict
+        @return: Head configuration.
+        """
+        config = self._get_base_head_config()
+        custom_config = self._get_custom_head_config()
+        deep_merge_dicts(config, custom_config)
+        return config
+        
+    def _get_base_head_config(self) -> dict:
+        """Get base head configuration.
+
+        @rtype: dict
+        @return: Base head configuration.
+        """
+        return {
+            "parser": self.parser,
+            "metadata": {
+                "is_softmax": self.is_softmaxed,
+                "classes": self.class_names,
+                "n_classes": self.n_classes,
+            }
+        }
+
+    @abstractmethod
+    def _get_custom_head_config(self) -> dict:
+        """Get custom head configuration.
+
+        @rtype: dict
+        @return: Custom head configuration.
+        """
+        ...
\ No newline at end of file
diff --git a/luxonis_train/utils/__init__.py b/luxonis_train/utils/__init__.py
index 2944dfde..ac6fa1b3 100644
--- a/luxonis_train/utils/__init__.py
+++ b/luxonis_train/utils/__init__.py
@@ -9,6 +9,7 @@
 from .dataset_metadata import DatasetMetadata
 from .exceptions import IncompatibleException
 from .general import (
+    deep_merge_dicts,
     get_with_default,
     infer_upscale_factor,
     make_divisible,
@@ -30,6 +31,7 @@
     "make_divisible",
     "infer_upscale_factor",
     "to_shape_packet",
+    "deep_merge_dicts",
     "get_with_default",
     "safe_download",
     "LuxonisTrackerPL",
diff --git a/luxonis_train/utils/general.py b/luxonis_train/utils/general.py
index 390ddfd2..020cb4bb 100644
--- a/luxonis_train/utils/general.py
+++ b/luxonis_train/utils/general.py
@@ -203,3 +203,15 @@ def clean_url(url: str) -> str:
 def url2file(url: str) -> str:
     """Convert URL to filename, i.e. https://url.com/file.txt?auth -> file.txt."""
     return Path(clean_url(url)).name
+
+
+def deep_merge_dicts(dict1: dict, dict2: dict) -> None:
+    for key, value in dict2.items():
+        if (
+            isinstance(value, dict)
+            and key in dict1
+            and isinstance(dict1[key], dict)
+        ):
+            deep_merge_dicts(dict1[key], value)
+        else:
+            dict1[key] = value

From 73f6451366ef8eb858ab6fc1ffa498bbeae58401 Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Tue, 26 Nov 2024 15:09:55 +0000
Subject: [PATCH 04/23] fix: remove archivable inetrface

---
 luxonis_train/nodes/interfaces/__init__.py   | 5 -----
 luxonis_train/nodes/interfaces/archivable.py | 6 ------
 2 files changed, 11 deletions(-)
 delete mode 100644 luxonis_train/nodes/interfaces/__init__.py
 delete mode 100644 luxonis_train/nodes/interfaces/archivable.py

diff --git a/luxonis_train/nodes/interfaces/__init__.py b/luxonis_train/nodes/interfaces/__init__.py
deleted file mode 100644
index 5a5c4300..00000000
--- a/luxonis_train/nodes/interfaces/__init__.py
+++ /dev/null
@@ -1,5 +0,0 @@
-from .archivable import Archivable
-
-__all__ = [
-    'Archivable',
-]
diff --git a/luxonis_train/nodes/interfaces/archivable.py b/luxonis_train/nodes/interfaces/archivable.py
deleted file mode 100644
index ceb48f9c..00000000
--- a/luxonis_train/nodes/interfaces/archivable.py
+++ /dev/null
@@ -1,6 +0,0 @@
-from abc import ABC, abstractmethod
-
-class Archivable(ABC):
-    @abstractmethod
-    def get_metadata(self) -> dict:
-        ...
\ No newline at end of file

From ac65e6f2e0e3c587f8ba34f591ebd757c7b2199c Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Tue, 26 Nov 2024 15:10:40 +0000
Subject: [PATCH 05/23] feat: update EfficientBBoxHead to inherit from BaseHead

---
 .../nodes/heads/efficient_bbox_head.py          | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/luxonis_train/nodes/heads/efficient_bbox_head.py b/luxonis_train/nodes/heads/efficient_bbox_head.py
index 0e6d344d..68aa1ce4 100644
--- a/luxonis_train/nodes/heads/efficient_bbox_head.py
+++ b/luxonis_train/nodes/heads/efficient_bbox_head.py
@@ -5,9 +5,8 @@
 from torch import Tensor, nn
 
 from luxonis_train.enums import TaskType
-from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import EfficientDecoupledBlock
-from luxonis_train.nodes.interfaces import Archivable
+from luxonis_train.nodes.heads import BaseHead
 from luxonis_train.utils import (
     Packet,
     anchors_for_fpn_features,
@@ -19,11 +18,12 @@
 
 
 class EfficientBBoxHead(
-    BaseNode[list[Tensor], tuple[list[Tensor], list[Tensor], list[Tensor]]],
-    Archivable,
+    BaseHead[list[Tensor], tuple[list[Tensor], list[Tensor], list[Tensor]]],
 ):
     in_channels: list[int]
     tasks: list[TaskType] = [TaskType.BOUNDINGBOX]
+    parser = "YOLO"
+    is_soxtmaxed: bool = False
 
     def __init__(
         self,
@@ -233,3 +233,12 @@ def _process_to_bbox(
             max_det=self.max_det,
             predicts_objectness=False,
         )
+
+    def _get_custom_head_config(self) -> dict:
+        """Returns custom head configuration."""
+        return {
+            "subtype": "yolov6",
+            "iou_threshold": self.iou_thres,
+            "conf_threshold": self.conf_thres,
+            "max_det": self.max_det,
+        }

From c0fe60b6b3d5743786bfa2f57ccf405fdbac23b5 Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Tue, 26 Nov 2024 15:20:04 +0000
Subject: [PATCH 06/23] style: formatting

---
 luxonis_train/nodes/heads/base_head.py | 34 ++++++++++++++++----------
 1 file changed, 21 insertions(+), 13 deletions(-)

diff --git a/luxonis_train/nodes/heads/base_head.py b/luxonis_train/nodes/heads/base_head.py
index 0c4de5dc..ee06e6ba 100644
--- a/luxonis_train/nodes/heads/base_head.py
+++ b/luxonis_train/nodes/heads/base_head.py
@@ -1,27 +1,35 @@
-
 from abc import abstractmethod
 from typing import Generic, Literal
 
-from luxonis_train.nodes.base_node import BaseNode, ForwardInputT, ForwardOutputT
+from luxonis_train.nodes.base_node import (
+    BaseNode,
+    ForwardInputT,
+    ForwardOutputT,
+)
 from luxonis_train.utils import deep_merge_dicts
 
 
-class BaseHead(BaseNode[ForwardInputT, ForwardOutputT], Generic[ForwardInputT, ForwardOutputT]):
+class BaseHead(
+    BaseNode[ForwardInputT, ForwardOutputT],
+    Generic[ForwardInputT, ForwardOutputT],
+):
     """Base class for all heads in the model.
 
     @type parser: str | None
     @ivar parser: Parser to use for the head.
-
     @type is_softmaxed: bool | None
     @ivar is_softmaxed: Whether the head uses softmax or not.
     """
 
-    parser: Literal[
-        "ClassificationParser",
-        "YOLO",
-        "YoloDetectionNetwork",
-        "SegmentationParser"
-    ] | None = None
+    parser: (
+        Literal[
+            "ClassificationParser",
+            "YOLO",
+            "YoloDetectionNetwork",
+            "SegmentationParser",
+        ]
+        | None
+    ) = None
     is_softmaxed: bool | None = None
 
     def __init__(self, **kwargs):
@@ -43,7 +51,7 @@ def get_head_config(self) -> dict:
         custom_config = self._get_custom_head_config()
         deep_merge_dicts(config, custom_config)
         return config
-        
+
     def _get_base_head_config(self) -> dict:
         """Get base head configuration.
 
@@ -56,7 +64,7 @@ def _get_base_head_config(self) -> dict:
                 "is_softmax": self.is_softmaxed,
                 "classes": self.class_names,
                 "n_classes": self.n_classes,
-            }
+            },
         }
 
     @abstractmethod
@@ -66,4 +74,4 @@ def _get_custom_head_config(self) -> dict:
         @rtype: dict
         @return: Custom head configuration.
         """
-        ...
\ No newline at end of file
+        ...

From 3dfed2409a48034ec4ae9be01fa91be43c847a95 Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Tue, 26 Nov 2024 15:24:35 +0000
Subject: [PATCH 07/23] fix: remove Archiable interface from head classes

---
 luxonis_train/nodes/heads/bisenet_head.py             | 3 +--
 luxonis_train/nodes/heads/classification_head.py      | 3 +--
 luxonis_train/nodes/heads/ddrnet_segmentation_head.py | 3 +--
 luxonis_train/nodes/heads/segmentation_head.py        | 3 +--
 4 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/luxonis_train/nodes/heads/bisenet_head.py b/luxonis_train/nodes/heads/bisenet_head.py
index 772e390a..b1fa57b3 100644
--- a/luxonis_train/nodes/heads/bisenet_head.py
+++ b/luxonis_train/nodes/heads/bisenet_head.py
@@ -5,11 +5,10 @@
 from luxonis_train.enums import TaskType
 from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import ConvModule
-from luxonis_train.nodes.interfaces import Archivable
 from luxonis_train.utils import infer_upscale_factor
 
 
-class BiSeNetHead(BaseNode[Tensor, Tensor], Archivable):
+class BiSeNetHead(BaseNode[Tensor, Tensor]):
     in_height: int
     in_width: int
     in_channels: int
diff --git a/luxonis_train/nodes/heads/classification_head.py b/luxonis_train/nodes/heads/classification_head.py
index ac09b36a..74033c2b 100644
--- a/luxonis_train/nodes/heads/classification_head.py
+++ b/luxonis_train/nodes/heads/classification_head.py
@@ -4,10 +4,9 @@
 
 from luxonis_train.enums import TaskType
 from luxonis_train.nodes.base_node import BaseNode
-from luxonis_train.nodes.interfaces import Archivable
 
 
-class ClassificationHead(BaseNode[Tensor, Tensor], Archivable):
+class ClassificationHead(BaseNode[Tensor, Tensor]):
     in_channels: int
     tasks: list[TaskType] = [TaskType.CLASSIFICATION]
 
diff --git a/luxonis_train/nodes/heads/ddrnet_segmentation_head.py b/luxonis_train/nodes/heads/ddrnet_segmentation_head.py
index f92c93fe..53912d7c 100644
--- a/luxonis_train/nodes/heads/ddrnet_segmentation_head.py
+++ b/luxonis_train/nodes/heads/ddrnet_segmentation_head.py
@@ -7,13 +7,12 @@
 
 from luxonis_train.enums import TaskType
 from luxonis_train.nodes.base_node import BaseNode
-from luxonis_train.nodes.interfaces import Archivable
 from luxonis_train.utils.general import infer_upscale_factor
 
 logger = logging.getLogger(__name__)
 
 
-class DDRNetSegmentationHead(BaseNode[Tensor, Tensor], Archivable):
+class DDRNetSegmentationHead(BaseNode[Tensor, Tensor]):
     attach_index: int = -1
     in_height: int
     in_width: int
diff --git a/luxonis_train/nodes/heads/segmentation_head.py b/luxonis_train/nodes/heads/segmentation_head.py
index 21c40498..31c79ab6 100644
--- a/luxonis_train/nodes/heads/segmentation_head.py
+++ b/luxonis_train/nodes/heads/segmentation_head.py
@@ -5,11 +5,10 @@
 from luxonis_train.enums import TaskType
 from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import UpBlock
-from luxonis_train.nodes.interfaces import Archivable
 from luxonis_train.utils import infer_upscale_factor
 
 
-class SegmentationHead(BaseNode[Tensor, Tensor], Archivable):
+class SegmentationHead(BaseNode[Tensor, Tensor]):
     in_height: int
     in_width: int
     in_channels: int

From 34b0c2949e79a80519ba98db3380b7419d21fde3 Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Wed, 27 Nov 2024 14:09:52 +0000
Subject: [PATCH 08/23] fix: improve BaseHead structure

---
 luxonis_train/core/utils/archive_utils.py     |  7 ++--
 luxonis_train/nodes/heads/base_head.py        | 32 ++++---------------
 .../nodes/heads/efficient_bbox_head.py        |  3 +-
 3 files changed, 13 insertions(+), 29 deletions(-)

diff --git a/luxonis_train/core/utils/archive_utils.py b/luxonis_train/core/utils/archive_utils.py
index 32faf279..0f79cc1f 100644
--- a/luxonis_train/core/utils/archive_utils.py
+++ b/luxonis_train/core/utils/archive_utils.py
@@ -143,8 +143,11 @@ def get_head_configs(
     for node_name, node in lightning_module.nodes.items():
         if not isinstance(node, BaseHead) or node.remove_on_export:
             continue
-
-        head_config = node.get_head_config()
+        try:
+            head_config = node.get_head_config()
+        except NotImplementedError as e:
+            logger.error(f"Failed to archive head `{node_name}`: {e}")
+            continue
         head_name = (
             node_name
             if node_name not in head_names
diff --git a/luxonis_train/nodes/heads/base_head.py b/luxonis_train/nodes/heads/base_head.py
index ee06e6ba..6fefb4fb 100644
--- a/luxonis_train/nodes/heads/base_head.py
+++ b/luxonis_train/nodes/heads/base_head.py
@@ -1,5 +1,4 @@
-from abc import abstractmethod
-from typing import Generic, Literal
+from typing import Generic
 
 from luxonis_train.nodes.base_node import (
     BaseNode,
@@ -17,30 +16,13 @@ class BaseHead(
 
     @type parser: str | None
     @ivar parser: Parser to use for the head.
-    @type is_softmaxed: bool | None
-    @ivar is_softmaxed: Whether the head uses softmax or not.
     """
 
-    parser: (
-        Literal[
-            "ClassificationParser",
-            "YOLO",
-            "YoloDetectionNetwork",
-            "SegmentationParser",
-        ]
-        | None
-    ) = None
-    is_softmaxed: bool | None = None
+    parser: str | None = None
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
-    def __init_subclass__(cls, **kwargs):
-        """Check if head has all required attributes."""
-        super().__init_subclass__()
-        if cls.parser is None:
-            raise ValueError(f"Head `{cls.__name__}` must define a parser.")
-
     def get_head_config(self) -> dict:
         """Get head configuration.
 
@@ -48,7 +30,7 @@ def get_head_config(self) -> dict:
         @return: Head configuration.
         """
         config = self._get_base_head_config()
-        custom_config = self._get_custom_head_config()
+        custom_config = self.get_custom_head_config()
         deep_merge_dicts(config, custom_config)
         return config
 
@@ -61,17 +43,17 @@ def _get_base_head_config(self) -> dict:
         return {
             "parser": self.parser,
             "metadata": {
-                "is_softmax": self.is_softmaxed,
                 "classes": self.class_names,
                 "n_classes": self.n_classes,
             },
         }
 
-    @abstractmethod
-    def _get_custom_head_config(self) -> dict:
+    def get_custom_head_config(self) -> dict:
         """Get custom head configuration.
 
         @rtype: dict
         @return: Custom head configuration.
         """
-        ...
+        raise NotImplementedError(
+            "get_custom_head_config method must be implemented."
+        )
diff --git a/luxonis_train/nodes/heads/efficient_bbox_head.py b/luxonis_train/nodes/heads/efficient_bbox_head.py
index 68aa1ce4..4fe1adf8 100644
--- a/luxonis_train/nodes/heads/efficient_bbox_head.py
+++ b/luxonis_train/nodes/heads/efficient_bbox_head.py
@@ -23,7 +23,6 @@ class EfficientBBoxHead(
     in_channels: list[int]
     tasks: list[TaskType] = [TaskType.BOUNDINGBOX]
     parser = "YOLO"
-    is_soxtmaxed: bool = False
 
     def __init__(
         self,
@@ -234,7 +233,7 @@ def _process_to_bbox(
             predicts_objectness=False,
         )
 
-    def _get_custom_head_config(self) -> dict:
+    def get_custom_head_config(self) -> dict:
         """Returns custom head configuration."""
         return {
             "subtype": "yolov6",

From 5ec6ca5da5f68fb58ebb6f4b63ae693b8d3af025 Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Wed, 27 Nov 2024 14:35:17 +0000
Subject: [PATCH 09/23] feat: update all heads to inherit from BaseHead

---
 luxonis_train/nodes/heads/bisenet_head.py         | 15 +++++++++++++--
 luxonis_train/nodes/heads/classification_head.py  | 15 +++++++++++++--
 .../nodes/heads/ddrnet_segmentation_head.py       | 15 +++++++++++++--
 .../heads/discsubnet_head/discsubnet_head.py      | 12 ++++++++++--
 luxonis_train/nodes/heads/efficient_bbox_head.py  |  6 +++++-
 .../nodes/heads/efficient_keypoint_bbox_head.py   | 14 ++++++++++++++
 luxonis_train/nodes/heads/segmentation_head.py    | 15 +++++++++++++--
 7 files changed, 81 insertions(+), 11 deletions(-)

diff --git a/luxonis_train/nodes/heads/bisenet_head.py b/luxonis_train/nodes/heads/bisenet_head.py
index b1fa57b3..1a57420f 100644
--- a/luxonis_train/nodes/heads/bisenet_head.py
+++ b/luxonis_train/nodes/heads/bisenet_head.py
@@ -3,17 +3,18 @@
 from torch import Tensor, nn
 
 from luxonis_train.enums import TaskType
-from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import ConvModule
+from luxonis_train.nodes.heads import BaseHead
 from luxonis_train.utils import infer_upscale_factor
 
 
-class BiSeNetHead(BaseNode[Tensor, Tensor]):
+class BiSeNetHead(BaseHead[Tensor, Tensor]):
     in_height: int
     in_width: int
     in_channels: int
 
     tasks: list[TaskType] = [TaskType.SEGMENTATION]
+    parser: str = "SegmentationParser"
 
     def __init__(self, intermediate_channels: int = 64, **kwargs: Any):
         """BiSeNet segmentation head.
@@ -56,3 +57,13 @@ def forward(self, inputs: Tensor) -> Tensor:
         x = self.conv_3x3(inputs)
         x = self.conv_1x1(x)
         return self.upscale(x)
+
+    def get_custom_head_config(self) -> dict:
+        """Returns custom head configuration.
+
+        @rtype: dict
+        @return: Custom head configuration.
+        """
+        return {
+            "is_softmax": False,
+        }
diff --git a/luxonis_train/nodes/heads/classification_head.py b/luxonis_train/nodes/heads/classification_head.py
index 74033c2b..46582354 100644
--- a/luxonis_train/nodes/heads/classification_head.py
+++ b/luxonis_train/nodes/heads/classification_head.py
@@ -3,12 +3,13 @@
 from torch import Tensor, nn
 
 from luxonis_train.enums import TaskType
-from luxonis_train.nodes.base_node import BaseNode
+from luxonis_train.nodes.heads import BaseHead
 
 
-class ClassificationHead(BaseNode[Tensor, Tensor]):
+class ClassificationHead(BaseHead[Tensor, Tensor]):
     in_channels: int
     tasks: list[TaskType] = [TaskType.CLASSIFICATION]
+    parser: str = "ClassificationParser"
 
     def __init__(self, dropout_rate: float = 0.2, **kwargs: Any):
         """Simple classification head.
@@ -31,3 +32,13 @@ def __init__(self, dropout_rate: float = 0.2, **kwargs: Any):
 
     def forward(self, inputs: Tensor) -> Tensor:
         return self.head(inputs)
+
+    def get_custom_head_config(self) -> dict:
+        """Returns custom head configuration.
+
+        @rtype: dict
+        @return: Custom head configuration.
+        """
+        return {
+            "is_softmax": False,
+        }
diff --git a/luxonis_train/nodes/heads/ddrnet_segmentation_head.py b/luxonis_train/nodes/heads/ddrnet_segmentation_head.py
index 53912d7c..001c52ab 100644
--- a/luxonis_train/nodes/heads/ddrnet_segmentation_head.py
+++ b/luxonis_train/nodes/heads/ddrnet_segmentation_head.py
@@ -6,19 +6,20 @@
 from torch import Tensor
 
 from luxonis_train.enums import TaskType
-from luxonis_train.nodes.base_node import BaseNode
+from luxonis_train.nodes.heads import BaseHead
 from luxonis_train.utils.general import infer_upscale_factor
 
 logger = logging.getLogger(__name__)
 
 
-class DDRNetSegmentationHead(BaseNode[Tensor, Tensor]):
+class DDRNetSegmentationHead(BaseHead[Tensor, Tensor]):
     attach_index: int = -1
     in_height: int
     in_width: int
     in_channels: int
 
     tasks: list[TaskType] = [TaskType.SEGMENTATION]
+    parser: str = "SegmentationParser"
 
     def __init__(
         self,
@@ -108,3 +109,13 @@ def forward(self, inputs: Tensor) -> Tensor:
             x = x.argmax(dim=1) if self.n_classes > 1 else x > 0
             return x.to(dtype=torch.int32)
         return x
+
+    def get_custom_head_config(self) -> dict:
+        """Returns custom head configuration.
+
+        @rtype: dict
+        @return: Custom head configuration.
+        """
+        return {
+            "is_softmax": False,
+        }
diff --git a/luxonis_train/nodes/heads/discsubnet_head/discsubnet_head.py b/luxonis_train/nodes/heads/discsubnet_head/discsubnet_head.py
index 70205ab8..6c8ff686 100644
--- a/luxonis_train/nodes/heads/discsubnet_head/discsubnet_head.py
+++ b/luxonis_train/nodes/heads/discsubnet_head/discsubnet_head.py
@@ -4,7 +4,7 @@
 from torch import Tensor
 
 from luxonis_train.enums import TaskType
-from luxonis_train.nodes.base_node import BaseNode
+from luxonis_train.nodes.heads import BaseHead
 from luxonis_train.utils import Packet
 
 from .blocks import Decoder, Encoder, NanoDecoder, NanoEncoder
@@ -27,7 +27,7 @@ def get_variant(variant: VariantLiteral) -> int:
     return variants[variant]
 
 
-class DiscSubNetHead(BaseNode[Tensor, Tensor]):
+class DiscSubNetHead(BaseHead[Tensor, Tensor]):
     in_channels: list[int] | int
     out_channels: int
     base_channels: int
@@ -109,3 +109,11 @@ def wrap(
                 "reconstructed": [recon],
                 "segmentation": [seg_out],
             }
+
+    def get_custom_head_config(self) -> dict:
+        """Returns custom head configuration.
+
+        @rtype: dict
+        @return: Custom head configuration.
+        """
+        return {}
diff --git a/luxonis_train/nodes/heads/efficient_bbox_head.py b/luxonis_train/nodes/heads/efficient_bbox_head.py
index 4fe1adf8..1881bfc3 100644
--- a/luxonis_train/nodes/heads/efficient_bbox_head.py
+++ b/luxonis_train/nodes/heads/efficient_bbox_head.py
@@ -234,7 +234,11 @@ def _process_to_bbox(
         )
 
     def get_custom_head_config(self) -> dict:
-        """Returns custom head configuration."""
+        """Returns custom head configuration.
+
+        @rtype: dict
+        @return: Custom head configuration.
+        """
         return {
             "subtype": "yolov6",
             "iou_threshold": self.iou_thres,
diff --git a/luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py b/luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py
index 0421275d..d855f831 100644
--- a/luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py
+++ b/luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py
@@ -17,6 +17,7 @@
 
 class EfficientKeypointBBoxHead(EfficientBBoxHead):
     tasks: list[TaskType] = [TaskType.KEYPOINTS, TaskType.BOUNDINGBOX]
+    parser: str = "YoloDetectionNetwork"
 
     def __init__(
         self,
@@ -214,3 +215,16 @@ def _process_to_bbox_and_kps(
             max_det=self.max_det,
             predicts_objectness=False,
         )
+
+    def get_custom_head_config(self) -> dict:
+        """Returns custom head configuration.
+
+        @rtype: dict
+        @return: Custom head configuration.
+        """
+        return {
+            "iou_threshold": self.iou_thres,
+            "conf_threshold": self.conf_thres,
+            "max_det": self.max_det,
+            "n_keypoints": self.n_keypoints,
+        }
diff --git a/luxonis_train/nodes/heads/segmentation_head.py b/luxonis_train/nodes/heads/segmentation_head.py
index 31c79ab6..3824336f 100644
--- a/luxonis_train/nodes/heads/segmentation_head.py
+++ b/luxonis_train/nodes/heads/segmentation_head.py
@@ -3,17 +3,18 @@
 from torch import Tensor, nn
 
 from luxonis_train.enums import TaskType
-from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import UpBlock
+from luxonis_train.nodes.heads import BaseHead
 from luxonis_train.utils import infer_upscale_factor
 
 
-class SegmentationHead(BaseNode[Tensor, Tensor]):
+class SegmentationHead(BaseHead[Tensor, Tensor]):
     in_height: int
     in_width: int
     in_channels: int
 
     tasks: list[TaskType] = [TaskType.SEGMENTATION]
+    parser: str = "SegmentationParser"
 
     def __init__(self, **kwargs: Any):
         """Basic segmentation FCN head.
@@ -40,3 +41,13 @@ def __init__(self, **kwargs: Any):
 
     def forward(self, inputs: Tensor) -> Tensor:
         return self.head(inputs)
+
+    def get_custom_head_config(self) -> dict:
+        """Returns custom head configuration.
+
+        @rtype: dict
+        @return: Custom head configuration.
+        """
+        return {
+            "is_softmax": False,
+        }

From 616c4c37836f6f8779674cbfe7426d89f992e87a Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Wed, 27 Nov 2024 14:42:15 +0000
Subject: [PATCH 10/23] feat: add a warning about not implemented
 get_custom_head_config in the head

---
 luxonis_train/models/luxonis_lightning.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/luxonis_train/models/luxonis_lightning.py b/luxonis_train/models/luxonis_lightning.py
index 08d0066f..e2d84701 100644
--- a/luxonis_train/models/luxonis_lightning.py
+++ b/luxonis_train/models/luxonis_lightning.py
@@ -32,6 +32,7 @@
 )
 from luxonis_train.config import AttachedModuleConfig, Config
 from luxonis_train.nodes import BaseNode
+from luxonis_train.nodes.heads import BaseHead
 from luxonis_train.utils import (
     DatasetMetadata,
     Kwargs,
@@ -353,6 +354,13 @@ def _initiate_nodes(
                 dataset_metadata=self.dataset_metadata,
                 **node_kwargs,
             )
+            if isinstance(node, BaseHead):
+                try:
+                    node.get_custom_head_config()
+                except NotImplementedError:
+                    logger.warning(
+                        f"Head {node_name} does not implement get_custom_head_config method. Archivation of this head will fail."
+                    )
             node_outputs = node.run(node_dummy_inputs)
 
             dummy_inputs[node_name] = node_outputs

From 35e9adcf0315d808820700e6acf6423e9e23aa46 Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Wed, 27 Nov 2024 16:53:27 +0000
Subject: [PATCH 11/23] fix: update test archive config

---
 tests/integration/parking_lot.json | 92 +++++++++++++-----------------
 1 file changed, 39 insertions(+), 53 deletions(-)

diff --git a/tests/integration/parking_lot.json b/tests/integration/parking_lot.json
index 37f8cb1a..197dbb1a 100644
--- a/tests/integration/parking_lot.json
+++ b/tests/integration/parking_lot.json
@@ -149,79 +149,50 @@
         ],
         "heads": [
             {
-                "name": "bbox-head",
-                "parser": "YOLO",
+                "name": "vehicle-type-segmentation-head",
+                "parser": "SegmentationParser",
                 "metadata": {
                     "postprocessor_path": null,
                     "classes": [
+                        "background",
                         "car",
                         "motorbike"
                     ],
-                    "n_classes": 2,
-                    "iou_threshold": 0.45,
-                    "conf_threshold": 0.25,
-                    "max_det": 300,
-                    "anchors": null,
-                    "subtype": "yolov6"
+                    "n_classes": 3
                 },
                 "outputs": [
-                    "output1_yolov6r2",
-                    "output2_yolov6r2",
-                    "output3_yolov6r2"
-                ]
-            },
-            {
-                "name": "motorbike-detection-head",
-                "parser": "YoloDetectionNetwork",
-                "metadata": {
-                    "postprocessor_path": null,
-                    "classes": [
-                        "motorbike"
-                    ],
-                    "n_classes": 1,
-                    "iou_threshold": 0.45,
-                    "conf_threshold": 0.25,
-                    "max_det": 300,
-                    "anchors": null,
-                    "n_keypoints": 3
-                },
-                "outputs": [
-                    "motorbike-detection-head/outputs/0",
-                    "motorbike-detection-head/outputs/1",
-                    "motorbike-detection-head/outputs/2"
+                    "vehicle-type-segmentation-head/vehicle_type-segmentation/0"
                 ]
             },
             {
-                "name": "color-segmentation-head",
+                "name": "any-vehicle-segmentation-head",
                 "parser": "SegmentationParser",
                 "metadata": {
                     "postprocessor_path": null,
                     "classes": [
-                        "background",
-                        "blue",
-                        "green",
-                        "red"
+                        "vehicle"
                     ],
-                    "n_classes": 4,
-                    "is_softmax": false
+                    "n_classes": 1
                 },
                 "outputs": [
-                    "color-segmentation-head/color-segmentation/0"
+                    "any-vehicle-segmentation-head/vehicle-segmentation/0"
                 ]
             },
             {
-                "name": "any-vehicle-segmentation-head",
-                "parser": "SegmentationParser",
+                "name": "bbox-head",
+                "parser": "YOLO",
                 "metadata": {
                     "postprocessor_path": null,
                     "classes": [
-                        "vehicle"
+                        "car",
+                        "motorbike"
                     ],
-                    "n_classes": 1,
-                    "is_softmax": false
+                    "n_classes": 2
                 },
                 "outputs": [
-                    "any-vehicle-segmentation-head/vehicle-segmentation/0"
+                    "output1_yolov6r2",
+                    "output2_yolov6r2",
+                    "output3_yolov6r2"
                 ]
             },
             {
@@ -254,28 +225,43 @@
                         "truimph",
                         "yamaha"
                     ],
-                    "n_classes": 23,
-                    "is_softmax": false
+                    "n_classes": 23
                 },
                 "outputs": [
                     "brand-segmentation-head/brand-segmentation/0"
                 ]
             },
             {
-                "name": "vehicle-type-segmentation-head",
+                "name": "color-segmentation-head",
                 "parser": "SegmentationParser",
                 "metadata": {
                     "postprocessor_path": null,
                     "classes": [
                         "background",
-                        "car",
+                        "blue",
+                        "green",
+                        "red"
+                    ],
+                    "n_classes": 4
+                },
+                "outputs": [
+                    "color-segmentation-head/color-segmentation/0"
+                ]
+            },
+            {
+                "name": "motorbike-detection-head",
+                "parser": "YoloDetectionNetwork",
+                "metadata": {
+                    "postprocessor_path": null,
+                    "classes": [
                         "motorbike"
                     ],
-                    "n_classes": 3,
-                    "is_softmax": false
+                    "n_classes": 1
                 },
                 "outputs": [
-                    "vehicle-type-segmentation-head/vehicle_type-segmentation/0"
+                    "motorbike-detection-head/outputs/0",
+                    "motorbike-detection-head/outputs/1",
+                    "motorbike-detection-head/outputs/2"
                 ]
             }
         ]

From 034891f828777e5309b374bfd374277b8777f165 Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Wed, 27 Nov 2024 17:08:14 +0000
Subject: [PATCH 12/23] test: sort fields in configs before comparison

---
 tests/integration/test_simple.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/tests/integration/test_simple.py b/tests/integration/test_simple.py
index e32980f2..72331fcc 100644
--- a/tests/integration/test_simple.py
+++ b/tests/integration/test_simple.py
@@ -123,6 +123,18 @@ def test_custom_tasks(
             ), "Config JSON not found in the archive."
             generated_config = json.loads(extracted_cfg.read().decode())
 
+        # Sort the fields in the config to make the comparison consistent
+        def sort_by_name(config, keys):
+            for key in keys:
+                if key in config["model"]:
+                    config["model"][key] = sorted(
+                        config["model"][key], key=lambda x: x["name"]
+                    )
+
+        keys_to_sort = ["inputs", "outputs", "heads"]
+        sort_by_name(generated_config, keys_to_sort)
+        sort_by_name(correct_archive_config, keys_to_sort)
+
         assert generated_config == correct_archive_config
 
 

From 105f82fcafe192fcd7483395d8adf40ab8af6834 Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Wed, 27 Nov 2024 17:11:17 +0000
Subject: [PATCH 13/23] fix: remove deprecated enums

---
 .../nodes/enums/head_categorization.py        | 23 -------------------
 1 file changed, 23 deletions(-)
 delete mode 100644 luxonis_train/nodes/enums/head_categorization.py

diff --git a/luxonis_train/nodes/enums/head_categorization.py b/luxonis_train/nodes/enums/head_categorization.py
deleted file mode 100644
index ce89a5d3..00000000
--- a/luxonis_train/nodes/enums/head_categorization.py
+++ /dev/null
@@ -1,23 +0,0 @@
-from enum import Enum
-
-
-class ImplementedHeads(Enum):
-    """Task categorization for the implemented heads."""
-
-    ClassificationHead = "ClassificationParser"
-    EfficientBBoxHead = "YOLO"
-    EfficientKeypointBBoxHead = "YoloDetectionNetwork"
-    SegmentationHead = "SegmentationParser"
-    BiSeNetHead = "SegmentationParser"
-    DDRNetSegmentationHead = "SegmentationParser"
-
-
-class ImplementedHeadsIsSoxtmaxed(Enum):
-    """Softmaxed output categorization for the implemented heads."""
-
-    ClassificationHead = False
-    EfficientBBoxHead = None
-    EfficientKeypointBBoxHead = None
-    SegmentationHead = False
-    BiSeNetHead = False
-    DDRNetSegmentationHead = False

From 7b310ecefa94eb0f64e6e26e769682cf6e0d38ac Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Wed, 27 Nov 2024 17:13:38 +0000
Subject: [PATCH 14/23] fix: add subtype to the keypoint head

---
 luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py b/luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py
index d855f831..e548417e 100644
--- a/luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py
+++ b/luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py
@@ -223,6 +223,7 @@ def get_custom_head_config(self) -> dict:
         @return: Custom head configuration.
         """
         return {
+            "subtype": "yolov6",
             "iou_threshold": self.iou_thres,
             "conf_threshold": self.conf_thres,
             "max_det": self.max_det,

From ba1512bc63c558432534fb4ac2a65cf17966dbfd Mon Sep 17 00:00:00 2001
From: Nikita Sokovnin <49622375+sokovninn@users.noreply.github.com>
Date: Wed, 27 Nov 2024 18:28:32 +0100
Subject: [PATCH 15/23] fix: update parser for the EfficientKeypointBBoxHead

Co-authored-by: KlemenSkrlj <47853619+klemen1999@users.noreply.github.com>
---
 luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py b/luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py
index d855f831..46965db6 100644
--- a/luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py
+++ b/luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py
@@ -17,7 +17,7 @@
 
 class EfficientKeypointBBoxHead(EfficientBBoxHead):
     tasks: list[TaskType] = [TaskType.KEYPOINTS, TaskType.BOUNDINGBOX]
-    parser: str = "YoloDetectionNetwork"
+    parser: str = "YOLOExtendedParser"
 
     def __init__(
         self,

From 27dff57b51dfe34d6f4422c43301a043330123fc Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Wed, 27 Nov 2024 19:47:18 +0000
Subject: [PATCH 16/23] fix: update test config

---
 tests/integration/parking_lot.json | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/tests/integration/parking_lot.json b/tests/integration/parking_lot.json
index 197dbb1a..f840f665 100644
--- a/tests/integration/parking_lot.json
+++ b/tests/integration/parking_lot.json
@@ -158,7 +158,8 @@
                         "car",
                         "motorbike"
                     ],
-                    "n_classes": 3
+                    "n_classes": 3,
+                    "is_softmax": false
                 },
                 "outputs": [
                     "vehicle-type-segmentation-head/vehicle_type-segmentation/0"
@@ -172,7 +173,8 @@
                     "classes": [
                         "vehicle"
                     ],
-                    "n_classes": 1
+                    "n_classes": 1,
+                    "is_softmax": false
                 },
                 "outputs": [
                     "any-vehicle-segmentation-head/vehicle-segmentation/0"
@@ -187,7 +189,12 @@
                         "car",
                         "motorbike"
                     ],
-                    "n_classes": 2
+                    "n_classes": 2,
+                    "iou_threshold": 0.45,
+                    "conf_threshold": 0.25,
+                    "max_det": 300,
+                    "anchors": null,
+                    "subtype": "yolov6"
                 },
                 "outputs": [
                     "output1_yolov6r2",
@@ -225,7 +232,8 @@
                         "truimph",
                         "yamaha"
                     ],
-                    "n_classes": 23
+                    "n_classes": 23,
+                    "is_softmax": false
                 },
                 "outputs": [
                     "brand-segmentation-head/brand-segmentation/0"
@@ -242,7 +250,8 @@
                         "green",
                         "red"
                     ],
-                    "n_classes": 4
+                    "n_classes": 4,
+                    "is_softmax": false
                 },
                 "outputs": [
                     "color-segmentation-head/color-segmentation/0"
@@ -250,13 +259,19 @@
             },
             {
                 "name": "motorbike-detection-head",
-                "parser": "YoloDetectionNetwork",
+                "parser": "YOLOExtendedParser",
                 "metadata": {
                     "postprocessor_path": null,
                     "classes": [
                         "motorbike"
                     ],
-                    "n_classes": 1
+                    "n_classes": 1,
+                    "iou_threshold": 0.45,
+                    "conf_threshold": 0.25,
+                    "max_det": 300,
+                    "anchors": null,
+                    "subtype": "yolov6",
+                    "n_keypoints": 3
                 },
                 "outputs": [
                     "motorbike-detection-head/outputs/0",

From a58163d6a4940efddedab36f1b6ae95275395298 Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Wed, 27 Nov 2024 19:48:40 +0000
Subject: [PATCH 17/23] fix: correct configs merge

---
 luxonis_train/nodes/heads/base_head.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/luxonis_train/nodes/heads/base_head.py b/luxonis_train/nodes/heads/base_head.py
index 6fefb4fb..24b5e3ed 100644
--- a/luxonis_train/nodes/heads/base_head.py
+++ b/luxonis_train/nodes/heads/base_head.py
@@ -5,7 +5,6 @@
     ForwardInputT,
     ForwardOutputT,
 )
-from luxonis_train.utils import deep_merge_dicts
 
 
 class BaseHead(
@@ -31,7 +30,7 @@ def get_head_config(self) -> dict:
         """
         config = self._get_base_head_config()
         custom_config = self.get_custom_head_config()
-        deep_merge_dicts(config, custom_config)
+        config["metadata"].update(custom_config)
         return config
 
     def _get_base_head_config(self) -> dict:

From 7d0ccc03cfee541d17beda39896ace12b62049cc Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Wed, 27 Nov 2024 19:49:59 +0000
Subject: [PATCH 18/23] fix: move EMACallback in the light detection model
 config

---
 configs/detection_light_model.yaml | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/configs/detection_light_model.yaml b/configs/detection_light_model.yaml
index 0fb3a085..41ca5cc2 100644
--- a/configs/detection_light_model.yaml
+++ b/configs/detection_light_model.yaml
@@ -10,7 +10,7 @@ model:
 
 loader:
   params:
-    dataset_name: coco_test
+    dataset_name: coco_10
 
 trainer:
   preprocessing:
@@ -26,6 +26,11 @@ trainer:
   n_log_images: 8
 
   callbacks:
+    - name: EMACallback
+      params:
+        decay: 0.9999 
+        use_dynamic_decay: True 
+        decay_tau: 2000
     - name: ExportOnTrainEnd
     - name: TestOnTrainEnd
 
@@ -43,10 +48,3 @@ trainer:
     params:
       T_max: *epochs         
       eta_min: 0.000384
-
-  callbacks:
-    - name: EMACallback
-      params:
-        decay: 0.9999 
-        use_dynamic_decay: True 
-        decay_tau: 2000

From 238fd35d448c33126a3835b5263202204f37efef Mon Sep 17 00:00:00 2001
From: Nikita Sokovnin <49622375+sokovninn@users.noreply.github.com>
Date: Thu, 28 Nov 2024 13:35:23 +0100
Subject: [PATCH 19/23] fix: revert dataset_name in the detection light config

---
 configs/detection_light_model.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/configs/detection_light_model.yaml b/configs/detection_light_model.yaml
index 41ca5cc2..13fdc54e 100644
--- a/configs/detection_light_model.yaml
+++ b/configs/detection_light_model.yaml
@@ -10,7 +10,7 @@ model:
 
 loader:
   params:
-    dataset_name: coco_10
+    dataset_name: coco_test
 
 trainer:
   preprocessing:

From 7f590f6ecd02049532df5b53136451eff0ab7361 Mon Sep 17 00:00:00 2001
From: Martin Kozlovsky <martin.kozlovsky@luxonis.com>
Date: Thu, 28 Nov 2024 13:43:15 +0100
Subject: [PATCH 20/23] added coverage rule

---
 pyproject.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pyproject.toml b/pyproject.toml
index 39c11a92..42457ed6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -83,6 +83,7 @@ exclude_also = [
     "def __str__",
     "assert",
     "raise NotImplementedError",
+    "except NotImplementedError",
     "except ImportError",
     "@abstractmethod",
     "@overload",

From 050ccc4455f430bfef35683b603063e4299a2b3f Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Thu, 28 Nov 2024 13:11:35 +0000
Subject: [PATCH 21/23] fix: remove redundant init in the BaseHead

---
 luxonis_train/nodes/heads/base_head.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/luxonis_train/nodes/heads/base_head.py b/luxonis_train/nodes/heads/base_head.py
index 24b5e3ed..931e006e 100644
--- a/luxonis_train/nodes/heads/base_head.py
+++ b/luxonis_train/nodes/heads/base_head.py
@@ -19,9 +19,6 @@ class BaseHead(
 
     parser: str | None = None
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
     def get_head_config(self) -> dict:
         """Get head configuration.
 

From dbc9027da198d7ff48119cdf74e06112f790cd22 Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Thu, 28 Nov 2024 13:20:08 +0000
Subject: [PATCH 22/23] fix: remove deep_merge_dicts

---
 luxonis_train/utils/__init__.py |  2 --
 luxonis_train/utils/general.py  | 12 ------------
 2 files changed, 14 deletions(-)

diff --git a/luxonis_train/utils/__init__.py b/luxonis_train/utils/__init__.py
index ac6fa1b3..2944dfde 100644
--- a/luxonis_train/utils/__init__.py
+++ b/luxonis_train/utils/__init__.py
@@ -9,7 +9,6 @@
 from .dataset_metadata import DatasetMetadata
 from .exceptions import IncompatibleException
 from .general import (
-    deep_merge_dicts,
     get_with_default,
     infer_upscale_factor,
     make_divisible,
@@ -31,7 +30,6 @@
     "make_divisible",
     "infer_upscale_factor",
     "to_shape_packet",
-    "deep_merge_dicts",
     "get_with_default",
     "safe_download",
     "LuxonisTrackerPL",
diff --git a/luxonis_train/utils/general.py b/luxonis_train/utils/general.py
index 020cb4bb..390ddfd2 100644
--- a/luxonis_train/utils/general.py
+++ b/luxonis_train/utils/general.py
@@ -203,15 +203,3 @@ def clean_url(url: str) -> str:
 def url2file(url: str) -> str:
     """Convert URL to filename, i.e. https://url.com/file.txt?auth -> file.txt."""
     return Path(clean_url(url)).name
-
-
-def deep_merge_dicts(dict1: dict, dict2: dict) -> None:
-    for key, value in dict2.items():
-        if (
-            isinstance(value, dict)
-            and key in dict1
-            and isinstance(dict1[key], dict)
-        ):
-            deep_merge_dicts(dict1[key], value)
-        else:
-            dict1[key] = value

From d2a05f460191cda687a25ea931a4de95ebb4f892 Mon Sep 17 00:00:00 2001
From: Nikita <nikitoos67@gmail.com>
Date: Thu, 28 Nov 2024 13:20:58 +0000
Subject: [PATCH 23/23] fix: remove inheritance from Generic in BaseHead

---
 luxonis_train/nodes/heads/base_head.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/luxonis_train/nodes/heads/base_head.py b/luxonis_train/nodes/heads/base_head.py
index 931e006e..2eb2377d 100644
--- a/luxonis_train/nodes/heads/base_head.py
+++ b/luxonis_train/nodes/heads/base_head.py
@@ -1,5 +1,3 @@
-from typing import Generic
-
 from luxonis_train.nodes.base_node import (
     BaseNode,
     ForwardInputT,
@@ -9,7 +7,6 @@
 
 class BaseHead(
     BaseNode[ForwardInputT, ForwardOutputT],
-    Generic[ForwardInputT, ForwardOutputT],
 ):
     """Base class for all heads in the model.