"""Anomalib CLI.""" # Copyright (C) 2023-2024 Intel Corporation # SPDX-License-Identifier: Apache-2.0 import logging from collections.abc import Callable, Sequence from functools import partial from pathlib import Path from types import MethodType from typing import Any from jsonargparse import ActionConfigFile, ArgumentParser, Namespace from jsonargparse._actions import _ActionSubCommands from rich import traceback from anomalib import TaskType, __version__ from anomalib.cli.utils.help_formatter import CustomHelpFormatter, get_short_docstring from anomalib.cli.utils.openvino import add_openvino_export_arguments from anomalib.loggers import configure_logger traceback.install() logger = logging.getLogger("anomalib.cli") _LIGHTNING_AVAILABLE = True try: from lightning.pytorch import Trainer from torch.utils.data import DataLoader, Dataset from anomalib.data import AnomalibDataModule from anomalib.engine import Engine from anomalib.metrics.threshold import BaseThreshold from anomalib.models import AnomalyModule from anomalib.utils.config import update_config except ImportError: _LIGHTNING_AVAILABLE = False class AnomalibCLI: """Implementation of a fully configurable CLI tool for anomalib. The advantage of this tool is its flexibility to configure the pipeline from both the CLI and a configuration file (.yaml or .json). It is even possible to use both the CLI and a configuration file simultaneously. For more details, the reader could refer to PyTorch Lightning CLI documentation. ``save_config_kwargs`` is set to ``overwrite=True`` so that the ``SaveConfigCallback`` overwrites the config if it already exists. """ def __init__(self, args: Sequence[str] | None = None) -> None: self.parser = self.init_parser() self.subcommand_parsers: dict[str, ArgumentParser] = {} self.subcommand_method_arguments: dict[str, list[str]] = {} self.add_subcommands() self.config = self.parser.parse_args(args=args) self.subcommand = self.config["subcommand"] if _LIGHTNING_AVAILABLE: self.before_instantiate_classes() self.instantiate_classes() self._run_subcommand() def init_parser(self, **kwargs) -> ArgumentParser: """Method that instantiates the argument parser.""" kwargs.setdefault("dump_header", [f"anomalib=={__version__}"]) parser = ArgumentParser(formatter_class=CustomHelpFormatter, **kwargs) parser.add_argument( "-c", "--config", action=ActionConfigFile, help="Path to a configuration file in json or yaml format.", ) return parser @staticmethod def subcommands() -> dict[str, set[str]]: """Skip predict subcommand as it is added later.""" return { "fit": {"model", "train_dataloaders", "val_dataloaders", "datamodule"}, "validate": {"model", "dataloaders", "datamodule"}, "test": {"model", "dataloaders", "datamodule"}, } @staticmethod def anomalib_subcommands() -> dict[str, dict[str, str]]: """Return a dictionary of subcommands and their description.""" return { "train": {"description": "Fit the model and then call test on the trained model."}, "predict": {"description": "Run inference on a model."}, "export": {"description": "Export the model to ONNX or OpenVINO format."}, } def add_subcommands(self, **kwargs) -> None: """Initialize base subcommands and add anomalib specific on top of it.""" parser_subcommands = self.parser.add_subcommands() # Extra subcommand: install self._set_install_subcommand(parser_subcommands) if not _LIGHTNING_AVAILABLE: # If environment is not configured to use pl, do not add a subcommand for Engine. return # Add Trainer subcommands for subcommand in self.subcommands(): sub_parser = self.init_parser(**kwargs) fn = getattr(Trainer, subcommand) # extract the first line description in the docstring for the subcommand help message description = get_short_docstring(fn) subparser_kwargs = kwargs.get(subcommand, {}) subparser_kwargs.setdefault("description", description) self.subcommand_parsers[subcommand] = sub_parser parser_subcommands.add_subcommand(subcommand, sub_parser, help=description) self.add_trainer_arguments(sub_parser, subcommand) # Add anomalib subcommands for subcommand in self.anomalib_subcommands(): sub_parser = self.init_parser(**kwargs) self.subcommand_parsers[subcommand] = sub_parser parser_subcommands.add_subcommand( subcommand, sub_parser, help=self.anomalib_subcommands()[subcommand]["description"], ) # add arguments to subcommand getattr(self, f"add_{subcommand}_arguments")(sub_parser) def add_arguments_to_parser(self, parser: ArgumentParser) -> None: """Extend trainer's arguments to add engine arguments. .. note:: Since ``Engine`` parameters are manually added, any change to the ``Engine`` class should be reflected manually. """ from anomalib.callbacks.normalization import get_normalization_callback parser.add_function_arguments(get_normalization_callback, "normalization") parser.add_argument("--task", type=TaskType | str, default=TaskType.SEGMENTATION) parser.add_argument( "--metrics.image", type=list[str] | str | dict[str, dict[str, Any]] | None, default=["F1Score", "AUROC"], ) parser.add_argument( "--metrics.pixel", type=list[str] | str | dict[str, dict[str, Any]] | None, default=None, required=False, ) parser.add_argument("--metrics.threshold", type=BaseThreshold | str, default="F1AdaptiveThreshold") parser.add_argument("--logging.log_graph", type=bool, help="Log the model to the logger", default=False) if hasattr(parser, "subcommand") and parser.subcommand not in ("export", "predict"): parser.link_arguments("task", "data.init_args.task") parser.add_argument( "--default_root_dir", type=Path, help="Path to save the results.", default=Path("./results"), ) parser.link_arguments("default_root_dir", "trainer.default_root_dir") # TODO(ashwinvaidya17): Tiling should also be a category of its own # CVS-122659 def add_trainer_arguments(self, parser: ArgumentParser, subcommand: str) -> None: """Add train arguments to the parser.""" self._add_default_arguments_to_parser(parser) self._add_trainer_arguments_to_parser(parser, add_optimizer=True, add_scheduler=True) parser.add_subclass_arguments( AnomalyModule, "model", fail_untyped=False, required=True, ) parser.add_subclass_arguments(AnomalibDataModule, "data") self.add_arguments_to_parser(parser) skip: set[str | int] = set(self.subcommands()[subcommand]) added = parser.add_method_arguments( Trainer, subcommand, skip=skip, ) self.subcommand_method_arguments[subcommand] = added def add_train_arguments(self, parser: ArgumentParser) -> None: """Add train arguments to the parser.""" self._add_default_arguments_to_parser(parser) self._add_trainer_arguments_to_parser(parser, add_optimizer=True, add_scheduler=True) parser.add_subclass_arguments( AnomalyModule, "model", fail_untyped=False, required=True, ) parser.add_subclass_arguments(AnomalibDataModule, "data") self.add_arguments_to_parser(parser) added = parser.add_method_arguments( Engine, "train", skip={"model", "datamodule", "val_dataloaders", "test_dataloaders", "train_dataloaders"}, ) self.subcommand_method_arguments["train"] = added def add_predict_arguments(self, parser: ArgumentParser) -> None: """Add predict arguments to the parser.""" self._add_default_arguments_to_parser(parser) self._add_trainer_arguments_to_parser(parser) parser.add_subclass_arguments( AnomalyModule, "model", fail_untyped=False, required=True, ) parser.add_argument( "--data", type=Dataset | AnomalibDataModule | DataLoader | str | Path, required=True, ) added = parser.add_method_arguments( Engine, "predict", skip={"model", "dataloaders", "datamodule", "dataset", "data_path"}, ) self.subcommand_method_arguments["predict"] = added self.add_arguments_to_parser(parser) def add_export_arguments(self, parser: ArgumentParser) -> None: """Add export arguments to the parser.""" self._add_default_arguments_to_parser(parser) self._add_trainer_arguments_to_parser(parser) parser.add_subclass_arguments( AnomalyModule, "model", fail_untyped=False, required=True, ) added = parser.add_method_arguments( Engine, "export", skip={"ov_args", "model"}, ) self.subcommand_method_arguments["export"] = added add_openvino_export_arguments(parser) self.add_arguments_to_parser(parser) def _set_install_subcommand(self, action_subcommand: _ActionSubCommands) -> None: sub_parser = ArgumentParser(formatter_class=CustomHelpFormatter) sub_parser.add_argument( "--option", help="Install the full or optional-dependencies.", default="full", type=str, choices=["full", "core", "dev", "loggers", "notebooks", "openvino"], ) sub_parser.add_argument( "-v", "--verbose", help="Set Logger level to INFO", action="store_true", ) self.subcommand_parsers["install"] = sub_parser action_subcommand.add_subcommand( "install", sub_parser, help="Install the full-package for anomalib.", ) def before_instantiate_classes(self) -> None: """Modify the configuration to properly instantiate classes and sets up tiler.""" subcommand = self.config["subcommand"] if subcommand in (*self.subcommands(), "train", "predict"): self.config[subcommand] = update_config(self.config[subcommand]) def instantiate_classes(self) -> None: """Instantiate classes depending on the subcommand. For trainer related commands it instantiates all the model, datamodule and trainer classes. But for subcommands we do not want to instantiate any trainer specific classes such as datamodule, model, etc This is because the subcommand is responsible for instantiating and executing code based on the passed config """ if self.config["subcommand"] in (*self.subcommands(), "predict"): # trainer commands # since all classes are instantiated, the LightningCLI also creates an unused ``Trainer`` object. # the minor change here is that engine is instantiated instead of trainer self.config_init = self.parser.instantiate_classes(self.config) self.datamodule = self._get(self.config_init, "data") if isinstance(self.datamodule, Dataset): self.datamodule = DataLoader(self.datamodule) self.model = self._get(self.config_init, "model") self._configure_optimizers_method_to_model() self.instantiate_engine() else: self.config_init = self.parser.instantiate_classes(self.config) subcommand = self.config["subcommand"] if subcommand in ("train", "export"): self.instantiate_engine() if "model" in self.config_init[subcommand]: self.model = self._get(self.config_init, "model") else: self.model = None if "data" in self.config_init[subcommand]: self.datamodule = self._get(self.config_init, "data") else: self.datamodule = None def instantiate_engine(self) -> None: """Instantiate the engine. .. note:: Most of the code in this method is taken from ``LightningCLI``'s ``instantiate_trainer`` method. Refer to that method for more details. """ from lightning.pytorch.cli import SaveConfigCallback from anomalib.callbacks import get_callbacks engine_args = { "normalization": self._get(self.config_init, "normalization.normalization_method"), "threshold": self._get(self.config_init, "metrics.threshold"), "task": self._get(self.config_init, "task"), "image_metrics": self._get(self.config_init, "metrics.image"), "pixel_metrics": self._get(self.config_init, "metrics.pixel"), } trainer_config = {**self._get(self.config_init, "trainer", default={}), **engine_args} key = "callbacks" if key in trainer_config: if trainer_config[key] is None: trainer_config[key] = [] elif not isinstance(trainer_config[key], list): trainer_config[key] = [trainer_config[key]] if not trainer_config.get("fast_dev_run", False): config_callback = SaveConfigCallback( self._parser(self.subcommand), self.config.get(str(self.subcommand), self.config), overwrite=True, ) trainer_config[key].append(config_callback) trainer_config[key].extend(get_callbacks(self.config[self.subcommand])) self.engine = Engine(**trainer_config) def _run_subcommand(self) -> None: """Run subcommand depending on the subcommand. This overrides the original ``_run_subcommand`` to run the ``Engine`` method rather than the ``Train`` method. """ if self.subcommand == "install": from anomalib.cli.install import anomalib_install install_kwargs = self.config.get("install", {}) anomalib_install(**install_kwargs) elif self.config["subcommand"] in (*self.subcommands(), "train", "export", "predict"): fn = getattr(self.engine, self.subcommand) fn_kwargs = self._prepare_subcommand_kwargs(self.subcommand) fn(**fn_kwargs) else: self.config_init = self.parser.instantiate_classes(self.config) getattr(self, f"{self.subcommand}")() @property def fit(self) -> Callable: """Fit the model using engine's fit method.""" return self.engine.fit @property def validate(self) -> Callable: """Validate the model using engine's validate method.""" return self.engine.validate @property def test(self) -> Callable: """Test the model using engine's test method.""" return self.engine.test @property def predict(self) -> Callable: """Predict using engine's predict method.""" return self.engine.predict @property def train(self) -> Callable: """Train the model using engine's train method.""" return self.engine.train @property def export(self) -> Callable: """Export the model using engine's export method.""" return self.engine.export def _add_trainer_arguments_to_parser( self, parser: ArgumentParser, add_optimizer: bool = False, add_scheduler: bool = False, ) -> None: """Add trainer arguments to the parser.""" parser.add_class_arguments(Trainer, "trainer", fail_untyped=False, instantiate=False, sub_configs=True) if add_optimizer: from torch.optim import Optimizer optim_kwargs = {"instantiate": False, "fail_untyped": False, "skip": {"params"}} parser.add_subclass_arguments( baseclass=(Optimizer,), nested_key="optimizer", **optim_kwargs, ) if add_scheduler: from lightning.pytorch.cli import LRSchedulerTypeTuple scheduler_kwargs = {"instantiate": False, "fail_untyped": False, "skip": {"optimizer"}} parser.add_subclass_arguments( baseclass=LRSchedulerTypeTuple, nested_key="lr_scheduler", **scheduler_kwargs, ) def _add_default_arguments_to_parser(self, parser: ArgumentParser) -> None: """Adds default arguments to the parser.""" parser.add_argument( "--seed_everything", type=bool | int, default=True, help=( "Set to an int to run seed_everything with this value before classes instantiation." "Set to True to use a random seed." ), ) def _get(self, config: Namespace, key: str, default: Any = None) -> Any: # noqa: ANN401 """Utility to get a config value which might be inside a subcommand.""" return config.get(str(self.subcommand), config).get(key, default) def _prepare_subcommand_kwargs(self, subcommand: str) -> dict[str, Any]: """Prepares the keyword arguments to pass to the subcommand to run.""" fn_kwargs = { k: v for k, v in self.config_init[subcommand].items() if k in self.subcommand_method_arguments[subcommand] } fn_kwargs["model"] = self.model if self.datamodule is not None: if isinstance(self.datamodule, AnomalibDataModule): fn_kwargs["datamodule"] = self.datamodule elif isinstance(self.datamodule, DataLoader): fn_kwargs["dataloaders"] = self.datamodule elif isinstance(self.datamodule, Path | str): fn_kwargs["data_path"] = self.datamodule return fn_kwargs def _parser(self, subcommand: str | None) -> ArgumentParser: if subcommand is None: return self.parser # return the subcommand parser for the subcommand passed return self.subcommand_parsers[subcommand] def _configure_optimizers_method_to_model(self) -> None: from lightning.pytorch.cli import LightningCLI, instantiate_class optimizer_cfg = self._get(self.config_init, "optimizer", None) if optimizer_cfg is None: return lr_scheduler_cfg = self._get(self.config_init, "lr_scheduler", {}) optimizer = instantiate_class(self.model.parameters(), optimizer_cfg) lr_scheduler = instantiate_class(optimizer, lr_scheduler_cfg) if lr_scheduler_cfg else None fn = partial(LightningCLI.configure_optimizers, optimizer=optimizer, lr_scheduler=lr_scheduler) # override the existing method self.model.configure_optimizers = MethodType(fn, self.model) def main() -> None: """Trainer via Anomalib CLI.""" configure_logger() AnomalibCLI() if __name__ == "__main__": main()