From 399cd76208a65df183d7eaa3c67d3be923ff250a Mon Sep 17 00:00:00 2001 From: klemen1999 Date: Fri, 7 Aug 2026 10:37:08 +0200 Subject: [PATCH 1/3] Added --output-dir option to CLI --- README.md | 2 ++ tests/test_unittests.py | 22 ++++++++++++ tools/conversion_registry.py | 52 ++++++++++++++++------------ tools/main.py | 7 ++++ tools/modules/exporter.py | 7 +++- tools/utils/config.py | 4 +++ tools/yolo/yolo26_exporter.py | 9 ++++- tools/yolo/yolov10_exporter.py | 2 ++ tools/yolo/yolov5_exporter.py | 2 ++ tools/yolo/yolov6_exporter.py | 2 ++ tools/yolo/yolov8_exporter.py | 2 ++ tools/yolov6r1/yolov6_r1_exporter.py | 2 ++ tools/yolov6r3/gold_yolo_exporter.py | 2 ++ tools/yolov6r3/yolov6_r3_exporter.py | 2 ++ tools/yolov7/yolov7_exporter.py | 2 ++ tools/yolox/yolox_exporter.py | 2 ++ 16 files changed, 96 insertions(+), 25 deletions(-) diff --git a/README.md b/README.md index da5228b..a6e2cd0 100644 --- a/README.md +++ b/README.md @@ -95,6 +95,8 @@ Tools CLI │ 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 │ diff --git a/tests/test_unittests.py b/tests/test_unittests.py index cdedb1a..0abc2df 100644 --- a/tests/test_unittests.py +++ b/tests/test_unittests.py @@ -224,6 +224,28 @@ 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" + 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 not Path(_output_dir(test_workspace)).exists() + + def test_wrong_explicit_class_names(test_workspace: Path): """Tests setting wrong explicit class names.""" model_name = "yolov8n" diff --git a/tools/conversion_registry.py b/tools/conversion_registry.py index 96afa11..fd0a033 100644 --- a/tools/conversion_registry.py +++ b/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 + ) diff --git a/tools/main.py b/tools/main.py index f81fd18..9e5acb8 100644 --- a/tools/main.py +++ b/tools/main.py @@ -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: diff --git a/tools/modules/exporter.py b/tools/modules/exporter.py index 484abdc..2b3c92e 100644 --- a/tools/modules/exporter.py +++ b/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 diff --git a/tools/utils/config.py b/tools/utils/config.py index 3e3f87f..e598543 100644 --- a/tools/utils/config.py +++ b/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." ) diff --git a/tools/yolo/yolo26_exporter.py b/tools/yolo/yolo26_exporter.py index 3b14a7c..8eb422a 100644 --- a/tools/yolo/yolo26_exporter.py +++ b/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() diff --git a/tools/yolo/yolov10_exporter.py b/tools/yolo/yolov10_exporter.py index 6d6adc9..8bbecf9 100644 --- a/tools/yolo/yolov10_exporter.py +++ b/tools/yolo/yolov10_exporter.py @@ -23,6 +23,7 @@ def __init__( model_path: str, imgsz: tuple[int, int], use_rvc2: bool, + output_dir: str | None = None, ): super().__init__( model_path, @@ -30,6 +31,7 @@ def __init__( use_rvc2, subtype="yolov10", output_names=["output1_yolov10", "output2_yolov10", "output3_yolov10"], + output_dir=output_dir, ) self.load_model() diff --git a/tools/yolo/yolov5_exporter.py b/tools/yolo/yolov5_exporter.py index 2f79586..9d10a91 100644 --- a/tools/yolo/yolov5_exporter.py +++ b/tools/yolo/yolov5_exporter.py @@ -91,6 +91,7 @@ def __init__( model_path: str, imgsz: tuple[int, int], use_rvc2: bool, + output_dir: str | None = None, ): super().__init__( model_path, @@ -98,6 +99,7 @@ def __init__( use_rvc2, subtype="yolov5", output_names=["output1_yolov5", "output2_yolov5", "output3_yolov5"], + output_dir=output_dir, ) self.load_model() diff --git a/tools/yolo/yolov6_exporter.py b/tools/yolo/yolov6_exporter.py index 283188f..693c776 100644 --- a/tools/yolo/yolov6_exporter.py +++ b/tools/yolo/yolov6_exporter.py @@ -45,6 +45,7 @@ def __init__( model_path: str, imgsz: tuple[int, int], use_rvc2: bool, + output_dir: str | None = None, ): super().__init__( model_path, @@ -52,6 +53,7 @@ def __init__( use_rvc2, subtype="yolov6r2", output_names=["output1_yolov6r2", "output2_yolov6r2", "output3_yolov6r2"], + output_dir=output_dir, ) self.load_model() diff --git a/tools/yolo/yolov8_exporter.py b/tools/yolo/yolov8_exporter.py index 1384245..a03b722 100644 --- a/tools/yolo/yolov8_exporter.py +++ b/tools/yolo/yolov8_exporter.py @@ -105,6 +105,7 @@ def __init__( model_path: str, imgsz: tuple[int, int], use_rvc2: bool, + output_dir: str | None = None, ): super().__init__( model_path, @@ -112,6 +113,7 @@ def __init__( use_rvc2, subtype="yolov8", output_names=["output1_yolov6r2", "output2_yolov6r2", "output3_yolov6r2"], + output_dir=output_dir, ) self.load_model() diff --git a/tools/yolov6r1/yolov6_r1_exporter.py b/tools/yolov6r1/yolov6_r1_exporter.py index 898264c..e5baf9d 100644 --- a/tools/yolov6r1/yolov6_r1_exporter.py +++ b/tools/yolov6r1/yolov6_r1_exporter.py @@ -44,6 +44,7 @@ def __init__( model_path: str, imgsz: tuple[int, int], use_rvc2: bool, + output_dir: str | None = None, ): super().__init__( model_path, @@ -51,6 +52,7 @@ def __init__( use_rvc2, subtype="yolov6", output_names=["output1_yolov6", "output2_yolov6", "output3_yolov6"], + output_dir=output_dir, ) self.load_model() diff --git a/tools/yolov6r3/gold_yolo_exporter.py b/tools/yolov6r3/gold_yolo_exporter.py index f3ac412..04fae40 100644 --- a/tools/yolov6r3/gold_yolo_exporter.py +++ b/tools/yolov6r3/gold_yolo_exporter.py @@ -49,6 +49,7 @@ def __init__( model_path: str, imgsz: tuple[int, int], use_rvc2: bool, + output_dir: str | None = None, ): super().__init__( model_path, @@ -56,6 +57,7 @@ def __init__( use_rvc2, subtype="yolov6r2", output_names=["output1_yolov6r2", "output2_yolov6r2", "output3_yolov6r2"], + output_dir=output_dir, ) self.load_model() diff --git a/tools/yolov6r3/yolov6_r3_exporter.py b/tools/yolov6r3/yolov6_r3_exporter.py index c71c19f..a84c0b2 100644 --- a/tools/yolov6r3/yolov6_r3_exporter.py +++ b/tools/yolov6r3/yolov6_r3_exporter.py @@ -56,6 +56,7 @@ def __init__( model_path: str, imgsz: tuple[int, int], use_rvc2: bool, + output_dir: str | None = None, ): super().__init__( model_path, @@ -63,6 +64,7 @@ def __init__( use_rvc2, subtype="yolov6r2", output_names=["output1_yolov6r2", "output2_yolov6r2", "output3_yolov6r2"], + output_dir=output_dir, ) self.load_model() diff --git a/tools/yolov7/yolov7_exporter.py b/tools/yolov7/yolov7_exporter.py index c43e549..ee8a61b 100644 --- a/tools/yolov7/yolov7_exporter.py +++ b/tools/yolov7/yolov7_exporter.py @@ -60,6 +60,7 @@ def __init__( model_path: str, imgsz: tuple[int, int], use_rvc2: bool, + output_dir: str | None = None, ): super().__init__( model_path, @@ -67,6 +68,7 @@ def __init__( use_rvc2, subtype="yolov7", output_names=["output1_yolov7", "output2_yolov7", "output3_yolov7"], + output_dir=output_dir, ) self.load_model() diff --git a/tools/yolox/yolox_exporter.py b/tools/yolox/yolox_exporter.py index b28bdff..8bbaaa3 100644 --- a/tools/yolox/yolox_exporter.py +++ b/tools/yolox/yolox_exporter.py @@ -174,6 +174,7 @@ def __init__( model_path: str, imgsz: tuple[int, int], use_rvc2: bool, + output_dir: str | None = None, ): super().__init__( model_path, @@ -182,6 +183,7 @@ def __init__( # YOLOX uses the same grid decode as this already-supported subtype. subtype="yolov6r1", output_names=self.output_names, + output_dir=output_dir, ) self.load_model() From 369fd7c53e3d5fe8e983aa3d8eb98c580ca842e2 Mon Sep 17 00:00:00 2001 From: klemen1999 Date: Fri, 7 Aug 2026 10:44:17 +0200 Subject: [PATCH 2/3] fix for --help --- README.md | 2 +- tests/test_unittests.py | 8 ++++++++ tools/main.py | 2 +- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index a6e2cd0..b0022f4 100644 --- a/README.md +++ b/README.md @@ -90,7 +90,7 @@ 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 │ diff --git a/tests/test_unittests.py b/tests/test_unittests.py index 0abc2df..ae7d2b3 100644 --- a/tests/test_unittests.py +++ b/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"), diff --git a/tools/main.py b/tools/main.py index 9e5acb8..24c889b 100644 --- a/tools/main.py +++ b/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, From b4f21056723af18e435b02ed7a00e0b0485dcefb Mon Sep 17 00:00:00 2001 From: klemen1999 Date: Fri, 7 Aug 2026 11:48:29 +0200 Subject: [PATCH 3/3] coderabbit fix --- tests/test_unittests.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/test_unittests.py b/tests/test_unittests.py index ae7d2b3..9dabe5d 100644 --- a/tests/test_unittests.py +++ b/tests/test_unittests.py @@ -236,6 +236,16 @@ 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, @@ -251,7 +261,15 @@ def test_explicit_output_dir(test_workspace: Path): pytest.fail(f"Exit code: {result.returncode}, Output: {result.stdout}") nn_archive_checker(output_dir=str(output_dir)) - assert not Path(_output_dir(test_workspace)).exists() + 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):