From 4b3f77cb7e88f9b83f1f5d667a8f3b0586f51d20 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Thu, 25 Jul 2024 22:07:42 +0200 Subject: [PATCH 1/3] added archive command to CLI --- modelconverter/__main__.py | 54 ++++++++++++++++++++++++++++-- modelconverter/utils/__init__.py | 2 ++ modelconverter/utils/nn_archive.py | 46 +++++++++++++++++++++++++ 3 files changed, 100 insertions(+), 2 deletions(-) diff --git a/modelconverter/__main__.py b/modelconverter/__main__.py index 3b945bd..84c3896 100644 --- a/modelconverter/__main__.py +++ b/modelconverter/__main__.py @@ -10,7 +10,7 @@ from luxonis_ml.nn_archive import ArchiveGenerator, is_nn_archive from luxonis_ml.nn_archive.config import Config as NNArchiveConfig from luxonis_ml.nn_archive.config_building_blocks import PreprocessingBlock -from luxonis_ml.utils import reset_logging, setup_logging +from luxonis_ml.utils import LuxonisFileSystem, reset_logging, setup_logging from typing_extensions import Annotated, TypeAlias from modelconverter.packages import ( @@ -20,6 +20,7 @@ ) from modelconverter.utils import ( ModelconverterException, + archive_from_model, docker_build, docker_exec, in_docker, @@ -38,6 +39,8 @@ ) from modelconverter.utils.types import Target +logger = logging.getLogger(__name__) + app = typer.Typer( help="Modelconverter CLI", add_completion=False, @@ -145,7 +148,6 @@ def get_configs( @return: Tuple of the parsed modelconverter L{Config}, L{NNArchiveConfig} and the main stage key. """ - logger = logging.getLogger(__name__) for p in [CONFIGS_DIR, MODELS_DIR, OUTPUTS_DIR, CALIBRATION_DIR]: logger.debug(f"Creating {p}") @@ -252,6 +254,7 @@ def infer( mult_cfg, _, _ = get_configs(path, opts) cfg = mult_cfg.get_stage_config(stage) Inferer = get_inferer(target) + assert output_dir is not None Inferer.from_config( model_path, input_path, Path(output_dir), cfg ).run() @@ -506,6 +509,53 @@ def convert( docker_exec(target.value, *args, tag=tag, use_gpu=gpu) +@app.command() +def archive( + path: Annotated[str, typer.Argument(help="Path to the model file.")], + save_path: Annotated[ + Optional[str], + typer.Option( + "-s", + "--save-path", + help="Path or URL to save the archive to. " + "By default, it is saved to the current directory " + "under the name of the model.", + ), + ] = None, + put_file_plugin: Annotated[ + Optional[str], + typer.Option( + help="The name of the plugin to use for uploading the file." + ), + ] = None, +): + setup_logging(use_rich=True) + model_path = resolve_path(path, MODELS_DIR) + cfg = archive_from_model(model_path) + save_path = save_path or f"{cfg.model.metadata.name}.tar.xz" + if not save_path.endswith(".tar.xz"): + save_path += f"/{cfg.model.metadata.name}.tar.xz" + archive_name = save_path.split("/")[-1] + protocol = LuxonisFileSystem.get_protocol(save_path) + if protocol != "file": + archive_save_path = "./" + else: + archive_save_path = str(Path(save_path).parent) + archive_save_path = ArchiveGenerator( + archive_name=archive_name, + save_path=archive_save_path, + cfg_dict=cfg.model_dump(), + executables_paths=[str(model_path)], + ).make_archive() + + if protocol != "file": + upload_file_to_remote(archive_save_path, save_path, put_file_plugin) + Path(archive_save_path).unlink() + logger.info(f"Archive uploaded to {save_path}") + else: + logger.info(f"Archive saved to {save_path}") + + def version_callback(value: bool): if value: typer.echo(f"ModelConverter Version: {version(__package__)}") diff --git a/modelconverter/utils/__init__.py b/modelconverter/utils/__init__.py index 549616f..5efd3ce 100644 --- a/modelconverter/utils/__init__.py +++ b/modelconverter/utils/__init__.py @@ -25,6 +25,7 @@ get_archive_input, modelconverter_config_to_nn, process_nn_archive, + archive_from_model, ) from .onnx_tools import onnx_attach_normalization_to_inputs from .subprocess import subprocess_run @@ -55,4 +56,5 @@ "make_default_layout", "Metadata", "get_metadata", + "archive_from_model" ] diff --git a/modelconverter/utils/nn_archive.py b/modelconverter/utils/nn_archive.py index 42e995d..c834965 100644 --- a/modelconverter/utils/nn_archive.py +++ b/modelconverter/utils/nn_archive.py @@ -218,3 +218,49 @@ def modelconverter_config_to_nn( f"{post_stage_key}{model_name.suffix}" ) return archive + + +def archive_from_model(model_path: Path) -> NNArchiveConfig: + metadata = get_metadata(model_path) + + archive_cfg = { + "config_version": "1.0", + "model": { + "metadata": { + "name": model_path.stem, + "path": model_path.name, + }, + "inputs": [], + "outputs": [], + "heads": [], + }, + } + + for name, shape in metadata.input_shapes.items(): + archive_cfg["model"]["inputs"].append( + { + "name": name, + "shape": shape, + "layout": make_default_layout(shape), + "dtype": "float32", + "input_type": "image", + "preprocessing": { + "mean": None, + "scale": None, + "reverse_channels": False, + "interleaved_to_planar": False, + }, + } + ) + + for name, shape in metadata.output_shapes.items(): + archive_cfg["model"]["outputs"].append( + { + "name": name, + "shape": shape, + "layout": make_default_layout(shape), + "dtype": "float32", + } + ) + + return NNArchiveConfig(**archive_cfg) From 76f9ef1da08ba92ba1a91085010b2e733c69b952 Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Thu, 25 Jul 2024 22:19:44 +0200 Subject: [PATCH 2/3] updated help --- modelconverter/__main__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/modelconverter/__main__.py b/modelconverter/__main__.py index 84c3896..f6a1763 100644 --- a/modelconverter/__main__.py +++ b/modelconverter/__main__.py @@ -511,7 +511,9 @@ def convert( @app.command() def archive( - path: Annotated[str, typer.Argument(help="Path to the model file.")], + path: Annotated[ + str, typer.Argument(help="Path or an URL of the model file.") + ], save_path: Annotated[ Optional[str], typer.Option( From 3eab824a594d779f8f6f0590ff37e0e2a4cd5edc Mon Sep 17 00:00:00 2001 From: Martin Kozlovsky Date: Thu, 25 Jul 2024 22:33:39 +0200 Subject: [PATCH 3/3] added compression --- modelconverter/__main__.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/modelconverter/__main__.py b/modelconverter/__main__.py index f6a1763..a596240 100644 --- a/modelconverter/__main__.py +++ b/modelconverter/__main__.py @@ -535,8 +535,17 @@ def archive( model_path = resolve_path(path, MODELS_DIR) cfg = archive_from_model(model_path) save_path = save_path or f"{cfg.model.metadata.name}.tar.xz" - if not save_path.endswith(".tar.xz"): - save_path += f"/{cfg.model.metadata.name}.tar.xz" + if save_path.endswith("tar.xz"): + compression = "xz" + elif save_path.endswith("tar.gz"): + compression = "gz" + elif save_path.endswith("tar.bz2"): + compression = "bz2" + else: + compression = "xz" + + if not save_path.endswith(f".tar.{compression}"): + save_path += f"/{cfg.model.metadata.name}.tar.{compression}" archive_name = save_path.split("/")[-1] protocol = LuxonisFileSystem.get_protocol(save_path) if protocol != "file": @@ -546,6 +555,7 @@ def archive( archive_save_path = ArchiveGenerator( archive_name=archive_name, save_path=archive_save_path, + compression=compression, cfg_dict=cfg.model_dump(), executables_paths=[str(model_path)], ).make_archive()