[Feat] Added name attribute to HeadType (#86)

Author: klemen1999

luxonis_train/core/utils/archive_utils.py

@@ -235,7 +235,7 @@ def get_heads(
     @param nodes: Dictionary of nodes.
     """
     heads = []
-
+    head_names = set()
     for node in cfg.model.nodes:
         node_name = node.name
         node_alias = node.alias or node_name
@@ -250,7 +250,13 @@ def get_heads(
                 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,

luxonis_train/models/luxonis_lightning.py

@@ -535,6 +535,14 @@ def export_onnx(self, save_path: str, **kwargs) -> list[str]:
             for node_name, output_name, i in output_order
         ]
 
+        if not self.cfg.exporter.output_names:
+            idx = 1
+            # Set to output names required by DAI
+            for i, output_name in enumerate(output_names):
+                if output_name.startswith("EfficientBBoxHead"):
+                    output_names[i] = f"output{idx}_yolov6r2"
+                    idx += 1
+
         old_forward = self.forward
 
         def export_forward(inputs) -> tuple[Tensor, ...]: