Added --output-dir option to CLI

README.md

@@ -90,11 +90,13 @@ Tools CLI
 │                           omitted, the command runs automatic version        │
 │                           detection. [default: None]                         │
 │ --encoding                Color encoding used by the input model. Must be    │
-│                           RGB or BGR. [choices: rgb, bgr] [default: rgb]     │
+│                           RGB or BGR. [choices: rgb, bgr]                     │
 │ --use-rvc2 --no-use-rvc2  Whether to target RVC2 instead of RVC3. [default:  │
 │                           True]                                              │
 │ --class-names             Comma-separated class names recognized by the      │
 │                           model. [default: None]                             │
+│ --output-dir              Directory where generated conversion artifacts are │
+│                           stored. [default: None]                            │
 │ --output-remote-url       Remote destination URL for uploading the generated │
 │                           NN archive. [default: None]                        │
 │ --put-file-plugin         Name of a function registered in PUT_FILE_REGISTRY │

tests/test_unittests.py

@@ -40,6 +40,14 @@ def _output_dir(test_workspace: Path) -> str:
     return str(test_workspace / "shared_with_container" / "outputs")
 
 
+def test_help(test_workspace: Path):
+    """Tests that CLI help rendering works."""
+    result = _run_tools(["tools", "--help"], test_workspace)
+
+    assert result.returncode == 0, result.stdout
+    assert "--version" in result.stdout
+
+
 MODEL_EXPLICIT_VERSION = [
     ("yolov5n", "yolov5"),
     ("yolov5nu", "yolov5u"),
@@ -224,6 +232,46 @@ def test_explicit_class_names(test_workspace: Path):
     )
 
 
+def test_explicit_output_dir(test_workspace: Path):
+    """Tests writing conversion artifacts to a custom output directory."""
+    model_path = _prepare_model("yolov8n", test_workspace)
+    output_dir = test_workspace / "custom-output"
+    default_output_dir = Path(_output_dir(test_workspace))
+    default_output_dir_state = (
+        default_output_dir.exists(),
+        sorted(
+            path.relative_to(default_output_dir)
+            for path in default_output_dir.rglob("*")
+        )
+        if default_output_dir.exists()
+        else [],
+    )
+    command = [
+        "tools",
+        model_path,
+        "--version",
+        "yolov8",
+        "--output-dir",
+        str(output_dir),
+    ]
+    logger.debug(f"CLI command: {command}")
+
+    result = _run_tools(command, test_workspace)
+    if result.returncode != 0:
+        pytest.fail(f"Exit code: {result.returncode}, Output: {result.stdout}")
+
+    nn_archive_checker(output_dir=str(output_dir))
+    assert (
+        default_output_dir.exists(),
+        sorted(
+            path.relative_to(default_output_dir)
+            for path in default_output_dir.rglob("*")
+        )
+        if default_output_dir.exists()
+        else [],
+    ) == default_output_dir_state
+
+
 def test_wrong_explicit_class_names(test_workspace: Path):
     """Tests setting wrong explicit class names."""
     model_name = "yolov8n"

tools/conversion_registry.py

@@ -22,7 +22,7 @@
     YOLOX_CONVERSION,
 )
 
-ExporterFactory = Callable[[str, tuple[int, int], bool], Any]
+ExporterFactory = Callable[[str, tuple[int, int], bool, str | None], Any]
 
 
 @dataclass(frozen=True)
@@ -32,83 +32,83 @@ class ConversionSpec:
 
 
 def _build_yolov5_exporter(
-    model_path: str, imgsz: tuple[int, int], use_rvc2: bool
+    model_path: str, imgsz: tuple[int, int], use_rvc2: bool, output_dir: str | None
 ) -> Any:
     from tools.yolo.yolov5_exporter import YoloV5Exporter
 
-    return YoloV5Exporter(model_path, imgsz, use_rvc2)
+    return YoloV5Exporter(model_path, imgsz, use_rvc2, output_dir)
 
 
 def _build_yolov6r1_exporter(
-    model_path: str, imgsz: tuple[int, int], use_rvc2: bool
+    model_path: str, imgsz: tuple[int, int], use_rvc2: bool, output_dir: str | None
 ) -> Any:
     from tools.yolov6r1.yolov6_r1_exporter import YoloV6R1Exporter
 
-    return YoloV6R1Exporter(model_path, imgsz, use_rvc2)
+    return YoloV6R1Exporter(model_path, imgsz, use_rvc2, output_dir)
 
 
 def _build_yolov6r3_exporter(
-    model_path: str, imgsz: tuple[int, int], use_rvc2: bool
+    model_path: str, imgsz: tuple[int, int], use_rvc2: bool, output_dir: str | None
 ) -> Any:
     from tools.yolov6r3.yolov6_r3_exporter import YoloV6R3Exporter
 
-    return YoloV6R3Exporter(model_path, imgsz, use_rvc2)
+    return YoloV6R3Exporter(model_path, imgsz, use_rvc2, output_dir)
 
 
 def _build_goldyolo_exporter(
-    model_path: str, imgsz: tuple[int, int], use_rvc2: bool
+    model_path: str, imgsz: tuple[int, int], use_rvc2: bool, output_dir: str | None
 ) -> Any:
     from tools.yolov6r3.gold_yolo_exporter import GoldYoloExporter
 
-    return GoldYoloExporter(model_path, imgsz, use_rvc2)
+    return GoldYoloExporter(model_path, imgsz, use_rvc2, output_dir)
 
 
 def _build_yolov6r4_exporter(
-    model_path: str, imgsz: tuple[int, int], use_rvc2: bool
+    model_path: str, imgsz: tuple[int, int], use_rvc2: bool, output_dir: str | None
 ) -> Any:
     from tools.yolo.yolov6_exporter import YoloV6R4Exporter
 
-    return YoloV6R4Exporter(model_path, imgsz, use_rvc2)
+    return YoloV6R4Exporter(model_path, imgsz, use_rvc2, output_dir)
 
 
 def _build_yolov7_exporter(
-    model_path: str, imgsz: tuple[int, int], use_rvc2: bool
+    model_path: str, imgsz: tuple[int, int], use_rvc2: bool, output_dir: str | None
 ) -> Any:
     from tools.yolov7.yolov7_exporter import YoloV7Exporter
 
-    return YoloV7Exporter(model_path, imgsz, use_rvc2)
+    return YoloV7Exporter(model_path, imgsz, use_rvc2, output_dir)
 
 
 def _build_yolov8_exporter(
-    model_path: str, imgsz: tuple[int, int], use_rvc2: bool
+    model_path: str, imgsz: tuple[int, int], use_rvc2: bool, output_dir: str | None
 ) -> Any:
     from tools.yolo.yolov8_exporter import YoloV8Exporter
 
-    return YoloV8Exporter(model_path, imgsz, use_rvc2)
+    return YoloV8Exporter(model_path, imgsz, use_rvc2, output_dir)
 
 
 def _build_yolo26_exporter(
-    model_path: str, imgsz: tuple[int, int], use_rvc2: bool
+    model_path: str, imgsz: tuple[int, int], use_rvc2: bool, output_dir: str | None
 ) -> Any:
     from tools.yolo.yolo26_exporter import Yolo26Exporter
 
-    return Yolo26Exporter(model_path, imgsz, use_rvc2)
+    return Yolo26Exporter(model_path, imgsz, use_rvc2, output_dir)
 
 
 def _build_yolov10_exporter(
-    model_path: str, imgsz: tuple[int, int], use_rvc2: bool
+    model_path: str, imgsz: tuple[int, int], use_rvc2: bool, output_dir: str | None
 ) -> Any:
     from tools.yolo.yolov10_exporter import YoloV10Exporter
 
-    return YoloV10Exporter(model_path, imgsz, use_rvc2)
+    return YoloV10Exporter(model_path, imgsz, use_rvc2, output_dir)
 
 
 def _build_yolox_exporter(
-    model_path: str, imgsz: tuple[int, int], use_rvc2: bool
+    model_path: str, imgsz: tuple[int, int], use_rvc2: bool, output_dir: str | None
 ) -> Any:
     from tools.yolox.yolox_exporter import YoloXExporter
 
-    return YoloXExporter(model_path, imgsz, use_rvc2)
+    return YoloXExporter(model_path, imgsz, use_rvc2, output_dir)
 
 
 CONVERSION_SPECS: dict[str, ConversionSpec] = {
@@ -144,6 +144,12 @@ def get_exporter_family(version: str) -> str:
 
 
 def create_exporter(
-    version: str, model_path: str, imgsz: tuple[int, int], use_rvc2: bool
+    version: str,
+    model_path: str,
+    imgsz: tuple[int, int],
+    use_rvc2: bool,
+    output_dir: str | None = None,
 ) -> Any:
-    return CONVERSION_SPECS[version].exporter_factory(model_path, imgsz, use_rvc2)
+    return CONVERSION_SPECS[version].exporter_factory(
+        model_path, imgsz, use_rvc2, output_dir
+    )

tools/main.py

@@ -63,7 +63,7 @@ def convert(
     ] = None,
     encoding: Annotated[
         Encoding | None,
-        Parameter(show_default=True),
+        Parameter(show_default=False),
     ] = None,
     use_rvc2: Annotated[
         bool,
@@ -73,6 +73,10 @@ def convert(
         str | None,
         Parameter(show_default=True),
     ] = None,
+    output_dir: Annotated[
+        str | None,
+        Parameter(show_default=True),
+    ] = None,
     output_remote_url: Annotated[
         str | None,
         Parameter(show_default=True),
@@ -98,6 +102,7 @@ def convert(
             ``BGR``. When omitted it is selected based on version.
         use_rvc2: Whether to target RVC2 instead of RVC3.
         class_names: Comma-separated class names recognized by the model.
+        output_dir: Directory where generated conversion artifacts are stored.
         output_remote_url: Remote destination URL for uploading the generated NN
             archive.
         put_file_plugin: Name of a function registered in
@@ -183,6 +188,7 @@ def convert(
                     "encoding": encoding,
                     "use_rvc2": use_rvc2,
                     "class_names": class_names_list,
+                    "output_dir": output_dir,
                     "output_remote_url": output_remote_url,
                     "put_file_plugin": put_file_plugin,
                 }
@@ -211,6 +217,7 @@ def convert(
                 str(model_path),
                 exporter_imgsz,
                 config.use_rvc2,
+                config.output_dir,
             )
             logger.info("Model loaded.")
         except Exception as e:

tools/modules/exporter.py

@@ -2,6 +2,7 @@
 
 import os
 from datetime import datetime
+from pathlib import Path
 from typing import Any
 
 import onnx
@@ -29,6 +30,7 @@ def __init__(
         subtype: str,
         output_names: list[str] | None = None,
         all_output_names: list[str] | None = None,
+        output_dir: str | Path | None = None,
     ):
         """Initialize the exporter state and output paths.
 
@@ -40,6 +42,8 @@ def __init__(
             output_names: Primary output tensor names.
             all_output_names: Complete output tensor names. When omitted,
                 ``output_names`` is reused.
+            output_dir: Root directory for generated artifacts. When omitted,
+                the default ``shared_with_container/outputs`` directory is used.
         """
         # Set up variables
         self.model_path = model_path
@@ -56,8 +60,9 @@ def __init__(
         self.all_output_names = (
             all_output_names if all_output_names is not None else output_names
         )
+        output_root = Path(output_dir) if output_dir is not None else OUTPUTS_DIR
         self.output_folder = (
-            OUTPUTS_DIR
+            output_root
             / f"{self.model_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
         ).resolve()
         # If output directory does not exist, create it

tools/utils/config.py

@@ -23,6 +23,10 @@ class Config(LuxonisConfig):
     )
     class_names: list[str] | None = Field(None, description="List of class names.")
     use_rvc2: Literal[False, True] = Field(True, description="Whether to use RVC2.")
+    output_dir: str | None = Field(
+        None,
+        description="Directory where generated conversion artifacts are stored.",
+    )
     output_remote_url: str | None = Field(
         None, description="URL to upload the output to."
     )

tools/yolo/yolo26_exporter.py

@@ -50,13 +50,20 @@ def get_yolo_output_names(mode: int = 0):
 
 
 class Yolo26Exporter(Exporter):
-    def __init__(self, model_path: str, imgsz: tuple[int, int], use_rvc2: bool):
+    def __init__(
+        self,
+        model_path: str,
+        imgsz: tuple[int, int],
+        use_rvc2: bool,
+        output_dir: str | None = None,
+    ):
         super().__init__(
             model_path,
             imgsz,
             use_rvc2,
             subtype="yolo26",
             output_names=["output_yolo26"],
+            output_dir=output_dir,
         )
         self.load_model()
 

tools/yolo/yolov10_exporter.py

@@ -23,13 +23,15 @@ def __init__(
         model_path: str,
         imgsz: tuple[int, int],
         use_rvc2: bool,
+        output_dir: str | None = None,
     ):
         super().__init__(
             model_path,
             imgsz,
             use_rvc2,
             subtype="yolov10",
             output_names=["output1_yolov10", "output2_yolov10", "output3_yolov10"],
+            output_dir=output_dir,
         )
         self.load_model()
 

tools/yolo/yolov5_exporter.py

@@ -91,13 +91,15 @@ def __init__(
         model_path: str,
         imgsz: tuple[int, int],
         use_rvc2: bool,
+        output_dir: str | None = None,
     ):
         super().__init__(
             model_path,
             imgsz,
             use_rvc2,
             subtype="yolov5",
             output_names=["output1_yolov5", "output2_yolov5", "output3_yolov5"],
+            output_dir=output_dir,
         )
         self.load_model()
 

tools/yolo/yolov6_exporter.py

@@ -45,13 +45,15 @@ def __init__(
         model_path: str,
         imgsz: tuple[int, int],
         use_rvc2: bool,
+        output_dir: str | None = None,
     ):
         super().__init__(
             model_path,
             imgsz,
             use_rvc2,
             subtype="yolov6r2",
             output_names=["output1_yolov6r2", "output2_yolov6r2", "output3_yolov6r2"],
+            output_dir=output_dir,
         )
         self.load_model()
 

tools/yolo/yolov8_exporter.py

@@ -105,13 +105,15 @@ def __init__(
         model_path: str,
         imgsz: tuple[int, int],
         use_rvc2: bool,
+        output_dir: str | None = None,
     ):
         super().__init__(
             model_path,
             imgsz,
             use_rvc2,
             subtype="yolov8",
             output_names=["output1_yolov6r2", "output2_yolov6r2", "output3_yolov6r2"],
+            output_dir=output_dir,
         )
         self.load_model()
 

tools/yolov6r1/yolov6_r1_exporter.py

@@ -44,13 +44,15 @@ def __init__(
         model_path: str,
         imgsz: tuple[int, int],
         use_rvc2: bool,
+        output_dir: str | None = None,
     ):
         super().__init__(
             model_path,
             imgsz,
             use_rvc2,
             subtype="yolov6",
             output_names=["output1_yolov6", "output2_yolov6", "output3_yolov6"],
+            output_dir=output_dir,
         )
         self.load_model()