pre-comit formatting

luxonis_train/__main__.py

@@ -207,6 +207,7 @@ def archive(config: ConfigType = None, opts: OptsType = None, executable: Path =
 
     Archiver(str(config), opts).archive(executable)
 
+
 def version_callback(value: bool):
     if value:
         typer.echo(f"LuxonisTrain Version: {version(__package__)}")

luxonis_train/callbacks/archive_on_train_end.py

@@ -1,5 +1,5 @@
-import os
 import logging
+import os
 from pathlib import Path
 from typing import cast
 
@@ -13,7 +13,8 @@
 @CALLBACKS.register_module()
 class ArchiveOnTrainEnd(pl.Callback):
     def __init__(self, upload_to_mlflow: bool = False):
-        """Callback that performs archiving of onnx or exported model at the end of training/export. TODO: description
+        """Callback that performs archiving of onnx or exported model at the end of
+        training/export. TODO: description.
 
         @type upload_to_mlflow: bool
         @param upload_to_mlflow: If set to True, overrides the upload url in Archiver
@@ -59,15 +60,13 @@ def on_train_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> No
                     "`upload_to_mlflow` is set to True, "
                     "but there is  no MLFlow active run, skipping."
                 )
-        
+
         onnx_path = str(Path(best_model_path).parent.with_suffix(".onnx"))
         if not os.path.exists(onnx_path):
-            raise FileNotFoundError("Model executable not found. Make sure to run exporter callback before archiver callback")
-            # TODO: if onnx model non-existent, should export be ran?
-            #from luxonis_train.core.exporter import Exporter
-            #exporter = Exporter(cfg=cfg)
-            #exporter.export(onnx_path=onnx_path)
+            raise FileNotFoundError(
+                "Model executable not found. Make sure to run exporter callback before archiver callback"
+            )
 
         archiver = Archiver(cfg=cfg)
-        
+
         archiver.archive(onnx_path)

luxonis_train/core/__init__.py

@@ -1,7 +1,7 @@
+from .archiver import Archiver
 from .exporter import Exporter
 from .inferer import Inferer
 from .trainer import Trainer
 from .tuner import Tuner
-from .archiver import Archiver
 
 __all__ = ["Exporter", "Trainer", "Tuner", "Inferer", "Archiver"]

luxonis_train/nodes/enums/head_categorization.py

@@ -1,12 +1,13 @@
 from enum import Enum
 
+
 class ImplementedHeads(Enum):
-    """ Task categorization for the implemented heads. """
+    """Task categorization for the implemented heads."""
 
     ClassificationHead = "Classification"
-    EfficientBBoxHead = "ObjectDetectionYOLO" # TODO ObjectDetectionYOLO/ObjectDetectionSSD?
+    EfficientBBoxHead = (
+        "ObjectDetectionYOLO"  # TODO ObjectDetectionYOLO/ObjectDetectionSSD?
+    )
     ImplicitKeypointBBoxHead = "KeypointDetection"
     SegmentationHead = "Segmentation"
-    BiSeNetHead = "Segmentation" # TODO: SemanticSegmentation?
-
-
+    BiSeNetHead = "Segmentation"  # TODO: SemanticSegmentation?