Archive without Export (#96)

luxonis_train/__main__.py

@@ -179,11 +179,11 @@ def inspect(
 @app.command()
 def archive(
     executable: Annotated[
-        str,
+        str | None,
         typer.Option(
             help="Path to the model file.", show_default=False, metavar="FILE"
         ),
-    ],
+    ] = None,
     config: ConfigType = None,
     opts: OptsType = None,
 ):

luxonis_train/callbacks/archive_on_train_end.py

@@ -25,16 +25,19 @@ def on_train_end(
         @param pl_module: Pytorch Lightning module.
         """
 
-        path = self.get_checkpoint(pl_module)
-        if path is None:  # pragma: no cover
-            logger.warning("Skipping model archiving.")
-            return
-
         onnx_path = pl_module.core._exported_models.get("onnx")
+        if onnx_path is None:  # pragma: no cover
+            checkpoint = self.get_checkpoint(pl_module)
+            if checkpoint is None:
+                logger.warning("Skipping model archiving.")
+                return
+            logger.info("Exported model not found. Exporting to ONNX...")
+            pl_module.core.export(weights=checkpoint)
+            onnx_path = pl_module.core._exported_models.get("onnx")
+
         if onnx_path is None:  # pragma: no cover
             logger.error(
-                "Model executable not found. "
-                "Make sure to run exporter callback before archiver callback. "
+                "Model executable not found and couldn't be created. "
                 "Skipping model archiving."
             )
             return

luxonis_train/callbacks/export_on_train_end.py

@@ -24,9 +24,9 @@ def on_train_end(
         @type pl_module: L{pl.LightningModule}
         @param pl_module: Pytorch Lightning module.
         """
-        path = self.get_checkpoint(pl_module)
-        if path is None:  # pragma: no cover
+        checkpoint = self.get_checkpoint(pl_module)
+        if checkpoint is None:  # pragma: no cover
             logger.warning("Skipping model export.")
             return
 
-        pl_module.core.export(weights=self.get_checkpoint(pl_module))
+        pl_module.core.export(weights=checkpoint)

luxonis_train/config/config.py

@@ -498,7 +498,7 @@ def get_config(
         cls,
         cfg: str | dict[str, Any] | None = None,
         overrides: dict[str, Any] | list[str] | tuple[str, ...] | None = None,
-    ):
+    ) -> "Config":
         instance = super().get_config(cfg, overrides)
         if not isinstance(cfg, str):
             return instance

luxonis_train/core/core.py

@@ -640,6 +640,7 @@ def archive(self, path: str | Path | None = None) -> Path:
         outputs = []
 
         if path is None:
+            logger.warning("No model executable specified for archiving.")
             if "onnx" not in self._exported_models:
                 logger.info("Exporting model to ONNX...")
                 self.export()

luxonis_train/core/utils/archive_utils.py

@@ -200,7 +200,8 @@ def _get_head_outputs(
     # TODO: Fix this, will require refactoring custom ONNX output names
     logger.error(
         "ONNX model uses custom output names, trying to determine outputs based on the head type. "
-        "This will likely result in incorrect archive for multi-head models."
+        "This will likely result in incorrect archive for multi-head models. "
+        "You can ignore this error if your model has only one head."
     )
 
     if head_type == "ClassificationHead":

luxonis_train/nodes/heads/ddrnet_segmentation_head.py

@@ -1,4 +1,5 @@
 import logging
+from typing import Literal
 
 import torch
 import torch.nn as nn
@@ -22,7 +23,15 @@ class DDRNetSegmentationHead(BaseNode[Tensor, Tensor]):
     def __init__(
         self,
         inter_channels: int = 64,
-        inter_mode: str = "bilinear",
+        inter_mode: Literal[
+            "nearest",
+            "linear",
+            "bilinear",
+            "bicubic",
+            "trilinear",
+            "area",
+            "pixel_shuffle",
+        ] = "bilinear",
         **kwargs,
     ):
         """DDRNet segmentation head.