option to source custom code in CLI

luxonis_train/__main__.py

@@ -214,13 +214,18 @@ def common(
             "--version", callback=version_callback, help="Show version and exit."
         ),
     ] = False,
+    source: Annotated[
+        Optional[Path],
+        typer.Option(
+            help="Path to a python file with custom components. "
+            "Will be sourced before running the command.",
+            metavar="FILE",
+        ),
+    ] = None,
 ):
-    ...
-
-
-def main():
-    app()
+    if source:
+        exec(source.read_text(), {}, {})
 
 
 if __name__ == "__main__":
-    main()
+    app()

luxonis_train/core/exporter.py

@@ -18,15 +18,12 @@
 
 
 class Exporter(Core):
-    """Main API which is used to create the model, setup pytorch lightning environment
-    and perform training based on provided arguments and config."""
-
     def __init__(
         self,
         cfg: str | dict[str, Any] | Config,
         opts: list[str] | tuple[str, ...] | dict[str, Any] | None = None,
     ):
-        """Constructs a new Exporter instance.
+        """Provides an interface for exporting models to .onnx and .blob formats.
 
         @type cfg: str | dict[str, Any] | Config
         @param cfg: Path to config file or config dict used to setup training.

pyproject.toml

@@ -19,7 +19,7 @@ classifiers = [
 ]
 
 [project.scripts]
-luxonis_train = "tools.main:main"
+luxonis_train = "luxonis_train.__main__:app"
 
 [project.urls]
 repository = "https://github.com/luxonis/luxonis-train"