sinapsis-anomalib 0.1.0__py3-none-any.whl

sinapsis_anomalib/templates/anomalib_export.py

@@ -0,0 +1,204 @@
+# -*- coding: utf-8 -*-
+from pathlib import Path
+from typing import Any
+
+from anomalib.deploy import CompressionType, ExportType
+from pydantic.dataclasses import dataclass
+from sinapsis_core.data_containers.data_packet import DataContainer
+from sinapsis_core.template_base import Template, TemplateAttributeType
+from sinapsis_core.template_base.dynamic_template_factory import make_dynamic_template
+from sinapsis_core.utils.env_var_keys import SINAPSIS_BUILD_DOCS
+from torchmetrics import Metric
+
+from sinapsis_anomalib.templates.anomalib_base import (
+    AnomalibBase,
+    AnomalibBaseAttributes,
+)
+from sinapsis_anomalib.templates.anomalib_train import AnomalibTrainDataClass
+
+
+@dataclass(frozen=True, slots=True)
+class AnomalibExportDataClass:
+    """Container for export results.
+
+    Attributes:
+        exported_model_path (Path | str): Path to the exported model file(s).
+            Can be either a Path object or string path.
+    """
+
+    exported_model_path: Path | str
+
+
+class AnomalibExportAttributes(AnomalibBaseAttributes):
+    """Export-specific attribute configuration for Anomalib models.
+
+    Attributes:
+        export_type (ExportType | str): Target export format. Defaults to TORCH.
+        export_root (str | Path | None): Root directory for exported files.
+        input_size (tuple[int, int] | None): Expected input dimensions (height, width).
+        compression_type (CompressionType | None): Quantization/compression method.
+        metric (Metric | str | None): Metric used for compression calibration.
+        ov_args (dict[str, Any] | None): OpenVINO-specific export arguments.
+        ckpt_path (str | None): Explicit path to model checkpoint.
+        generic_key_chkpt (str | None): Key to retrieve training results.
+    """
+
+    export_type: ExportType = ExportType.TORCH
+    export_root: str | Path | None = None
+    input_size: tuple[int, int] | None = None
+    compression_type: CompressionType | None = None
+    metric: Metric | str | None = None
+    ov_args: dict[str, Any] | None = None
+    ckpt_path: str | None = None
+    generic_key_chkpt: str | None = None
+
+
+class AnomalibExport(AnomalibBase):
+    """Export functionality for trained Anomalib models.
+
+    Usage example:
+
+        agent:
+        name: my_test_agent
+        templates:
+        - template_name: InputTemplate
+        class_name: InputTemplate
+        attributes: {}
+        - template_name: CfaExport
+        class_name: CfaExport
+        template_input: InputTemplate
+        attributes:
+            folder_attributes_config_path: null
+            generic_key: 'my_generic_key'
+            callbacks: null
+            normalization: null
+            threshold: null
+            task: null
+            image_metrics: null
+            pixel_metrics: null
+            logger: null
+            default_root_dir: null
+            callback_configs: null
+            logger_configs: null
+            export_type: 'openvino'
+            export_root: null
+            input_size: null
+            transform: null
+            compression_type: null
+            metric: null
+            ov_args: null
+            ckpt_path: null
+            generic_key_chkpt: null
+            cfa_init:
+                backbone: wide_resnet50_2
+                gamma_c: 1
+                gamma_d: 1
+                num_nearest_neighbors: 3
+                num_hard_negative_features: 3
+                radius: 1.0e-05
+
+    For a full list of options use the sinapsis cli: sinapsis info --all-template-names
+    If you want to see all available models, please visit:
+    https://anomalib.readthedocs.io/en/v1.2.0/markdown/guides/reference/models/image/index.html
+
+    Notes:
+        - Supports multiple export formats via ExportType
+        - Enables model compression through CompressionType.
+        - Can load models from checkpoints or previous training results.
+    """
+
+    AttributesBaseModel = AnomalibExportAttributes
+    SUFFIX = "Export"
+
+    def __init__(self, attributes: TemplateAttributeType) -> None:
+        super().__init__(attributes)
+        self.data_module = None
+        if self.attributes.compression_type in (
+            CompressionType.INT8_ACQ,
+            CompressionType.INT8_PTQ,
+        ):
+            self.data_module = self.setup_data_loader()
+
+    def _get_checkpoint_path(self, container: DataContainer) -> str | Path:
+        """Resolves the checkpoint path for model export.
+
+        Args:
+            container (DataContainer): Container with potential training results
+
+        Raises:
+            ValueError: If no valid checkpoint path is found
+
+        Returns:
+            str | Path : Path to model checkpoint
+
+        Note:
+            Priority order:
+            1. Explicit ckpt_path from attributes
+            2. Training results from container (if generic_key_chkpt provided)
+        """
+        if self.attributes.ckpt_path:
+            return self.attributes.ckpt_path
+
+        generic_data = self._get_generic_data(container, self.attributes.generic_key_chkpt)
+        if generic_data and isinstance(generic_data, AnomalibTrainDataClass):
+            return generic_data.checkpoint_path
+
+        raise ValueError("No checkpoint path found")
+
+    def export_model(self, container: DataContainer) -> AnomalibExportDataClass:
+        """Exports the model to specified format.
+
+        Args:
+            container (DataContainer): Input data container
+
+        Returns:
+            AnomalibExportDataClass: Contains exported model path
+        """
+        ckpt_path = self._get_checkpoint_path(container)
+
+        exported_path = self.engine.export(
+            model=self.model,
+            export_type=self.attributes.export_type,
+            export_root=self.attributes.export_root,
+            input_size=self.attributes.input_size,
+            compression_type=self.attributes.compression_type,
+            datamodule=self.data_module,
+            metric=self.attributes.metric,
+            ov_args=self.attributes.ov_args,
+            ckpt_path=ckpt_path,
+        )
+
+        return AnomalibExportDataClass(exported_model_path=exported_path)
+
+    def execute(self, container: DataContainer) -> DataContainer:
+        """Performs model export and stores results.
+
+        Args:
+            container (DataContainer): Input data container
+
+        Returns:
+            DataContainer: Container with export results stored as generic data
+        """
+        result = self.export_model(container)
+        self._set_generic_data(container, result)
+        return container
+
+
+def __getattr__(name: str) -> Template:
+    """Creat dynamic templates.
+
+    Only create a template if it's imported, this avoids creating all the base models for all templates
+    and potential import errors due to not available packages.
+    """
+    if name in AnomalibExport.WrapperEntry.module_att_names:
+        return make_dynamic_template(name, AnomalibExport)
+    raise AttributeError(f"template `{name}` not found in {__name__}")
+
+
+__all__ = AnomalibExport.WrapperEntry.module_att_names
+
+if SINAPSIS_BUILD_DOCS:
+    dynamic_templates = [__getattr__(template_name) for template_name in __all__]
+    for template in dynamic_templates:
+        globals()[template.__name__] = template
+        del template

sinapsis_anomalib/templates/anomalib_openvino_inference.py

@@ -0,0 +1,93 @@
+# -*- coding: utf-8 -*-
+import os
+from typing import Literal
+
+import cv2 as cv
+import numpy as np
+import torchvision.transforms.v2 as T
+from anomalib.deploy import OpenVINOInferencer
+
+from sinapsis_anomalib.templates.anomalib_base_inference import (
+    AnomalibBaseInference,
+    AnomalibInferenceAttributes,
+)
+
+
+class AnomalibOpenVINOInferenceAttributes(AnomalibInferenceAttributes):
+    """OpenVINO-specific inference attribute configuration.
+
+    Attributes:
+        device (Literal["CPU", "GPU"]): Target hardware accelerator for inference.
+            Must be either 'CPU' or 'GPU'.
+        model_height (int): The image height expected by OV model.
+        model_width (int): The image width expected by OV model.
+    """
+
+    device: Literal["CPU", "GPU"]
+    model_height: int
+    model_width: int
+
+
+class AnomalibOpenVINOInference(AnomalibBaseInference):
+    """OpenVINO-specific inference implementation for Anomalib models.
+
+    Extends base inference to provide optimized model execution using OpenVINO toolkit.
+
+    Usage example:
+
+        agent:
+        name: my_test_agent
+        templates:
+        - template_name: InputTemplate
+        class_name: InputTemplate
+        attributes: {}
+        - template_name: AnomalibOpenVINOInference
+        class_name: AnomalibOpenVINOInference
+        template_input: InputTemplate
+        attributes:
+            model_path: 'path/to/model.xml'
+            transforms: null
+            device: CPU
+            model_height: 256
+            model_width: 256
+    """
+
+    AttributesBaseModel = AnomalibOpenVINOInferenceAttributes
+
+    def get_inferencer(self) -> OpenVINOInferencer:
+        """Initialize OpenVINO inferencer with model and metadata.
+
+        Returns:
+            OpenVINOInferencer: Inferencer instance with model and metadata loaded.
+        """
+        model_path: str = os.path.expanduser(self.attributes.model_path)
+
+        return OpenVINOInferencer(
+            path=model_path,
+            device=self.attributes.device,
+        )
+
+    @staticmethod
+    def postprocess_segmentation_mask(binary_mask: np.ndarray, image: np.ndarray) -> np.ndarray:
+        """Apply resizing and squeezing to mask.
+
+        Args:
+            binary_mask (np.ndarray): Mask produced by OpenVinoInfencer.
+            image (np.ndarray): Original image.
+
+        Returns:
+            np.ndarray: Postprocessed mask.
+        """
+        height, width = image.shape[:2]
+        binary_mask = np.squeeze(binary_mask).astype(np.uint8)
+        return cv.resize(binary_mask, (width, height))
+
+    def get_transformation_list(self) -> list[T.Transform]:
+        """Construct the list of transformations for the OpenVinoInferencer.
+
+        Returns:
+            list[T.Transform]: List of transformations.
+        """
+        transforms = [T.ToImage(), T.Resize([self.attributes.model_height, self.attributes.model_width])]
+        transforms.extend(self._convert_transform_names())
+        return transforms

sinapsis_anomalib/templates/anomalib_torch_inference.py

@@ -0,0 +1,83 @@
+# -*- coding: utf-8 -*-
+import os
+from typing import Literal
+
+import numpy as np
+import torch
+import torchvision.transforms.v2 as T
+from anomalib.deploy import TorchInferencer
+
+from sinapsis_anomalib.templates.anomalib_base_inference import (
+    AnomalibBaseInference,
+    AnomalibInferenceAttributes,
+)
+
+
+class AnomalibTorchInferenceAttributes(AnomalibInferenceAttributes):
+    """PyTorch-specific inference attribute configuration.
+
+    Attributes:
+        device (Literal["cuda", "cpu"]): Target device for inference (either 'cuda' or 'cpu').
+    """
+
+    device: Literal["cuda", "cpu"]
+
+
+class AnomalibTorchInference(AnomalibBaseInference):
+    """PyTorch-specific inference implementation for Anomalib models.
+
+    Extends base inference to provide native PyTorch model execution.
+
+    Usage example:
+
+        agent:
+        name: my_test_agent
+        templates:
+        - template_name: InputTemplate
+        class_name: InputTemplate
+        attributes: {}
+        - template_name: AnomalibTorchInference
+        class_name: AnomalibTorchInference
+        template_input: InputTemplate
+        attributes:
+            model_path: 'path/to/model.pt'
+            transforms: null
+            device: 'cuda'
+    """
+
+    AttributesBaseModel = AnomalibTorchInferenceAttributes
+
+    def get_inferencer(self) -> TorchInferencer:
+        """Get PyTorch Inferencer instance.
+
+        Returns:
+            TorchInferencer: Inferencer instance with model loaded on specified device.
+        """
+        model_path: str = os.path.expanduser(self.attributes.model_path)
+        return TorchInferencer(path=model_path, device=self.attributes.device)
+
+    @staticmethod
+    def postprocess_segmentation_mask(binary_mask: torch.TensorType, image: np.ndarray) -> np.ndarray:
+        """Apply rescaling, squeezing and conversion from torch to numpy array format.
+
+        Args:
+            binary_mask (np.ndarray): Mask produced by TorchInferencer.
+            image (np.ndarray): Input image.
+
+        Returns:
+            np.ndarray: Postprocessed mask.
+        """
+        height, width = image.shape[:2]
+        rescaled_mask = torch.squeeze(T.Resize(size=[height, width])(binary_mask))
+        return rescaled_mask.cpu().numpy().astype(np.uint8)
+
+    def get_transformation_list(self) -> list[T.Transform]:
+        """Construct the list of transformations for the TorchInferencer.
+
+        Returns:
+            list[T.Transform]: List of transformations.
+        """
+        transforms = [T.ToImage(), T.ToDtype(torch.float32, scale=True)]
+
+        transforms.extend(self._convert_transform_names())
+        return transforms

sinapsis_anomalib/templates/anomalib_train.py

@@ -0,0 +1,205 @@
+# -*- coding: utf-8 -*-
+from pathlib import Path
+from typing import Any, Literal
+
+from anomalib.engine.engine import _TrainerArgumentsCache
+from pydantic import Field
+from pydantic.dataclasses import dataclass
+from sinapsis_core.data_containers.data_packet import DataContainer
+from sinapsis_core.template_base import Template, TemplateAttributeType
+from sinapsis_core.template_base.dynamic_template_factory import make_dynamic_template
+from sinapsis_core.utils.env_var_keys import SINAPSIS_BUILD_DOCS
+
+from sinapsis_anomalib.templates.anomalib_base import (
+    AnomalibBase,
+    AnomalibBaseAttributes,
+)
+
+
+@dataclass(frozen=True)
+class AnomalibTrainKeys:
+    """Constants for accessing training-related configuration keys.
+
+    Attributes:
+        ACCELERATOR (str): Key for accessing accelerator field in trainer args dict.
+        DEVICES: (str): Key for accessing devices field in trainer args dict.
+        CALLBACK_METRICS (str): Key for accessing callback metrics ('callback_metrics')
+        MAX_EPOCHS (str): Key for maximum epochs setting ('max_epochs')
+        CKPT_PATH (str): Key for checkpoint path ('ckpt_path')
+        BEST_MODEL_PATH (str): Key for best model path ('best_model_path')
+    """
+
+    ACCELERATOR: str = "accelerator"
+    CALLBACK_METRICS: str = "callback_metrics"
+    MAX_EPOCHS: str = "max_epochs"
+    CKPT_PATH: str = "ckpt_path"
+    BEST_MODEL_PATH: str = "best_model_path"
+
+
+@dataclass(frozen=True, slots=True)
+class AnomalibTrainDataClass:
+    """Container for training results and artifacts.
+
+    Attributes:
+        metrics (dict[str, float]): Dictionary of training metrics (metric_name: value)
+        checkpoint_path (Path | str | None): Path to the best model checkpoint,
+            or None if no checkpoint was saved
+    """
+
+    metrics: dict[str, float]
+    checkpoint_path: Path | str | None
+
+
+class AnomalibTrainAttributes(AnomalibBaseAttributes):
+    """Training-specific configuration attributes.
+
+    Attributes:
+        max_epochs (int | None): Maximum number of training epochs.
+            If None, uses default from model configuration.
+        accelerator: (Literal["cpu", "gpu", "tpu", "hpu", "auto"]): Define the device to be used during training.
+            Defaults to "cpu".
+        ckpt_path (str | Path | None): Path to checkpoint for resuming training.
+            If None, starts training from scratch.
+        trainer_args: (dict[str, Any]): General trainer arguments. For more details see:
+            https://lightning.ai/docs/pytorch/stable/common/trainer.html#trainer-flags
+    """
+
+    max_epochs: int | None = None
+    accelerator: Literal["cpu", "gpu", "tpu", "hpu", "auto"] = "cpu"
+    ckpt_path: str | Path | None = None
+    trainer_args: dict[str, Any] = Field(default_factory=dict)
+
+
+class AnomalibTrain(AnomalibBase):
+    """Training implementation for Anomalib models.
+
+    Usage example:
+
+        agent:
+        name: my_test_agent
+        templates:
+        - template_name: InputTemplate
+        class_name: InputTemplate
+        attributes: {}
+        - template_name: CfaTrain
+        class_name: CfaTrain
+        template_input: InputTemplate
+        attributes:
+            folder_attributes_config_path: 'path/to/config.yaml'
+            generic_key: 'my_generic_key'
+            callbacks: null
+            normalization: null
+            threshold: null
+            image_metrics: null
+            pixel_metrics: null
+            logger: null
+            default_root_dir: null
+            callback_configs: null
+            logger_configs: null
+            max_epochs: null
+            ckpt_path: null
+            accelerator: gpu
+            trainer_args:
+                devices: "0"
+            cfa_init:
+                backbone: wide_resnet50_2
+                gamma_c: 1
+                gamma_d: 1
+                num_nearest_neighbors: 3
+                num_hard_negative_features: 3
+                radius: 1.0e-05
+    For a full list of options use the sinapsis cli: sinapsis info --all-template-names.
+    If you want to see all available models, please visit:
+    https://anomalib.readthedocs.io/en/v1.2.0/markdown/guides/reference/models/image/index.html.
+    """
+
+    SUFFIX = "Train"
+    AttributesBaseModel = AnomalibTrainAttributes
+
+    def __init__(self, attributes: TemplateAttributeType) -> None:
+        super().__init__(attributes)
+        self.data_module = self.setup_data_loader()
+
+    def _update_trainer_args(self) -> None:
+        """Updates the trainer configuration with current settings.
+
+        Specifically sets the maximum number of training epochs
+        based on the attributes configuration.
+        """
+        existing_args = self.engine._cache.args
+        existing_args[AnomalibTrainKeys.MAX_EPOCHS] = self.attributes.max_epochs
+
+        self.attributes.trainer_args[AnomalibTrainKeys.ACCELERATOR] = self.attributes.accelerator
+        existing_args.update(self.attributes.trainer_args)
+        self.engine._cache = _TrainerArgumentsCache(**existing_args)
+
+    def _get_training_metrics(self) -> dict[str, Any]:
+        """Extracts training metrics from the model.
+
+        Returns:
+            dict[str, float]: Dictionary of metric names and their values.
+                Returns empty dict if no metrics available.
+        """
+        if not hasattr(self.engine.trainer, AnomalibTrainKeys.CALLBACK_METRICS):
+            return {}
+        return {k: v.item() if hasattr(v, "item") else v for k, v in self.engine.trainer.callback_metrics.items()}
+
+    def train_model(self) -> AnomalibTrainDataClass:
+        """Executes model training process.
+
+        Returns:
+            AnomalibTrainDataClass: Contains:
+                - metrics: Collected training metrics
+                - checkpoint_path: Path to best model checkpoint
+
+        Note:
+            - Updates trainer configuration before starting
+            - Uses provided checkpoint if available for resuming
+        """
+        self._update_trainer_args()
+
+        self.engine.train(
+            model=self.model,
+            datamodule=self.data_module,
+            ckpt_path=self.attributes.ckpt_path,
+        )
+
+        metrics = self._get_training_metrics()
+        checkpoint_path = (
+            self.engine.best_model_path if hasattr(self.engine, AnomalibTrainKeys.BEST_MODEL_PATH) else None
+        )
+
+        return AnomalibTrainDataClass(metrics=metrics, checkpoint_path=checkpoint_path)
+
+    def execute(self, container: DataContainer) -> DataContainer:
+        """Executes training and stores results.
+
+        Args:
+            container: Input data container
+
+        Returns:
+            Modified container with training results stored under generic_key
+        """
+        result = self.train_model()
+        self._set_generic_data(container, result)
+        return container
+
+
+def __getattr__(name: str) -> Template:
+    """Creat dynamic templates.
+
+    Only create a template if it's imported, this avoids creating all the base models for all templates
+    and potential import errors due to not available packages.
+    """
+    if name in AnomalibTrain.WrapperEntry.module_att_names:
+        return make_dynamic_template(name, AnomalibTrain)
+    raise AttributeError(f"template `{name}` not found in {__name__}")
+
+
+__all__ = AnomalibTrain.WrapperEntry.module_att_names
+
+if SINAPSIS_BUILD_DOCS:
+    dynamic_templates = [__getattr__(template_name) for template_name in __all__]
+    for template in dynamic_templates:
+        globals()[template.__name__] = template
+        del template