[Feat] Added name attribute to HeadType

luxonis_train/core/utils/archive_utils.py

@@ -235,7 +235,7 @@
     @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 @@
                 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

@@ -298,10 +298,15 @@ def _initiate_nodes(
             for source_name, shape in shapes.items()
         }
 
-        for node_name, (
-            Node,
-            node_kwargs,
-        ), node_input_names, _ in traverse_graph(self.graph, nodes):
+        for (
+            node_name,
+            (
+                Node,
+                node_kwargs,
+            ),
+            node_input_names,
+            _,
+        ) in traverse_graph(self.graph, nodes):
             node_dummy_inputs: list[Packet[Tensor]] = []
             """List of dummy input packets for the node.
 
@@ -530,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, ...]: