Add core libraries: anomalib, dinov2, open_clip_local

anomalib/__init__.py

@@ -0,0 +1,24 @@
+"""Anomalib library for research and benchmarking."""
+
+# Copyright (C) 2022-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from enum import Enum
+
+__version__ = "1.1.0dev"
+
+
+class LearningType(str, Enum):
+    """Learning type defining how the model learns from the dataset samples."""
+
+    ONE_CLASS = "one_class"
+    ZERO_SHOT = "zero_shot"
+    FEW_SHOT = "few_shot"
+
+
+class TaskType(str, Enum):
+    """Task type used when generating predictions on the dataset."""
+
+    CLASSIFICATION = "classification"
+    DETECTION = "detection"
+    SEGMENTATION = "segmentation"

anomalib/callbacks/__init__.py

@@ -0,0 +1,64 @@
+"""Callbacks for Anomalib models."""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+import logging
+from importlib import import_module
+from pathlib import Path
+
+import yaml
+from jsonargparse import Namespace
+from lightning.pytorch.callbacks import Callback
+from omegaconf import DictConfig, ListConfig, OmegaConf
+
+from .checkpoint import ModelCheckpoint
+from .graph import GraphLogger
+from .model_loader import LoadModelCallback
+from .tiler_configuration import TilerConfigurationCallback
+from .timer import TimerCallback
+
+__all__ = [
+    "ModelCheckpoint",
+    "GraphLogger",
+    "LoadModelCallback",
+    "TilerConfigurationCallback",
+    "TimerCallback",
+]
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_callbacks(config: DictConfig | ListConfig | Namespace) -> list[Callback]:
+    """Return base callbacks for all the lightning models.
+
+    Args:
+        config (DictConfig | ListConfig | Namespace): Model config
+
+    Return:
+        (list[Callback]): List of callbacks.
+    """
+    logger.info("Loading the callbacks")
+
+    callbacks: list[Callback] = []
+
+    if "ckpt_path" in config.trainer and config.ckpt_path is not None:
+        load_model = LoadModelCallback(config.ckpt_path)
+        callbacks.append(load_model)
+
+    if "optimization" in config and "nncf" in config.optimization and config.optimization.nncf.apply:
+        # NNCF wraps torch's jit which conflicts with kornia's jit calls.
+        # Hence, nncf is imported only when required
+        nncf_module = import_module("anomalib.utils.callbacks.nncf.callback")
+        nncf_callback = nncf_module.NNCFCallback
+        nncf_config = yaml.safe_load(OmegaConf.to_yaml(config.optimization.nncf))
+        callbacks.append(
+            nncf_callback(
+                config=nncf_config,
+                export_dir=str(Path(config.project.path) / "compressed"),
+            ),
+        )
+
+    return callbacks

anomalib/callbacks/checkpoint.py

@@ -0,0 +1,58 @@
+"""Anomalib Model Checkpoint Callback."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from lightning.pytorch import Trainer
+from lightning.pytorch.callbacks import ModelCheckpoint as LightningCheckpoint
+from lightning.pytorch.trainer.states import TrainerFn
+
+from anomalib import LearningType
+
+
+class ModelCheckpoint(LightningCheckpoint):
+    """Anomalib Model Checkpoint Callback.
+
+    This class overrides the Lightning ModelCheckpoint callback to enable saving checkpoints without running any
+    training steps. This is useful for zero-/few-shot models, where the fit sequence only consists of validation.
+
+    To enable saving checkpoints without running any training steps, we need to override two checks which are being
+    called in the ``on_validation_end`` method of the parent class:
+    - ``_should_save_on_train_epoch_end``: This method checks whether the checkpoint should be saved at the end of a
+        training epoch, or at the end of the validation sequence. We modify this method to default to saving at the end
+        of the validation sequence when the model is of zero- or few-shot type, unless ``save_on_train_epoch_end`` is
+        specifically set by the user.
+    - ``_should_skip_saving_checkpoint``: This method checks whether the checkpoint should be saved at all. We modify
+        this method to allow saving during both the ``FITTING`` and ``VALIDATING`` states. In addition, we allow saving
+        if the global step has not changed since the last checkpoint, but only for zero- and few-shot models. This is
+        needed because both the last global step and the last checkpoint remain unchanged during zero-/few-shot
+        training, which would otherwise prevent saving checkpoints during validation.
+    """
+
+    def _should_skip_saving_checkpoint(self, trainer: Trainer) -> bool:
+        """Checks whether the checkpoint should be saved.
+
+        Overrides the parent method to allow saving during both the ``FITTING`` and ``VALIDATING`` states, and to allow
+        saving when the global step and last_global_step_saved are both 0 (only for zero-/few-shot models).
+        """
+        is_zero_or_few_shot = trainer.model.learning_type in [LearningType.ZERO_SHOT, LearningType.FEW_SHOT]
+        return (
+            bool(trainer.fast_dev_run)  # disable checkpointing with fast_dev_run
+            or trainer.state.fn not in [TrainerFn.FITTING, TrainerFn.VALIDATING]  # don't save anything during non-fit
+            or trainer.sanity_checking  # don't save anything during sanity check
+            or (self._last_global_step_saved == trainer.global_step and not is_zero_or_few_shot)
+        )
+
+    def _should_save_on_train_epoch_end(self, trainer: Trainer) -> bool:
+        """Checks whether the checkpoint should be saved at the end of a training epoch or validation sequence.
+
+        Overrides the parent method to default to saving at the end of the validation sequence when the model is of
+        zero- or few-shot type, unless ``save_on_train_epoch_end`` is specifically set by the user.
+        """
+        if self._save_on_train_epoch_end is not None:
+            return self._save_on_train_epoch_end
+
+        if trainer.model.learning_type in [LearningType.ZERO_SHOT, LearningType.FEW_SHOT]:
+            return False
+
+        return super()._should_save_on_train_epoch_end(trainer)

anomalib/callbacks/graph.py

@@ -0,0 +1,61 @@
+"""Log model graph to respective logger."""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+from lightning.pytorch import Callback, LightningModule, Trainer
+
+from anomalib.loggers import AnomalibCometLogger, AnomalibTensorBoardLogger, AnomalibWandbLogger
+
+
+class GraphLogger(Callback):
+    """Log model graph to respective logger.
+
+    Examples:
+        Log model graph to Tensorboard
+
+        >>> from anomalib.callbacks import GraphLogger
+        >>> from anomalib.loggers import AnomalibTensorBoardLogger
+        >>> from anomalib.engine import Engine
+        ...
+        >>> logger = AnomalibTensorBoardLogger()
+        >>> callbacks = [GraphLogger()]
+        >>> engine = Engine(logger=logger, callbacks=callbacks)
+
+        Log model graph to Comet
+
+        >>> from anomalib.loggers import AnomalibCometLogger
+        >>> from anomalib.engine import Engine
+        ...
+        >>> logger = AnomalibCometLogger()
+        >>> callbacks = [GraphLogger()]
+        >>> engine = Engine(logger=logger, callbacks=callbacks)
+    """
+
+    def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        """Log model graph to respective logger.
+
+        Args:
+            trainer: Trainer object which contans reference to loggers.
+            pl_module: LightningModule object which is logged.
+        """
+        for logger in trainer.loggers:
+            if isinstance(logger, AnomalibWandbLogger):
+                # NOTE: log graph gets populated only after one backward pass. This won't work for models which do not
+                # require training such as Padim
+                logger.watch(pl_module, log_graph=True, log="all")
+                break
+
+    def on_train_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        """Unwatch model if configured for wandb and log it model graph in Tensorboard if specified.
+
+        Args:
+            trainer: Trainer object which contans reference to loggers.
+            pl_module: LightningModule object which is logged.
+        """
+        for logger in trainer.loggers:
+            if isinstance(logger, AnomalibCometLogger | AnomalibTensorBoardLogger):
+                logger.log_graph(pl_module, input_array=torch.ones((1, 3, 256, 256)))
+            elif isinstance(logger, AnomalibWandbLogger):
+                logger.experiment.unwatch(pl_module)

anomalib/callbacks/metrics.py

@@ -0,0 +1,226 @@
+"""MetricsManager callback."""
+
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+import logging
+from enum import Enum
+from typing import Any
+
+import torch
+from lightning.pytorch import Callback, Trainer
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+
+from anomalib import TaskType
+from anomalib.metrics import create_metric_collection
+from anomalib.models import AnomalyModule
+
+logger = logging.getLogger(__name__)
+
+
+class Device(str, Enum):
+    """Device on which to compute metrics."""
+
+    CPU = "cpu"
+    GPU = "gpu"
+
+
+class _MetricsCallback(Callback):
+    """Create image and pixel-level AnomalibMetricsCollection.
+
+    This callback creates AnomalibMetricsCollection based on the
+        list of strings provided for image and pixel-level metrics.
+    After these MetricCollections are created, the callback assigns
+    these to the lightning module.
+
+    Args:
+        task (TaskType | str): Task type of the current run.
+        image_metrics (list[str] | str | dict[str, dict[str, Any]] | None): List of image-level metrics.
+        pixel_metrics (list[str] | str | dict[str, dict[str, Any]] | None): List of pixel-level metrics.
+        device (str): Whether to compute metrics on cpu or gpu. Defaults to cpu.
+    """
+
+    def __init__(
+        self,
+        task: TaskType | str = TaskType.SEGMENTATION,
+        image_metrics: list[str] | str | dict[str, dict[str, Any]] | None = None,
+        pixel_metrics: list[str] | str | dict[str, dict[str, Any]] | None = None,
+        device: Device = Device.CPU,
+    ) -> None:
+        super().__init__()
+        self.task = TaskType(task)
+        self.image_metric_names = image_metrics
+        self.pixel_metric_names = pixel_metrics
+        self.device = device
+
+    def setup(
+        self,
+        trainer: Trainer,
+        pl_module: AnomalyModule,
+        stage: str | None = None,
+    ) -> None:
+        """Set image and pixel-level AnomalibMetricsCollection within Anomalib Model.
+
+        Args:
+            trainer (pl.Trainer): PyTorch Lightning Trainer
+            pl_module (AnomalyModule): Anomalib Model that inherits pl LightningModule.
+            stage (str | None, optional): fit, validate, test or predict. Defaults to None.
+        """
+        del stage, trainer  # this variable is not used.
+        image_metric_names = [] if self.image_metric_names is None else self.image_metric_names
+        if isinstance(image_metric_names, str):
+            image_metric_names = [image_metric_names]
+
+        pixel_metric_names: list[str] | dict[str, dict[str, Any]]
+        if self.pixel_metric_names is None:
+            pixel_metric_names = []
+        elif self.task == TaskType.CLASSIFICATION:
+            pixel_metric_names = []
+            logger.warning(
+                "Cannot perform pixel-level evaluation when task type is classification. "
+                "Ignoring the following pixel-level metrics: %s",
+                self.pixel_metric_names,
+            )
+        else:
+            pixel_metric_names = (
+                self.pixel_metric_names.copy()
+                if not isinstance(self.pixel_metric_names, str)
+                else [self.pixel_metric_names]
+            )
+
+        # create a separate metric collection for metrics that operate over the semantic segmentation mask
+        # (segmentation mask with a separate channel for each defect type)
+        semantic_pixel_metric_names: list[str] | dict[str, dict[str, Any]] = []
+        # currently only SPRO metric is supported as semantic segmentation metric
+        if "SPRO" in pixel_metric_names:
+            if isinstance(pixel_metric_names, list):
+                pixel_metric_names.remove("SPRO")
+                semantic_pixel_metric_names = ["SPRO"]
+            elif isinstance(pixel_metric_names, dict):
+                spro_metric = pixel_metric_names.pop("SPRO")
+                semantic_pixel_metric_names = {"SPRO": spro_metric}
+            else:
+                logger.warning("Unexpected type for pixel_metric_names: %s", type(pixel_metric_names))
+
+        if isinstance(pl_module, AnomalyModule):
+            pl_module.image_metrics = create_metric_collection(image_metric_names, "image_")
+            if hasattr(pl_module, "pixel_metrics"):  # incase metrics are loaded from model checkpoint
+                new_metrics = create_metric_collection(pixel_metric_names)
+                for name in new_metrics:
+                    if name not in pl_module.pixel_metrics:
+                        pl_module.pixel_metrics.add_metrics(new_metrics[name])
+            else:
+                pl_module.pixel_metrics = create_metric_collection(pixel_metric_names, "pixel_")
+            pl_module.semantic_pixel_metrics = create_metric_collection(semantic_pixel_metric_names, "pixel_")
+            self._set_threshold(pl_module)
+
+    def on_validation_epoch_start(
+        self,
+        trainer: Trainer,
+        pl_module: AnomalyModule,
+    ) -> None:
+        del trainer  # Unused argument.
+
+        pl_module.image_metrics.reset()
+        pl_module.pixel_metrics.reset()
+        pl_module.semantic_pixel_metrics.reset()
+
+    def on_validation_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: AnomalyModule,
+        outputs: STEP_OUTPUT | None,
+        batch: Any,  # noqa: ANN401
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        del trainer, batch, batch_idx, dataloader_idx  # Unused arguments.
+
+        if outputs is not None:
+            self._outputs_to_device(outputs)
+            self._update_metrics(pl_module, outputs)
+
+    def on_validation_epoch_end(
+        self,
+        trainer: Trainer,
+        pl_module: AnomalyModule,
+    ) -> None:
+        del trainer  # Unused argument.
+
+        self._set_threshold(pl_module)
+        self._log_metrics(pl_module)
+
+    def on_test_epoch_start(
+        self,
+        trainer: Trainer,
+        pl_module: AnomalyModule,
+    ) -> None:
+        del trainer  # Unused argument.
+
+        pl_module.image_metrics.reset()
+        pl_module.pixel_metrics.reset()
+        pl_module.semantic_pixel_metrics.reset()
+
+    def on_test_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: AnomalyModule,
+        outputs: STEP_OUTPUT | None,
+        batch: Any,  # noqa: ANN401
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        del trainer, batch, batch_idx, dataloader_idx  # Unused arguments.
+
+        if outputs is not None:
+            self._outputs_to_device(outputs)
+            self._update_metrics(pl_module, outputs)
+
+    def on_test_epoch_end(
+        self,
+        trainer: Trainer,
+        pl_module: AnomalyModule,
+    ) -> None:
+        del trainer  # Unused argument.
+
+        self._log_metrics(pl_module)
+
+    def _set_threshold(self, pl_module: AnomalyModule) -> None:
+        pl_module.image_metrics.set_threshold(pl_module.image_threshold.value.item())
+        pl_module.pixel_metrics.set_threshold(pl_module.pixel_threshold.value.item())
+        pl_module.semantic_pixel_metrics.set_threshold(pl_module.pixel_threshold.value.item())
+
+    def _update_metrics(
+        self,
+        pl_module: AnomalyModule,
+        output: STEP_OUTPUT,
+    ) -> None:
+        pl_module.image_metrics.to(self.device)
+        pl_module.image_metrics.update(output["pred_scores"], output["label"].int())
+        if "mask" in output and "anomaly_maps" in output:
+            pl_module.pixel_metrics.to(self.device)
+            pl_module.pixel_metrics.update(torch.squeeze(output["anomaly_maps"]), torch.squeeze(output["mask"].int()))
+        if "semantic_mask" in output and "anomaly_maps" in output:
+            pl_module.semantic_pixel_metrics.to(self.device)
+            pl_module.semantic_pixel_metrics.update(torch.squeeze(output["anomaly_maps"]), output["semantic_mask"])
+
+    def _outputs_to_device(self, output: STEP_OUTPUT) -> STEP_OUTPUT | dict[str, Any]:
+        if isinstance(output, dict):
+            for key, value in output.items():
+                output[key] = self._outputs_to_device(value)
+        elif isinstance(output, torch.Tensor):
+            output = output.to(self.device)
+        elif isinstance(output, list):
+            for i, value in enumerate(output):
+                output[i] = self._outputs_to_device(value)
+        return output
+
+    @staticmethod
+    def _log_metrics(pl_module: AnomalyModule) -> None:
+        """Log computed performance metrics."""
+        pl_module.log_dict(pl_module.image_metrics, prog_bar=True)
+        if pl_module.pixel_metrics.update_called:
+            pl_module.log_dict(pl_module.pixel_metrics, prog_bar=False)
+        if pl_module.semantic_pixel_metrics.update_called:
+            pl_module.log_dict(pl_module.semantic_pixel_metrics, prog_bar=False)

anomalib/callbacks/model_loader.py

@@ -0,0 +1,39 @@
+"""Callback that loads model weights from the state dict."""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+import logging
+
+import torch
+from lightning.pytorch import Callback, Trainer
+
+from anomalib.models.components import AnomalyModule
+
+logger = logging.getLogger(__name__)
+
+
+class LoadModelCallback(Callback):
+    """Callback that loads the model weights from the state dict.
+
+    Examples:
+        >>> from anomalib.callbacks import LoadModelCallback
+        >>> from anomalib.engine import Engine
+        ...
+        >>> callbacks = [LoadModelCallback(weights_path="path/to/weights.pt")]
+        >>> engine = Engine(callbacks=callbacks)
+    """
+
+    def __init__(self, weights_path: str) -> None:
+        self.weights_path = weights_path
+
+    def setup(self, trainer: Trainer, pl_module: AnomalyModule, stage: str | None = None) -> None:
+        """Call when inference begins.
+
+        Loads the model weights from ``weights_path`` into the PyTorch module.
+        """
+        del trainer, stage  # These variables are not used.
+
+        logger.info("Loading the model from %s", self.weights_path)
+        pl_module.load_state_dict(torch.load(self.weights_path, map_location=pl_module.device)["state_dict"])

anomalib/callbacks/nncf/__init__.py

@@ -0,0 +1,4 @@
+"""Integration NNCF."""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0

anomalib/callbacks/nncf/callback.py

@@ -0,0 +1,106 @@
+"""Callbacks for NNCF optimization."""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+import subprocess  # nosec B404
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import lightning.pytorch as pl
+from lightning.pytorch import Callback
+from nncf import NNCFConfig
+from nncf.torch import register_default_init_args
+
+from anomalib.callbacks.nncf.utils import InitLoader, wrap_nncf_model
+
+if TYPE_CHECKING:
+    from nncf.api.compression import CompressionAlgorithmController
+
+
+class NNCFCallback(Callback):
+    """Callback for NNCF compression.
+
+    Assumes that the pl module contains a 'model' attribute, which is
+    the PyTorch module that must be compressed.
+
+    Args:
+        config (dict): NNCF Configuration
+        export_dir (Str): Path where the export `onnx` and the OpenVINO `xml` and `bin` IR are saved.
+                          If None model will not be exported.
+    """
+
+    def __init__(self, config: dict, export_dir: str | None = None) -> None:
+        self.export_dir = export_dir
+        self.config = NNCFConfig(config)
+        self.nncf_ctrl: CompressionAlgorithmController | None = None
+
+    def setup(self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: str | None = None) -> None:
+        """Call when fit or test begins.
+
+        Takes the pytorch model and wraps it using the compression controller
+        so that it is ready for nncf fine-tuning.
+        """
+        del stage  # `stage` variable is not used.
+
+        if self.nncf_ctrl is not None:
+            return
+
+        # Get validate subset to initialize quantization,
+        # because train subset does not contain anomalous images.
+        init_loader = InitLoader(trainer.datamodule.val_dataloader())
+        config = register_default_init_args(self.config, init_loader)
+
+        self.nncf_ctrl, pl_module.model = wrap_nncf_model(
+            model=pl_module.model,
+            config=config,
+            dataloader=trainer.datamodule.train_dataloader(),
+            init_state_dict=None,  # type: ignore[arg-type]
+        )
+
+    def on_train_batch_start(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        batch: Any,  # noqa: ANN401
+        batch_idx: int,
+        unused: int = 0,
+    ) -> None:
+        """Call when the train batch begins.
+
+        Prepare compression method to continue training the model in the next step.
+        """
+        del trainer, pl_module, batch, batch_idx, unused  # These variables are not used.
+
+        if self.nncf_ctrl:
+            self.nncf_ctrl.scheduler.step()
+
+    def on_train_epoch_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        """Call when the train epoch starts.
+
+        Prepare compression method to continue training the model in the next epoch.
+        """
+        del trainer, pl_module  # `trainer` and `pl_module` variables are not used.
+
+        if self.nncf_ctrl:
+            self.nncf_ctrl.scheduler.epoch_step()
+
+    def on_train_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        """Call when the train ends.
+
+        Exports onnx model and if compression controller is not None, uses the onnx model to generate the OpenVINO IR.
+        """
+        del trainer, pl_module  # `trainer` and `pl_module` variables are not used.
+
+        if self.export_dir is None or self.nncf_ctrl is None:
+            return
+
+        Path(self.export_dir).mkdir(parents=True, exist_ok=True)
+        onnx_path = str(Path(self.export_dir) / "model_nncf.onnx")
+        self.nncf_ctrl.export_model(onnx_path)
+
+        optimize_command = ["mo", "--input_model", onnx_path, "--output_dir", self.export_dir]
+        # TODO(samet-akcay): Check if mo can be done via python API
+        # CVS-122665
+        subprocess.run(optimize_command, check=True)  # noqa: S603  # nosec B603

anomalib/callbacks/nncf/utils.py

@@ -0,0 +1,243 @@
+"""Utils for NNCf optimization."""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+import logging
+from copy import copy
+from typing import TYPE_CHECKING, Any
+
+import torch
+from nncf import NNCFConfig
+from nncf.api.compression import CompressionAlgorithmController
+from nncf.torch import create_compressed_model, load_state, register_default_init_args
+from nncf.torch.initialization import PTInitializingDataLoader
+from nncf.torch.nncf_network import NNCFNetwork
+from torch import nn
+from torch.utils.data.dataloader import DataLoader
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+
+logger = logging.getLogger(name="NNCF compression")
+
+
+class InitLoader(PTInitializingDataLoader):
+    """Initializing data loader for NNCF to be used with unsupervised training algorithms."""
+
+    def __init__(self, data_loader: DataLoader) -> None:
+        super().__init__(data_loader)
+        self._data_loader_iter: Iterator
+
+    def __iter__(self) -> "InitLoader":
+        """Create iterator for dataloader."""
+        self._data_loader_iter = iter(self._data_loader)
+        return self
+
+    def __next__(self) -> torch.Tensor:
+        """Return next item from dataloader iterator."""
+        loaded_item = next(self._data_loader_iter)
+        return loaded_item["image"]
+
+    def get_inputs(self, dataloader_output: dict[str, str | torch.Tensor]) -> tuple[tuple, dict]:
+        """Get input to model.
+
+        Returns:
+            (dataloader_output,), {}: tuple[tuple, dict]: The current model call to be made during
+            the initialization process
+        """
+        return (dataloader_output,), {}
+
+    def get_target(self, _):  # noqa: ANN001, ANN201
+        """Return structure for ground truth in loss criterion based on dataloader output.
+
+        This implementation does not do anything and is a placeholder.
+
+        Returns:
+            None
+        """
+        return
+
+
+def wrap_nncf_model(
+    model: nn.Module,
+    config: dict,
+    dataloader: DataLoader,
+    init_state_dict: dict,
+) -> tuple[CompressionAlgorithmController, NNCFNetwork]:
+    """Wrap model by NNCF.
+
+    :param model: Anomalib model.
+    :param config: NNCF config.
+    :param dataloader: Dataloader for initialization of NNCF model.
+    :param init_state_dict: Opti
+    :return: compression controller, compressed model
+    """
+    nncf_config = NNCFConfig.from_dict(config)
+
+    if not dataloader and not init_state_dict:
+        logger.warning(
+            "Either dataloader or NNCF pre-trained "
+            "model checkpoint should be set. Without this, "
+            "quantizers will not be initialized",
+        )
+
+    compression_state = None
+    resuming_state_dict = None
+    if init_state_dict:
+        resuming_state_dict = init_state_dict.get("model")
+        compression_state = init_state_dict.get("compression_state")
+
+    if dataloader:
+        init_loader = InitLoader(dataloader)
+        nncf_config = register_default_init_args(nncf_config, init_loader)
+
+    nncf_ctrl, nncf_model = create_compressed_model(
+        model=model,
+        config=nncf_config,
+        dump_graphs=False,
+        compression_state=compression_state,
+    )
+
+    if resuming_state_dict:
+        load_state(nncf_model, resuming_state_dict, is_resume=True)
+
+    return nncf_ctrl, nncf_model
+
+
+def is_state_nncf(state: dict) -> bool:
+    """Check if state is the result of NNCF-compressed model."""
+    return bool(state.get("meta", {}).get("nncf_enable_compression", False))
+
+
+def compose_nncf_config(nncf_config: dict, enabled_options: list[str]) -> dict:
+    """Compose NNCf config by selected options.
+
+    :param nncf_config:
+    :param enabled_options:
+    :return: config
+    """
+    optimisation_parts = nncf_config
+    optimisation_parts_to_choose = []
+    if "order_of_parts" in optimisation_parts:
+        # The result of applying the changes from optimisation parts
+        # may depend on the order of applying the changes
+        # (e.g. if for nncf_quantization it is sufficient to have `total_epochs=2`,
+        #  but for sparsity it is required `total_epochs=50`)
+        # So, user can define `order_of_parts` in the optimisation_config
+        # to specify the order of applying the parts.
+        order_of_parts = optimisation_parts["order_of_parts"]
+        if not isinstance(order_of_parts, list):
+            msg = 'The field "order_of_parts" in optimization config should be a list'
+            raise TypeError(msg)
+
+        for part in enabled_options:
+            if part not in order_of_parts:
+                msg = f"The part {part} is selected, but it is absent in order_of_parts={order_of_parts}"
+                raise ValueError(msg)
+
+        optimisation_parts_to_choose = [part for part in order_of_parts if part in enabled_options]
+
+    if "base" not in optimisation_parts:
+        msg = 'Error: the optimisation config does not contain the "base" part'
+        raise KeyError(msg)
+    nncf_config_part = optimisation_parts["base"]
+
+    for part in optimisation_parts_to_choose:
+        if part not in optimisation_parts:
+            msg = f'Error: the optimisation config does not contain the part "{part}"'
+            raise KeyError(msg)
+        optimisation_part_dict = optimisation_parts[part]
+        try:
+            nncf_config_part = merge_dicts_and_lists_b_into_a(nncf_config_part, optimisation_part_dict)
+        except AssertionError as cur_error:
+            err_descr = (
+                f"Error during merging the parts of nncf configs:\n"
+                f"the current part={part}, "
+                f"the order of merging parts into base is {optimisation_parts_to_choose}.\n"
+                f"The error is:\n{cur_error}"
+            )
+            raise RuntimeError(err_descr) from None
+
+    return nncf_config_part
+
+
+def merge_dicts_and_lists_b_into_a(
+    a: dict[Any, Any] | list[Any],
+    b: dict[Any, Any] | list[Any],
+) -> dict[Any, Any] | list[Any]:
+    """Merge dict configs.
+
+    Args:
+        a (dict[Any, Any] | list[Any]): First dict or list.
+        b (dict[Any, Any] | list[Any]): Second dict or list.
+
+    Returns:
+        dict[Any, Any] | list[Any]: Merged dict or list.
+    """
+    return _merge_dicts_and_lists_b_into_a(a, b, "")
+
+
+def _merge_dicts_and_lists_b_into_a(
+    a: dict[Any, Any] | list[Any],
+    b: dict[Any, Any] | list[Any],
+    cur_key: int | str | None = None,
+) -> dict[Any, Any] | list[Any]:
+    """Merge dict configs.
+
+        * works with usual dicts and lists and derived types
+        * supports merging of lists (by concatenating the lists)
+        * makes recursive merging for dict + dict case
+        * overwrites when merging scalar into scalar
+        Note that we merge b into a (whereas Config makes merge a into b),
+        since otherwise the order of list merging is counter-intuitive.
+
+    Args:
+        a (dict[Any, Any] | list[Any]): First dict or list.
+        b (dict[Any, Any] | list[Any]): Second dict or list.
+        cur_key (int | str | None, optional): key for current level of recursion. Defaults to None.
+
+    Returns:
+        dict[Any, Any] | list[Any]: Merged dict or list.
+    """
+
+    def _err_str(_a: dict | list, _b: dict | list, _key: int | str | None = None) -> str:
+        _key_str = "of whole structures" if _key is None else f"during merging for key=`{_key}`"
+        return (
+            f"Error in merging parts of config: different types {_key_str},"
+            f" type(a) = {type(_a)},"
+            f" type(b) = {type(_b)}"
+        )
+
+    if not (isinstance(a, dict | list)):
+        msg = f"Can merge only dicts and lists, whereas type(a)={type(a)}"
+        raise TypeError(msg)
+
+    if not (isinstance(b, dict | list)):
+        raise TypeError(_err_str(a, b, cur_key))
+
+    if (isinstance(a, list) and not isinstance(b, list)) or (isinstance(b, list) and not isinstance(a, list)):
+        raise TypeError(_err_str(a, b, cur_key))
+
+    if isinstance(a, list) and isinstance(b, list):
+        # the main diff w.r.t. mmcf.Config -- merging of lists
+        return a + b
+
+    a = copy(a)
+    for k in b:
+        if k not in a:
+            a[k] = copy(b[k])
+            continue
+        new_cur_key = str(cur_key) + "." + k if cur_key else k
+        if isinstance(a[k], dict | list):
+            a[k] = _merge_dicts_and_lists_b_into_a(a[k], b[k], new_cur_key)
+            continue
+
+        if any(isinstance(b[k], t) for t in [dict, list]):
+            raise TypeError(_err_str(a[k], b[k], new_cur_key))
+
+        # suppose here that a[k] and b[k] are scalars, just overwrite
+        a[k] = b[k]
+    return a

anomalib/callbacks/normalization/__init__.py

@@ -0,0 +1,12 @@
+"""Normalization callbacks.
+
+Note: These callbacks are used within the Engine.
+"""
+
+# Copyright (C) 2023-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .min_max_normalization import _MinMaxNormalizationCallback
+from .utils import get_normalization_callback
+
+__all__ = ["get_normalization_callback", "_MinMaxNormalizationCallback"]

anomalib/callbacks/normalization/base.py

@@ -0,0 +1,29 @@
+"""Base Normalization Callback."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+
+from lightning.pytorch import Callback
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+
+from anomalib.models.components import AnomalyModule
+
+
+class NormalizationCallback(Callback, ABC):
+    """Base normalization callback."""
+
+    @staticmethod
+    @abstractmethod
+    def _normalize_batch(batch: STEP_OUTPUT, pl_module: AnomalyModule) -> None:
+        """Normalize an output batch.
+
+        Args:
+            batch (dict[str, torch.Tensor]): Output batch.
+            pl_module (AnomalyModule): AnomalyModule instance.
+
+        Returns:
+            dict[str, torch.Tensor]: Normalized batch.
+        """
+        raise NotImplementedError

anomalib/callbacks/normalization/min_max_normalization.py

@@ -0,0 +1,109 @@
+"""Anomaly Score Normalization Callback that uses min-max normalization."""
+
+# Copyright (C) 2022-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+from typing import Any
+
+import torch
+from lightning.pytorch import Trainer
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+
+from anomalib.metrics import MinMax
+from anomalib.models.components import AnomalyModule
+from anomalib.utils.normalization.min_max import normalize
+
+from .base import NormalizationCallback
+
+
+class _MinMaxNormalizationCallback(NormalizationCallback):
+    """Callback that normalizes the image-level and pixel-level anomaly scores using min-max normalization.
+
+    Note: This callback is set within the Engine.
+    """
+
+    def setup(self, trainer: Trainer, pl_module: AnomalyModule, stage: str | None = None) -> None:
+        """Add min_max metrics to normalization metrics."""
+        del trainer, stage  # These variables are not used.
+
+        if not hasattr(pl_module, "normalization_metrics"):
+            pl_module.normalization_metrics = MinMax().cpu()
+        elif not isinstance(pl_module.normalization_metrics, MinMax):
+            msg = f"Expected normalization_metrics to be of type MinMax, got {type(pl_module.normalization_metrics)}"
+            raise AttributeError(
+                msg,
+            )
+
+    def on_test_start(self, trainer: Trainer, pl_module: AnomalyModule) -> None:
+        """Call when the test begins."""
+        del trainer  # `trainer` variable is not used.
+
+        for metric in (pl_module.image_metrics, pl_module.pixel_metrics, pl_module.semantic_pixel_metrics):
+            if metric is not None:
+                metric.set_threshold(0.5)
+
+    def on_validation_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: AnomalyModule,
+        outputs: STEP_OUTPUT,
+        batch: Any,  # noqa: ANN401
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        """Call when the validation batch ends, update the min and max observed values."""
+        del trainer, batch, batch_idx, dataloader_idx  # These variables are not used.
+
+        if "anomaly_maps" in outputs:
+            pl_module.normalization_metrics(outputs["anomaly_maps"])
+        elif "box_scores" in outputs:
+            pl_module.normalization_metrics(torch.cat(outputs["box_scores"]))
+        elif "pred_scores" in outputs:
+            pl_module.normalization_metrics(outputs["pred_scores"])
+        else:
+            msg = "No values found for normalization, provide anomaly maps, bbox scores, or image scores"
+            raise ValueError(msg)
+
+    def on_test_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: AnomalyModule,
+        outputs: STEP_OUTPUT | None,
+        batch: Any,  # noqa: ANN401
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        """Call when the test batch ends, normalizes the predicted scores and anomaly maps."""
+        del trainer, batch, batch_idx, dataloader_idx  # These variables are not used.
+
+        self._normalize_batch(outputs, pl_module)
+
+    def on_predict_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: AnomalyModule,
+        outputs: Any,  # noqa: ANN401
+        batch: Any,  # noqa: ANN401
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        """Call when the predict batch ends, normalizes the predicted scores and anomaly maps."""
+        del trainer, batch, batch_idx, dataloader_idx  # These variables are not used.
+
+        self._normalize_batch(outputs, pl_module)
+
+    @staticmethod
+    def _normalize_batch(outputs: Any, pl_module: AnomalyModule) -> None:  # noqa: ANN401
+        """Normalize a batch of predictions."""
+        image_threshold = pl_module.image_threshold.value.cpu()
+        pixel_threshold = pl_module.pixel_threshold.value.cpu()
+        stats = pl_module.normalization_metrics.cpu()
+        if "pred_scores" in outputs:
+            outputs["pred_scores"] = normalize(outputs["pred_scores"], image_threshold, stats.min, stats.max)
+        if "anomaly_maps" in outputs:
+            outputs["anomaly_maps"] = normalize(outputs["anomaly_maps"], pixel_threshold, stats.min, stats.max)
+        if "box_scores" in outputs:
+            outputs["box_scores"] = [
+                normalize(scores, pixel_threshold, stats.min, stats.max) for scores in outputs["box_scores"]
+            ]

anomalib/callbacks/normalization/utils.py

@@ -0,0 +1,78 @@
+"""Normalization callback utils."""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import importlib
+
+from lightning.pytorch import Callback
+from omegaconf import DictConfig
+
+from anomalib.utils.normalization import NormalizationMethod
+from anomalib.utils.types import NORMALIZATION
+
+from .min_max_normalization import _MinMaxNormalizationCallback
+
+
+def get_normalization_callback(
+    normalization_method: NORMALIZATION = NormalizationMethod.MIN_MAX,
+) -> Callback | None:
+    """Return normalization object.
+
+    normalization_method is an instance of ``Callback``, it is returned as is.
+
+    if normalization_method is of type ``NormalizationMethod``, then a new class is created based on the type of
+    normalization_method.
+
+    Otherwise it expects a dictionary containing class_path and init_args.
+        normalization_method:
+            class_path: MinMaxNormalizer
+            init_args:
+                -
+                -
+
+    Example:
+        >>> normalizer = get_normalization_callback(NormalizationMethod.MIN_MAX)
+        or
+        >>> normalizer = get_normalization_callback("min_max")
+        or
+        >>> normalizer = get_normalization_callback({"class_path": "MinMaxNormalizationCallback", "init_args": {}})
+        or
+        >>> normalizer = get_normalization_callback(MinMaxNormalizationCallback())
+    """
+    normalizer: Callback | None
+    if isinstance(normalization_method, NormalizationMethod | str):
+        normalizer = _get_normalizer_from_method(NormalizationMethod(normalization_method))
+    elif isinstance(normalization_method, Callback):
+        normalizer = normalization_method
+    elif isinstance(normalization_method, DictConfig):
+        normalizer = _parse_normalizer_config(normalization_method)
+    else:
+        msg = f"Unknown normalizer type {normalization_method}"
+        raise TypeError(msg)
+    return normalizer
+
+
+def _get_normalizer_from_method(normalization_method: NormalizationMethod | str) -> Callback | None:
+    if normalization_method == NormalizationMethod.NONE:
+        normalizer = None
+    elif normalization_method == NormalizationMethod.MIN_MAX:
+        normalizer = _MinMaxNormalizationCallback()
+    else:
+        msg = f"Unknown normalization method {normalization_method}"
+        raise ValueError(msg)
+    return normalizer
+
+
+def _parse_normalizer_config(normalization_method: DictConfig) -> Callback:
+    class_path = normalization_method.class_path
+    init_args = normalization_method.init_args
+
+    if len(class_path.split(".")) == 1:
+        module_path = "anomalib.utils.callbacks.normalization"
+    else:
+        module_path = ".".join(class_path.split(".")[:-1])
+        class_path = class_path.split(".")[-1]
+    module = importlib.import_module(module_path)
+    class_ = getattr(module, class_path)
+    return class_(**init_args)

anomalib/callbacks/post_processor.py

@@ -0,0 +1,125 @@
+"""Callback that attaches necessary pre/post-processing to the model."""
+
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+from typing import Any
+
+import torch
+from lightning import Callback
+from lightning.pytorch import Trainer
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+
+from anomalib.data.utils import boxes_to_anomaly_maps, boxes_to_masks, masks_to_boxes
+from anomalib.models import AnomalyModule
+
+
+class _PostProcessorCallback(Callback):
+    """Applies post-processing to the model outputs.
+
+    Note: This callback is set within the Engine.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def on_validation_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: AnomalyModule,
+        outputs: STEP_OUTPUT | None,
+        batch: Any,  # noqa: ANN401
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        del batch, batch_idx, dataloader_idx  # Unused arguments.
+
+        if outputs is not None:
+            self.post_process(trainer, pl_module, outputs)
+
+    def on_test_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: AnomalyModule,
+        outputs: STEP_OUTPUT | None,
+        batch: Any,  # noqa: ANN401
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        del batch, batch_idx, dataloader_idx  # Unused arguments.
+
+        if outputs is not None:
+            self.post_process(trainer, pl_module, outputs)
+
+    def on_predict_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: AnomalyModule,
+        outputs: Any,  # noqa: ANN401
+        batch: Any,  # noqa: ANN401
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        del batch, batch_idx, dataloader_idx  # Unused arguments.
+
+        if outputs is not None:
+            self.post_process(trainer, pl_module, outputs)
+
+    def post_process(self, trainer: Trainer, pl_module: AnomalyModule, outputs: STEP_OUTPUT) -> None:
+        if isinstance(outputs, dict):
+            self._post_process(outputs)
+            if trainer.predicting or trainer.testing:
+                self._compute_scores_and_labels(pl_module, outputs)
+
+    @staticmethod
+    def _compute_scores_and_labels(
+        pl_module: AnomalyModule,
+        outputs: dict[str, Any],
+    ) -> None:
+        if "pred_scores" in outputs:
+            outputs["pred_labels"] = outputs["pred_scores"] >= pl_module.image_threshold.value
+        if "anomaly_maps" in outputs:
+            outputs["pred_masks"] = outputs["anomaly_maps"] >= pl_module.pixel_threshold.value
+            if "pred_boxes" not in outputs:
+                outputs["pred_boxes"], outputs["box_scores"] = masks_to_boxes(
+                    outputs["pred_masks"],
+                    outputs["anomaly_maps"],
+                )
+                outputs["box_labels"] = [torch.ones(boxes.shape[0]) for boxes in outputs["pred_boxes"]]
+        # apply thresholding to boxes
+        if "box_scores" in outputs and "box_labels" not in outputs:
+            # apply threshold to assign normal/anomalous label to boxes
+            is_anomalous = [scores > pl_module.pixel_threshold.value for scores in outputs["box_scores"]]
+            outputs["box_labels"] = [labels.int() for labels in is_anomalous]
+
+    @staticmethod
+    def _post_process(outputs: STEP_OUTPUT) -> None:
+        """Compute labels based on model predictions."""
+        if isinstance(outputs, dict):
+            if "pred_scores" not in outputs and "anomaly_maps" in outputs:
+                # infer image scores from anomaly maps
+                outputs["pred_scores"] = (
+                    outputs["anomaly_maps"]  # noqa: PD011
+                    .reshape(outputs["anomaly_maps"].shape[0], -1)
+                    .max(dim=1)
+                    .values
+                )
+            elif "pred_scores" not in outputs and "box_scores" in outputs and "label" in outputs:
+                # infer image score from bbox confidence scores
+                outputs["pred_scores"] = torch.zeros_like(outputs["label"]).float()
+                for idx, (boxes, scores) in enumerate(zip(outputs["pred_boxes"], outputs["box_scores"], strict=True)):
+                    if boxes.numel():
+                        outputs["pred_scores"][idx] = scores.max().item()
+
+            if "pred_boxes" in outputs and "anomaly_maps" not in outputs:
+                # create anomaly maps from bbox predictions for thresholding and evaluation
+                image_size: tuple[int, int] = outputs["image"].shape[-2:]
+                pred_boxes: torch.Tensor = outputs["pred_boxes"]
+                box_scores: torch.Tensor = outputs["box_scores"]
+
+                outputs["anomaly_maps"] = boxes_to_anomaly_maps(pred_boxes, box_scores, image_size)
+
+                if "boxes" in outputs:
+                    true_boxes: list[torch.Tensor] = outputs["boxes"]
+                    outputs["mask"] = boxes_to_masks(true_boxes, image_size)

anomalib/callbacks/thresholding.py

@@ -0,0 +1,197 @@
+"""Thresholding callback."""
+
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import importlib
+from typing import Any
+
+import torch
+from lightning.pytorch import Callback, Trainer
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from omegaconf import DictConfig, ListConfig
+
+from anomalib.metrics.threshold import BaseThreshold
+from anomalib.models import AnomalyModule
+from anomalib.utils.types import THRESHOLD
+
+
+class _ThresholdCallback(Callback):
+    """Setup/apply thresholding.
+
+    Note: This callback is set within the Engine.
+    """
+
+    def __init__(
+        self,
+        threshold: THRESHOLD = "F1AdaptiveThreshold",
+    ) -> None:
+        super().__init__()
+        self._initialize_thresholds(threshold)
+        self.image_threshold: BaseThreshold
+        self.pixel_threshold: BaseThreshold
+
+    def setup(self, trainer: Trainer, pl_module: AnomalyModule, stage: str) -> None:
+        del trainer, stage  # Unused arguments.
+        if not hasattr(pl_module, "image_threshold"):
+            pl_module.image_threshold = self.image_threshold
+        if not hasattr(pl_module, "pixel_threshold"):
+            pl_module.pixel_threshold = self.pixel_threshold
+
+    def on_validation_epoch_start(self, trainer: Trainer, pl_module: AnomalyModule) -> None:
+        del trainer  # Unused argument.
+        self._reset(pl_module)
+
+    def on_validation_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: AnomalyModule,
+        outputs: STEP_OUTPUT | None,
+        batch: Any,  # noqa: ANN401
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        del trainer, batch, batch_idx, dataloader_idx  # Unused arguments.
+        if outputs is not None:
+            self._outputs_to_cpu(outputs)
+            self._update(pl_module, outputs)
+
+    def on_validation_epoch_end(self, trainer: Trainer, pl_module: AnomalyModule) -> None:
+        del trainer  # Unused argument.
+        self._compute(pl_module)
+
+    def _initialize_thresholds(
+        self,
+        threshold: THRESHOLD,
+    ) -> None:
+        """Initialize ``self.image_threshold`` and ``self.pixel_threshold``.
+
+        Args:
+            threshold (THRESHOLD):
+                Threshold configuration
+
+        Example:
+            >>> _initialize_thresholds(F1AdaptiveThreshold())
+            or
+            >>> _initialize_thresholds((ManualThreshold(0.5), ManualThreshold(0.5)))
+            or configuration
+
+        For more details on configuration see :fun:`_load_from_config`
+
+        Raises:
+            ValueError: Unknown threshold class or incorrect configuration
+        """
+        # TODO(djdameln): Add tests for each case
+        # CVS-122661
+        # When only a single threshold class is passed.
+        # This initializes image and pixel thresholds with the same class
+        # >>> _initialize_thresholds(F1AdaptiveThreshold())
+        if isinstance(threshold, BaseThreshold):
+            self.image_threshold = threshold
+            self.pixel_threshold = threshold.clone()
+
+        # When a tuple of threshold classes are passed
+        # >>> _initialize_thresholds((ManualThreshold(0.5), ManualThreshold(0.5)))
+        elif isinstance(threshold, tuple) and isinstance(threshold[0], BaseThreshold):
+            self.image_threshold = threshold[0]
+            self.pixel_threshold = threshold[1]
+        # When the passed threshold is not an instance of a Threshold class.
+        elif isinstance(threshold, str | DictConfig | ListConfig | list):
+            self._load_from_config(threshold)
+        else:
+            msg = f"Invalid threshold type {type(threshold)}"
+            raise TypeError(msg)
+
+    def _load_from_config(self, threshold: DictConfig | str | ListConfig | list[dict[str, str | float]]) -> None:
+        """Load the thresholding class based on the config.
+
+        Example:
+            threshold: F1AdaptiveThreshold
+            or
+            threshold:
+                class_path: F1AdaptiveThreshold
+                init_args:
+                    -
+            or
+            threshold:
+                - F1AdaptiveThreshold
+                - F1AdaptiveThreshold
+            or
+            threshold:
+                - class_path: F1AdaptiveThreshold
+                    init_args:
+                        -
+                - class_path: F1AdaptiveThreshold
+        """
+        if isinstance(threshold, str | DictConfig):
+            self.image_threshold = self._get_threshold_from_config(threshold)
+            self.pixel_threshold = self.image_threshold.clone()
+        elif isinstance(threshold, ListConfig | list):
+            self.image_threshold = self._get_threshold_from_config(threshold[0])
+            self.pixel_threshold = self._get_threshold_from_config(threshold[1])
+        else:
+            msg = f"Invalid threshold config {threshold}"
+            raise TypeError(msg)
+
+    def _get_threshold_from_config(self, threshold: DictConfig | str | dict[str, str | float]) -> BaseThreshold:
+        """Return the instantiated threshold object.
+
+        Example:
+            >>> _get_threshold_from_config(F1AdaptiveThreshold)
+            or
+            >>> config = DictConfig({
+            ...    "class_path": "ManualThreshold",
+            ...    "init_args": {"default_value": 0.7}
+            ... })
+            >>> __get_threshold_from_config(config)
+            or
+            >>> config = DictConfig({
+            ...    "class_path": "anomalib.metrics.threshold.F1AdaptiveThreshold"
+            ... })
+            >>> __get_threshold_from_config(config)
+
+        Returns:
+            (BaseThreshold): Instance of threshold object.
+        """
+        if isinstance(threshold, str):
+            threshold = DictConfig({"class_path": threshold})
+
+        class_path = threshold["class_path"]
+        init_args = threshold.get("init_args", {})
+
+        if len(class_path.split(".")) == 1:
+            module_path = "anomalib.metrics.threshold"
+
+        else:
+            module_path = ".".join(class_path.split(".")[:-1])
+            class_path = class_path.split(".")[-1]
+
+        module = importlib.import_module(module_path)
+        class_ = getattr(module, class_path)
+        return class_(**init_args)
+
+    def _reset(self, pl_module: AnomalyModule) -> None:
+        pl_module.image_threshold.reset()
+        pl_module.pixel_threshold.reset()
+
+    def _outputs_to_cpu(self, output: STEP_OUTPUT) -> STEP_OUTPUT | dict[str, Any]:
+        if isinstance(output, dict):
+            for key, value in output.items():
+                output[key] = self._outputs_to_cpu(value)
+        elif isinstance(output, torch.Tensor):
+            output = output.cpu()
+        return output
+
+    def _update(self, pl_module: AnomalyModule, outputs: STEP_OUTPUT) -> None:
+        pl_module.image_threshold.cpu()
+        pl_module.image_threshold.update(outputs["pred_scores"], outputs["label"].int())
+        if "mask" in outputs and "anomaly_maps" in outputs:
+            pl_module.pixel_threshold.cpu()
+            pl_module.pixel_threshold.update(outputs["anomaly_maps"], outputs["mask"].int())
+
+    def _compute(self, pl_module: AnomalyModule) -> None:
+        pl_module.image_threshold.compute()
+        if pl_module.pixel_threshold._update_called:  # noqa: SLF001
+            pl_module.pixel_threshold.compute()
+        else:
+            pl_module.pixel_threshold.value = pl_module.image_threshold.value

anomalib/callbacks/tiler_configuration.py

@@ -0,0 +1,74 @@
+"""Tiler Callback."""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+from collections.abc import Sequence
+
+import lightning.pytorch as pl
+from lightning.pytorch.callbacks import Callback
+
+from anomalib.data.utils.tiler import ImageUpscaleMode, Tiler
+from anomalib.models.components import AnomalyModule
+
+__all__ = ["TilerConfigurationCallback"]
+
+
+class TilerConfigurationCallback(Callback):
+    """Tiler Configuration Callback."""
+
+    def __init__(
+        self,
+        enable: bool = False,
+        tile_size: int | Sequence = 256,
+        stride: int | Sequence | None = None,
+        remove_border_count: int = 0,
+        mode: ImageUpscaleMode = ImageUpscaleMode.PADDING,
+    ) -> None:
+        """Set tiling configuration from the command line.
+
+        Args:
+            enable (bool): Boolean to enable tiling operation.
+                Defaults to False.
+            tile_size ([int | Sequence]): Tile size.
+                Defaults to 256.
+            stride ([int | Sequence]): Stride to move tiles on the image.
+            remove_border_count (int, optional): Number of pixels to remove from the image before
+                tiling. Defaults to 0.
+            mode (str, optional): Up-scaling mode when untiling overlapping tiles.
+                Defaults to "padding".
+            tile_count (SupportsIndex, optional): Number of random tiles to sample from the image.
+                Defaults to 4.
+        """
+        self.enable = enable
+        self.tile_size = tile_size
+        self.stride = stride
+        self.remove_border_count = remove_border_count
+        self.mode = mode
+
+    def setup(self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: str | None = None) -> None:
+        """Set Tiler object within Anomalib Model.
+
+        Args:
+            trainer (pl.Trainer): PyTorch Lightning Trainer
+            pl_module (pl.LightningModule): Anomalib Model that inherits pl LightningModule.
+            stage (str | None, optional): fit, validate, test or predict. Defaults to None.
+
+        Raises:
+            ValueError: When Anomalib Model doesn't contain ``Tiler`` object, it means the model
+                doesn not support tiling operation.
+        """
+        del trainer, stage  # These variables are not used.
+
+        if self.enable:
+            if isinstance(pl_module, AnomalyModule) and hasattr(pl_module.model, "tiler"):
+                pl_module.model.tiler = Tiler(
+                    tile_size=self.tile_size,
+                    stride=self.stride,
+                    remove_border_count=self.remove_border_count,
+                    mode=self.mode,
+                )
+            else:
+                msg = "Model does not support tiling."
+                raise ValueError(msg)

anomalib/callbacks/timer.py

@@ -0,0 +1,109 @@
+"""Callback to measure training and testing time of a PyTorch Lightning module."""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import time
+
+import torch
+from lightning.pytorch import Callback, LightningModule, Trainer
+
+logger = logging.getLogger(__name__)
+
+
+class TimerCallback(Callback):
+    """Callback that measures the training and testing time of a PyTorch Lightning module.
+
+    Examples:
+        >>> from anomalib.callbacks import TimerCallback
+        >>> from anomalib.engine import Engine
+        ...
+        >>> callbacks = [TimerCallback()]
+        >>> engine = Engine(callbacks=callbacks)
+    """
+
+    def __init__(self) -> None:
+        self.start: float
+        self.num_images: int = 0
+
+    def on_fit_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        """Call when fit begins.
+
+        Sets the start time to the time training started.
+
+        Args:
+            trainer (Trainer): PyTorch Lightning trainer.
+            pl_module (LightningModule): Current training module.
+
+        Returns:
+            None
+        """
+        del trainer, pl_module  # These variables are not used.
+
+        self.start = time.time()
+
+    def on_fit_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        """Call when fit ends.
+
+        Prints the time taken for training.
+
+        Args:
+            trainer (Trainer): PyTorch Lightning trainer.
+            pl_module (LightningModule): Current training module.
+
+        Returns:
+            None
+        """
+        del trainer, pl_module  # Unused arguments.
+        logger.info("Training took %5.2f seconds", (time.time() - self.start))
+
+    def on_test_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        """Call when the test begins.
+
+        Sets the start time to the time testing started.
+        Goes over all the test dataloaders and adds the number of images in each.
+
+        Args:
+            trainer (Trainer): PyTorch Lightning trainer.
+            pl_module (LightningModule): Current training module.
+
+        Returns:
+            None
+        """
+        del pl_module  # Unused argument.
+
+        self.start = time.time()
+        self.num_images = 0
+
+        if trainer.test_dataloaders is not None:  # Check to placate Mypy.
+            if isinstance(trainer.test_dataloaders, torch.utils.data.dataloader.DataLoader):
+                self.num_images += len(trainer.test_dataloaders.dataset)
+            else:
+                for dataloader in trainer.test_dataloaders:
+                    self.num_images += len(dataloader.dataset)
+
+    def on_test_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        """Call when the test ends.
+
+        Prints the time taken for testing and the throughput in frames per second.
+
+        Args:
+            trainer (Trainer): PyTorch Lightning trainer.
+            pl_module (LightningModule): Current training module.
+
+        Returns:
+            None
+        """
+        del pl_module  # Unused argument.
+
+        testing_time = time.time() - self.start
+        output = f"Testing took {testing_time} seconds\nThroughput "
+        if trainer.test_dataloaders is not None:
+            if isinstance(trainer.test_dataloaders, torch.utils.data.dataloader.DataLoader):
+                test_data_loader = trainer.test_dataloaders
+            else:
+                test_data_loader = trainer.test_dataloaders[0]
+            output += f"(batch_size={test_data_loader.batch_size})"
+        output += f" : {self.num_images/testing_time} FPS"
+        logger.info(output)

anomalib/callbacks/visualizer.py

@@ -0,0 +1,182 @@
+"""Visualizer Callback.
+
+This is assigned by Anomalib Engine internally.
+"""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from pathlib import Path
+from typing import Any, cast
+
+from lightning.pytorch import Callback, Trainer
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+
+from anomalib.data.utils.image import save_image, show_image
+from anomalib.loggers import AnomalibWandbLogger
+from anomalib.loggers.base import ImageLoggerBase
+from anomalib.models import AnomalyModule
+from anomalib.utils.visualization import (
+    BaseVisualizer,
+    GeneratorResult,
+    VisualizationStep,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class _VisualizationCallback(Callback):
+    """Callback for visualization that is used internally by the Engine.
+
+    Args:
+        visualizers (BaseVisualizer | list[BaseVisualizer]):
+            Visualizer objects that are used for computing the visualizations. Defaults to None.
+        save (bool, optional): Save the image. Defaults to False.
+        root (Path | None, optional): The path to save the images. Defaults to None.
+        log (bool, optional): Log the images into the loggers. Defaults to False.
+        show (bool, optional): Show the images. Defaults to False.
+
+    Example:
+        >>> visualizers = [ImageVisualizer(), MetricsVisualizer()]
+        >>> visualization_callback = _VisualizationCallback(
+        ... visualizers=visualizers,
+        ...   save=True,
+        ...   root="results/images"
+        ... )
+
+        CLI
+        $ anomalib train --model Padim --data MVTec \
+            --visualization.visualizers ImageVisualizer \
+            --visualization.visualizers+=MetricsVisualizer
+        or
+        $ anomalib train --model Padim --data MVTec \
+            --visualization.visualizers '[ImageVisualizer, MetricsVisualizer]'
+
+    Raises:
+        ValueError: Incase `root` is None and `save` is True.
+    """
+
+    def __init__(
+        self,
+        visualizers: BaseVisualizer | list[BaseVisualizer],
+        save: bool = False,
+        root: Path | None = None,
+        log: bool = False,
+        show: bool = False,
+    ) -> None:
+        self.save = save
+        if save and root is None:
+            msg = "`root` must be provided if save is True"
+            raise ValueError(msg)
+        self.root: Path = root if root is not None else Path()  # need this check for mypy
+        self.log = log
+        self.show = show
+        self.generators = visualizers if isinstance(visualizers, list) else [visualizers]
+
+    def on_test_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: AnomalyModule,
+        outputs: STEP_OUTPUT | None,
+        batch: Any,  # noqa: ANN401
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        for generator in self.generators:
+            if generator.visualize_on == VisualizationStep.BATCH:
+                for result in generator(
+                    trainer=trainer,
+                    pl_module=pl_module,
+                    outputs=outputs,
+                    batch=batch,
+                    batch_idx=batch_idx,
+                    dataloader_idx=dataloader_idx,
+                ):
+                    if self.save:
+                        if result.file_name is None:
+                            msg = "``save`` is set to ``True`` but file name is ``None``"
+                            raise ValueError(msg)
+
+                        # Get the filename to save the image.
+                        # Filename is split based on the datamodule name and category.
+                        # For example, if the filename is `MVTec/bottle/000.png`, then the
+                        # filename is split based on `MVTec/bottle` and `000.png` is saved.
+                        if trainer.datamodule is not None:
+                            filename = str(result.file_name).split(
+                                sep=f"{trainer.datamodule.name}/{trainer.datamodule.category}",
+                            )[-1]
+                        else:
+                            filename = Path(result.file_name).name
+                        save_image(image=result.image, root=self.root, filename=filename)
+                    if self.show:
+                        show_image(image=result.image, title=str(result.file_name))
+                    if self.log:
+                        self._add_to_logger(result, pl_module, trainer)
+
+    def on_test_end(self, trainer: Trainer, pl_module: AnomalyModule) -> None:
+        for generator in self.generators:
+            if generator.visualize_on == VisualizationStep.STAGE_END:
+                for result in generator(trainer=trainer, pl_module=pl_module):
+                    if self.save:
+                        if result.file_name is None:
+                            msg = "``save`` is set to ``True`` but file name is ``None``"
+                            raise ValueError(msg)
+                        save_image(image=result.image, root=self.root, filename=result.file_name)
+                    if self.show:
+                        show_image(image=result.image, title=str(result.file_name))
+                    if self.log:
+                        self._add_to_logger(result, pl_module, trainer)
+
+        for logger in trainer.loggers:
+            if isinstance(logger, AnomalibWandbLogger):
+                logger.save()
+
+    def on_predict_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: AnomalyModule,
+        outputs: STEP_OUTPUT | None,
+        batch: Any,  # noqa: ANN401
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        return self.on_test_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx)
+
+    def on_predict_end(self, trainer: Trainer, pl_module: AnomalyModule) -> None:
+        return self.on_test_end(trainer, pl_module)
+
+    def _add_to_logger(
+        self,
+        result: GeneratorResult,
+        module: AnomalyModule,
+        trainer: Trainer,
+    ) -> None:
+        """Add image to logger.
+
+        Args:
+            result (GeneratorResult): Output from the generators.
+            module (AnomalyModule): LightningModule from which the global step is extracted.
+            trainer (Trainer): Trainer object.
+        """
+        # Store names of logger and the logger in a dict
+        available_loggers = {
+            type(logger).__name__.lower().replace("logger", "").replace("anomalib", ""): logger
+            for logger in trainer.loggers
+        }
+        # save image to respective logger
+        if result.file_name is None:
+            msg = "File name is None"
+            raise ValueError(msg)
+        filename = result.file_name
+        image = result.image
+        for log_to in available_loggers:
+            # check if logger object is same as the requested object
+            if isinstance(available_loggers[log_to], ImageLoggerBase):
+                logger: ImageLoggerBase = cast(ImageLoggerBase, available_loggers[log_to])  # placate mypy
+                _name = filename.parent.name + "_" + filename.name if isinstance(filename, Path) else filename
+                logger.add_image(
+                    image=image,
+                    name=_name,
+                    global_step=module.global_step,
+                )

anomalib/cli/__init__.py

@@ -0,0 +1,8 @@
+"""Anomalib CLI."""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .cli import AnomalibCLI
+
+__all__ = ["AnomalibCLI"]

anomalib/cli/cli.py

@@ -0,0 +1,483 @@
+"""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()

anomalib/cli/install.py

@@ -0,0 +1,81 @@
+"""Anomalib install subcommand code."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+from pkg_resources import Requirement
+from rich.console import Console
+from rich.logging import RichHandler
+
+from anomalib.cli.utils.installation import (
+    get_requirements,
+    get_torch_install_args,
+    parse_requirements,
+)
+
+logger = logging.getLogger("pip")
+logger.setLevel(logging.WARNING)  # setLevel: CRITICAL, ERROR, WARNING, INFO, DEBUG, NOTSET
+console = Console()
+handler = RichHandler(
+    console=console,
+    show_level=False,
+    show_path=False,
+)
+logger.addHandler(handler)
+
+
+def anomalib_install(option: str = "full", verbose: bool = False) -> int:
+    """Install Anomalib requirements.
+
+    Args:
+        option (str | None): Optional-dependency to install requirements for.
+        verbose (bool): Set pip logger level to INFO
+
+    Raises:
+        ValueError: When the task is not supported.
+
+    Returns:
+        int: Status code of the pip install command.
+    """
+    from pip._internal.commands import create_command
+
+    requirements_dict = get_requirements("anomalib")
+
+    requirements = []
+    if option == "full":
+        for extra in requirements_dict:
+            requirements.extend(requirements_dict[extra])
+    elif option in requirements_dict:
+        requirements.extend(requirements_dict[option])
+    elif option is not None:
+        requirements.append(Requirement.parse(option))
+
+    # Parse requirements into torch and other requirements.
+    # This is done to parse the correct version of torch (cpu/cuda).
+    torch_requirement, other_requirements = parse_requirements(requirements, skip_torch=option not in ("full", "core"))
+
+    # Get install args for torch to install it from a specific index-url
+    install_args: list[str] = []
+    torch_install_args = []
+    if option in ("full", "core") and torch_requirement is not None:
+        torch_install_args = get_torch_install_args(torch_requirement)
+
+    # Combine torch and other requirements.
+    install_args = other_requirements + torch_install_args
+
+    # Install requirements.
+    with console.status("[bold green]Installing packages...  This may take a few minutes.\n") as status:
+        if verbose:
+            logger.setLevel(logging.INFO)
+            status.stop()
+        console.log(f"Installation list: [yellow]{install_args}[/yellow]")
+        status_code = create_command("install").main(install_args)
+        if status_code == 0:
+            console.log(f"Installation Complete: {install_args}")
+
+    if status_code == 0:
+        console.print("Anomalib Installation [bold green]Complete.[/bold green]")
+
+    return status_code

anomalib/cli/utils/__init__.py

@@ -0,0 +1,8 @@
+"""Anomalib CLI Utils."""
+
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .help_formatter import CustomHelpFormatter
+
+__all__ = ["CustomHelpFormatter"]

anomalib/cli/utils/help_formatter.py

@@ -0,0 +1,268 @@
+"""Custom Help Formatters for Anomalib CLI."""
+
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import re
+import sys
+from typing import TypeVar
+
+import docstring_parser
+from jsonargparse import DefaultHelpFormatter
+from rich.markdown import Markdown
+from rich.panel import Panel
+from rich_argparse import RichHelpFormatter
+
+REQUIRED_ARGUMENTS = {
+    "train": {"model", "model.help", "data", "data.help", "ckpt_path", "config"},
+    "fit": {"model", "model.help", "data", "data.help", "ckpt_path", "config"},
+    "validate": {"model", "model.help", "data", "data.help", "ckpt_path", "config"},
+    "test": {"model", "model.help", "data", "data.help", "ckpt_path", "config"},
+    "predict": {"model", "model.help", "data", "data.help", "ckpt_path", "config"},
+    "export": {"model", "model.help", "export_type", "ckpt_path", "config"},
+}
+
+try:
+    from anomalib.engine import Engine
+
+    DOCSTRING_USAGE = {
+        "train": Engine.train,
+        "fit": Engine.fit,
+        "validate": Engine.validate,
+        "test": Engine.test,
+        "predict": Engine.predict,
+        "export": Engine.export,
+    }
+except ImportError:
+    print("To use other subcommand using `anomalib install`")
+
+
+def get_short_docstring(component: TypeVar) -> str:
+    """Get the short description from the docstring.
+
+    Args:
+        component (TypeVar): The component to get the docstring from
+
+    Returns:
+        str: The short description
+    """
+    if component.__doc__ is None:
+        return ""
+    docstring = docstring_parser.parse(component.__doc__)
+    return docstring.short_description
+
+
+def get_verbosity_subcommand() -> dict:
+    """Parse command line arguments and returns a dictionary of key-value pairs.
+
+    Returns:
+        A dictionary containing the parsed command line arguments.
+
+    Examples:
+        >>> import sys
+        >>> sys.argv = ['anomalib', 'train', '-h', '-v']
+        >>> get_verbosity_subcommand()
+        {'subcommand': 'train', 'help': True, 'verbosity': 1}
+    """
+    arguments: dict = {"subcommand": None, "help": False, "verbosity": 2}
+    if len(sys.argv) >= 2 and sys.argv[1] not in ("--help", "-h"):
+        arguments["subcommand"] = sys.argv[1]
+    if "--help" in sys.argv or "-h" in sys.argv:
+        arguments["help"] = True
+        if arguments["subcommand"] in REQUIRED_ARGUMENTS:
+            arguments["verbosity"] = 0
+            if "-v" in sys.argv or "--verbose" in sys.argv:
+                arguments["verbosity"] = 1
+            if "-vv" in sys.argv:
+                arguments["verbosity"] = 2
+    return arguments
+
+
+def get_intro() -> Markdown:
+    """Return a Markdown object containing the introduction text for Anomalib CLI Guide.
+
+    The introduction text includes a brief description of the guide and links to the Github repository and documentation
+
+    Returns:
+        A Markdown object containing the introduction text for Anomalib CLI Guide.
+    """
+    intro_markdown = (
+        "# Anomalib CLI Guide\n\n"
+        "Github Repository: [https://github.com/openvinotoolkit/anomalib](https://github.com/openvinotoolkit/anomalib)."
+        "\n\n"
+        "A better guide is provided by the [documentation](https://anomalib.readthedocs.io/en/latest/index.html)."
+    )
+    return Markdown(intro_markdown)
+
+
+def get_verbose_usage(subcommand: str = "train") -> str:
+    """Return a string containing verbose usage information for the specified subcommand.
+
+    Args:
+    ----
+        subcommand (str): The name of the subcommand to get verbose usage information for. Defaults to "train".
+
+    Returns:
+    -------
+        str: A string containing verbose usage information for the specified subcommand.
+    """
+    return (
+        "To get more overridable argument information, run the command below.\n"
+        "```python\n"
+        "# Verbosity Level 1\n"
+        f"anomalib {subcommand} [optional_arguments] -h -v\n"
+        "# Verbosity Level 2\n"
+        f"anomalib {subcommand} [optional_arguments] -h -vv\n"
+        "```"
+    )
+
+
+def get_cli_usage_docstring(component: object | None) -> str | None:
+    r"""Get the cli usage from the docstring.
+
+    Args:
+    ----
+        component (Optional[object]): The component to get the docstring from
+
+    Returns:
+    -------
+        Optional[str]: The quick-start guide as Markdown format.
+
+    Example:
+    -------
+        component.__doc__ = '''
+            <Prev Section>
+
+            CLI Usage:
+                1. First Step.
+                2. Second Step.
+
+            <Next Section>
+        '''
+        >>> get_cli_usage_docstring(component)
+        "1. First Step.\n2. Second Step."
+    """
+    if component is None or component.__doc__ is None or "CLI Usage" not in component.__doc__:
+        return None
+
+    pattern = r"CLI Usage:(.*?)(?=\n{2,}|\Z)"
+    match = re.search(pattern, component.__doc__, re.DOTALL)
+
+    if match:
+        contents = match.group(1).strip().split("\n")
+        return "\n".join([content.strip() for content in contents])
+    return None
+
+
+def render_guide(subcommand: str | None = None) -> list:
+    """Render a guide for the specified subcommand.
+
+    Args:
+    ----
+        subcommand (Optional[str]): The subcommand to render the guide for.
+
+    Returns:
+    -------
+        list: A list of contents to be displayed in the guide.
+    """
+    if subcommand is None or subcommand not in DOCSTRING_USAGE:
+        return []
+    contents = [get_intro()]
+    target_command = DOCSTRING_USAGE[subcommand]
+    cli_usage = get_cli_usage_docstring(target_command)
+    if cli_usage is not None:
+        cli_usage += f"\n{get_verbose_usage(subcommand)}"
+        quick_start = Panel(Markdown(cli_usage), border_style="dim", title="Quick-Start", title_align="left")
+        contents.append(quick_start)
+    return contents
+
+
+class CustomHelpFormatter(RichHelpFormatter, DefaultHelpFormatter):
+    """A custom help formatter for Anomalib CLI.
+
+    This formatter extends the RichHelpFormatter and DefaultHelpFormatter classes to provide
+    a more detailed and customizable help output for Anomalib CLI.
+
+    Attributes:
+    verbosity_level : int
+        The level of verbosity for the help output.
+    subcommand : str | None
+        The subcommand to render the guide for.
+
+    Methods:
+    add_usage(usage, actions, *args, **kwargs)
+        Add usage information to the help output.
+    add_argument(action)
+        Add an argument to the help output.
+    format_help()
+        Format the help output.
+    """
+
+    verbosity_dict = get_verbosity_subcommand()
+    verbosity_level = verbosity_dict["verbosity"]
+    subcommand = verbosity_dict["subcommand"]
+
+    def add_usage(self, usage: str | None, actions: list, *args, **kwargs) -> None:
+        """Add usage information to the formatter.
+
+        Args:
+        ----
+            usage (str | None): A string describing the usage of the program.
+            actions (list): An list of argparse.Action objects.
+            *args (Any): Additional positional arguments to pass to the superclass method.
+            **kwargs (Any): Additional keyword arguments to pass to the superclass method.
+
+        Returns:
+        -------
+            None
+        """
+        if self.subcommand in REQUIRED_ARGUMENTS:
+            if self.verbosity_level == 0:
+                actions = []
+            elif self.verbosity_level == 1:
+                actions = [action for action in actions if action.dest in REQUIRED_ARGUMENTS[self.subcommand]]
+
+        super().add_usage(usage, actions, *args, **kwargs)
+
+    def add_argument(self, action: argparse.Action) -> None:
+        """Add an argument to the help formatter.
+
+        If the verbose level is set to 0, the argument is not added.
+        If the verbose level is set to 1 and the argument is not in the non-skip list, the argument is not added.
+
+        Args:
+        ----
+            action (argparse.Action): The action to add to the help formatter.
+        """
+        if self.subcommand in REQUIRED_ARGUMENTS:
+            if self.verbosity_level == 0:
+                return
+            if self.verbosity_level == 1 and action.dest not in REQUIRED_ARGUMENTS[self.subcommand]:
+                return
+        super().add_argument(action)
+
+    def format_help(self) -> str:
+        """Format the help message for the current command and returns it as a string.
+
+        The help message includes information about the command's arguments and options,
+        as well as any additional information provided by the command's help guide.
+
+        Returns:
+            str: A string containing the formatted help message.
+        """
+        with self.console.capture() as capture:
+            section = self._root_section
+            if self.subcommand in REQUIRED_ARGUMENTS and self.verbosity_level in (0, 1) and len(section.rich_items) > 1:
+                contents = render_guide(self.subcommand)
+                for content in contents:
+                    self.console.print(content)
+            if self.verbosity_level > 0:
+                if len(section.rich_items) > 1:
+                    section = Panel(section, border_style="dim", title="Arguments", title_align="left")
+                self.console.print(section, highlight=False, soft_wrap=True)
+        help_msg = capture.get()
+
+        if help_msg:
+            help_msg = self._long_break_matcher.sub("\n\n", help_msg).rstrip() + "\n"
+        return help_msg

anomalib/cli/utils/installation.py

@@ -0,0 +1,430 @@
+"""Anomalib installation util functions."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import json
+import os
+import platform
+import re
+from importlib.metadata import requires
+from pathlib import Path
+from warnings import warn
+
+from pkg_resources import Requirement
+
+AVAILABLE_TORCH_VERSIONS = {
+    "2.0.0": {"torchvision": "0.15.1", "cuda": ("11.7", "11.8")},
+    "2.0.1": {"torchvision": "0.15.2", "cuda": ("11.7", "11.8")},
+    "2.1.1": {"torchvision": "0.16.1", "cuda": ("11.8", "12.1")},
+    "2.1.2": {"torchvision": "0.16.2", "cuda": ("11.8", "12.1")},
+    "2.2.0": {"torchvision": "0.16.2", "cuda": ("11.8", "12.1")},
+}
+
+
+def get_requirements(module: str = "anomalib") -> dict[str, list[Requirement]]:
+    """Get requirements of module from importlib.metadata.
+
+    This function returns list of required packages from importlib_metadata.
+
+    Example:
+        >>> get_requirements("anomalib")
+        {
+            "base": ["jsonargparse==4.27.1", ...],
+            "core": ["torch==2.1.1", ...],
+            ...
+        }
+
+    Returns:
+        dict[str, list[Requirement]]: List of required packages for each optional-extras.
+    """
+    requirement_list: list[str] | None = requires(module)
+    extra_requirement: dict[str, list[Requirement]] = {}
+    if requirement_list is None:
+        return extra_requirement
+    for requirement in requirement_list:
+        extra = "core"
+        requirement_extra: list[str] = requirement.replace(" ", "").split(";")
+        if isinstance(requirement_extra, list) and len(requirement_extra) > 1:
+            extra = requirement_extra[-1].split("==")[-1].strip("'\"")
+        _requirement_name = requirement_extra[0]
+        _requirement = Requirement.parse(_requirement_name)
+        if extra in extra_requirement:
+            extra_requirement[extra].append(_requirement)
+        else:
+            extra_requirement[extra] = [_requirement]
+    return extra_requirement
+
+
+def parse_requirements(
+    requirements: list[Requirement],
+    skip_torch: bool = False,
+) -> tuple[str | None, list[str]]:
+    """Parse requirements and returns torch and other requirements.
+
+    Args:
+        requirements (list[Requirement]): List of requirements.
+        skip_torch (bool): Whether to skip torch requirement. Defaults to False.
+
+    Raises:
+        ValueError: If torch requirement is not found.
+
+    Examples:
+        >>> requirements = [
+        ...     Requirement.parse("torch==1.13.0"),
+        ...     Requirement.parse("onnx>=1.8.1"),
+        ... ]
+        >>> parse_requirements(requirements=requirements)
+        (Requirement.parse("torch==1.13.0"),
+        Requirement.parse("onnx>=1.8.1"))
+
+    Returns:
+        tuple[str, list[str], list[str]]: Tuple of torch and other requirements.
+    """
+    torch_requirement: str | None = None
+    other_requirements: list[str] = []
+
+    for requirement in requirements:
+        if requirement.unsafe_name == "torch":
+            torch_requirement = str(requirement)
+            if len(requirement.specs) > 1:
+                warn(
+                    "requirements.txt contains. Please remove other versions of torch from requirements.",
+                    stacklevel=2,
+                )
+
+        # Rest of the requirements are task requirements.
+        # Other torch-related requirements such as `torchvision` are to be excluded.
+        # This is because torch-related requirements are already handled in torch_requirement.
+        else:
+            # if not requirement.unsafe_name.startswith("torch"):
+            other_requirements.append(str(requirement))
+
+    if not skip_torch and not torch_requirement:
+        msg = "Could not find torch requirement. Anoamlib depends on torch. Please add torch to your requirements."
+        raise ValueError(msg)
+
+    # Get the unique list of the requirements.
+    other_requirements = list(set(other_requirements))
+
+    return torch_requirement, other_requirements
+
+
+def get_cuda_version() -> str | None:
+    """Get CUDA version installed on the system.
+
+    Examples:
+        >>> # Assume that CUDA version is 11.2
+        >>> get_cuda_version()
+        "11.2"
+
+        >>> # Assume that CUDA is not installed on the system
+        >>> get_cuda_version()
+        None
+
+    Returns:
+        str | None: CUDA version installed on the system.
+    """
+    # 1. Check CUDA_HOME Environment variable
+    cuda_home = os.environ.get("CUDA_HOME", "/usr/local/cuda")
+
+    if Path(cuda_home).exists():
+        # Check $CUDA_HOME/version.json file.
+        version_file = Path(cuda_home) / "version.json"
+        if version_file.is_file():
+            with Path(version_file).open() as file:
+                data = json.load(file)
+                cuda_version = data.get("cuda", {}).get("version", None)
+                if cuda_version is not None:
+                    cuda_version_parts = cuda_version.split(".")
+                    return ".".join(cuda_version_parts[:2])
+    # 2. 'nvcc --version' check & without version.json case
+    try:
+        result = os.popen(cmd="nvcc --version")
+        output = result.read()
+
+        cuda_version_pattern = r"cuda_(\d+\.\d+)"
+        cuda_version_match = re.search(cuda_version_pattern, output)
+
+        if cuda_version_match is not None:
+            return cuda_version_match.group(1)
+    except OSError:
+        msg = "Could not find cuda-version. Instead, the CPU version of torch will be installed."
+        warn(msg, stacklevel=2)
+    return None
+
+
+def update_cuda_version_with_available_torch_cuda_build(cuda_version: str, torch_version: str) -> str:
+    """Update the installed CUDA version with the highest supported CUDA version by PyTorch.
+
+    Args:
+        cuda_version (str): The installed CUDA version.
+        torch_version (str): The PyTorch version.
+
+    Raises:
+        Warning: If the installed CUDA version is not supported by PyTorch.
+
+    Examples:
+        >>> update_cuda_version_with_available_torch_cuda_builds("11.1", "1.13.0")
+        "11.6"
+
+        >>> update_cuda_version_with_available_torch_cuda_builds("11.7", "1.13.0")
+        "11.7"
+
+        >>> update_cuda_version_with_available_torch_cuda_builds("11.8", "1.13.0")
+        "11.7"
+
+        >>> update_cuda_version_with_available_torch_cuda_builds("12.1", "2.0.1")
+        "11.8"
+
+    Returns:
+        str: The updated CUDA version.
+    """
+    max_supported_cuda = max(AVAILABLE_TORCH_VERSIONS[torch_version]["cuda"])
+    min_supported_cuda = min(AVAILABLE_TORCH_VERSIONS[torch_version]["cuda"])
+    bounded_cuda_version = max(min(cuda_version, max_supported_cuda), min_supported_cuda)
+
+    if cuda_version != bounded_cuda_version:
+        warn(
+            f"Installed CUDA version is v{cuda_version}. \n"
+            f"v{min_supported_cuda} <= Supported CUDA version <= v{max_supported_cuda}.\n"
+            f"This script will use CUDA v{bounded_cuda_version}.\n"
+            f"However, this may not be safe, and you are advised to install the correct version of CUDA.\n"
+            f"For more details, refer to https://pytorch.org/get-started/locally/",
+            stacklevel=2,
+        )
+        cuda_version = bounded_cuda_version
+
+    return cuda_version
+
+
+def get_cuda_suffix(cuda_version: str) -> str:
+    """Get CUDA suffix for PyTorch versions.
+
+    Args:
+        cuda_version (str): CUDA version installed on the system.
+
+    Note:
+        The CUDA version of PyTorch is not always the same as the CUDA version
+            that is installed on the system. For example, the latest PyTorch
+            version (1.10.0) supports CUDA 11.3, but the latest CUDA version
+            that is available for download is 11.2. Therefore, we need to use
+            the latest available CUDA version for PyTorch instead of the CUDA
+            version that is installed on the system. Therefore, this function
+            shoudl be regularly updated to reflect the latest available CUDA.
+
+    Examples:
+        >>> get_cuda_suffix(cuda_version="11.2")
+        "cu112"
+
+        >>> get_cuda_suffix(cuda_version="11.8")
+        "cu118"
+
+    Returns:
+        str: CUDA suffix for PyTorch or mmX version.
+    """
+    return f"cu{cuda_version.replace('.', '')}"
+
+
+def get_hardware_suffix(with_available_torch_build: bool = False, torch_version: str | None = None) -> str:
+    """Get hardware suffix for PyTorch or mmX versions.
+
+    Args:
+        with_available_torch_build (bool): Whether to use the latest available
+            PyTorch build or not. If True, the latest available PyTorch build
+            will be used. If False, the installed PyTorch build will be used.
+            Defaults to False.
+        torch_version (str | None): PyTorch version. This is only used when the
+            ``with_available_torch_build`` is True.
+
+    Examples:
+        >>> # Assume that CUDA version is 11.2
+        >>> get_hardware_suffix()
+        "cu112"
+
+        >>> # Assume that CUDA is not installed on the system
+        >>> get_hardware_suffix()
+        "cpu"
+
+        Assume that that installed CUDA version is 12.1.
+        However, the latest available CUDA version for PyTorch v2.0 is 11.8.
+        Therefore, we use 11.8 instead of 12.1. This is because PyTorch does not
+        support CUDA 12.1 yet. In this case, we could correct the CUDA version
+        by setting `with_available_torch_build` to True.
+
+        >>> cuda_version = get_cuda_version()
+        "12.1"
+        >>> get_hardware_suffix(with_available_torch_build=True, torch_version="2.0.1")
+        "cu118"
+
+    Returns:
+        str: Hardware suffix for PyTorch or mmX version.
+    """
+    cuda_version = get_cuda_version()
+    if cuda_version:
+        if with_available_torch_build:
+            if torch_version is None:
+                msg = "``torch_version`` must be provided when with_available_torch_build is True."
+                raise ValueError(msg)
+            cuda_version = update_cuda_version_with_available_torch_cuda_build(cuda_version, torch_version)
+        hardware_suffix = get_cuda_suffix(cuda_version)
+    else:
+        hardware_suffix = "cpu"
+
+    return hardware_suffix
+
+
+def add_hardware_suffix_to_torch(
+    requirement: Requirement,
+    hardware_suffix: str | None = None,
+    with_available_torch_build: bool = False,
+) -> str:
+    """Add hardware suffix to the torch requirement.
+
+    Args:
+        requirement (Requirement): Requirement object comprising requirement
+            details.
+        hardware_suffix (str | None): Hardware suffix. If None, it will be set
+            to the correct hardware suffix. Defaults to None.
+        with_available_torch_build (bool): To check whether the installed
+            CUDA version is supported by the latest available PyTorch build.
+            Defaults to False.
+
+    Examples:
+        >>> from pkg_resources import Requirement
+        >>> req = "torch>=1.13.0, <=2.0.1"
+        >>> requirement = Requirement.parse(req)
+        >>> requirement.name, requirement.specs
+        ('torch', [('>=', '1.13.0'), ('<=', '2.0.1')])
+
+        >>> add_hardware_suffix_to_torch(requirement)
+        'torch>=1.13.0+cu121, <=2.0.1+cu121'
+
+        ``with_available_torch_build=True`` will use the latest available PyTorch build.
+        >>> req = "torch==2.0.1"
+        >>> requirement = Requirement.parse(req)
+        >>> add_hardware_suffix_to_torch(requirement, with_available_torch_build=True)
+        'torch==2.0.1+cu118'
+
+        It is possible to pass the ``hardware_suffix`` manually.
+        >>> req = "torch==2.0.1"
+        >>> requirement = Requirement.parse(req)
+        >>> add_hardware_suffix_to_torch(requirement, hardware_suffix="cu121")
+        'torch==2.0.1+cu111'
+
+    Raises:
+        ValueError: When the requirement has more than two version criterion.
+
+    Returns:
+        str: Updated torch package with the right cuda suffix.
+    """
+    name = requirement.unsafe_name
+    updated_specs: list[str] = []
+
+    for operator, version in requirement.specs:
+        hardware_suffix = hardware_suffix or get_hardware_suffix(with_available_torch_build, version)
+        updated_version = version + f"+{hardware_suffix}" if not version.startswith(("2.1", "2.2")) else version
+
+        # ``specs`` contains operators and versions as follows:
+        # These are to be concatenated again for the updated version.
+        updated_specs.append(operator + updated_version)
+
+    updated_requirement: str = ""
+
+    if updated_specs:
+        # This is the case when specs are e.g. ['<=1.9.1+cu111']
+        if len(updated_specs) == 1:
+            updated_requirement = name + updated_specs[0]
+        # This is the case when specs are e.g., ['<=1.9.1+cu111', '>=1.8.1+cu111']
+        elif len(updated_specs) == 2:
+            updated_requirement = name + updated_specs[0] + ", " + updated_specs[1]
+        else:
+            msg = (
+                "Requirement version can be a single value or a range. \n"
+                "For example it could be torch>=1.8.1 "
+                "or torch>=1.8.1, <=1.9.1\n"
+                f"Got {updated_specs} instead."
+            )
+            raise ValueError(msg)
+    return updated_requirement
+
+
+def get_torch_install_args(requirement: str | Requirement) -> list[str]:
+    """Get the install arguments for Torch requirement.
+
+    This function will return the install arguments for the Torch requirement
+    and its corresponding torchvision requirement.
+
+    Args:
+        requirement (str | Requirement): The torch requirement.
+
+    Raises:
+        RuntimeError: If the OS is not supported.
+
+    Example:
+        >>> from pkg_resources import Requirement
+        >>> requriment = "torch>=1.13.0"
+        >>> get_torch_install_args(requirement)
+        ['--extra-index-url', 'https://download.pytorch.org/whl/cpu',
+        'torch==1.13.0+cpu', 'torchvision==0.14.0+cpu']
+
+    Returns:
+        list[str]: The install arguments.
+    """
+    if isinstance(requirement, str):
+        requirement = Requirement.parse(requirement)
+
+    # NOTE: This does not take into account if the requirement has multiple versions
+    #   such as torch<2.0.1,>=1.13.0
+    if len(requirement.specs) < 1:
+        return [str(requirement)]
+    select_spec_idx = 0
+    for i, spec in enumerate(requirement.specs):
+        if "=" in spec[0]:
+            select_spec_idx = i
+            break
+    operator, version = requirement.specs[select_spec_idx]
+    if version not in AVAILABLE_TORCH_VERSIONS:
+        version = max(AVAILABLE_TORCH_VERSIONS.keys())
+        warn(
+            f"Torch Version will be selected as {version}.",
+            stacklevel=2,
+        )
+    install_args: list[str] = []
+
+    if platform.system() in ("Linux", "Windows"):
+        # Get the hardware suffix (eg., +cpu, +cu116 and +cu118 etc.)
+        hardware_suffix = get_hardware_suffix(with_available_torch_build=True, torch_version=version)
+
+        # Create the PyTorch Index URL to download the correct wheel.
+        index_url = f"https://download.pytorch.org/whl/{hardware_suffix}"
+
+        # Create the PyTorch version depending on the CUDA version. For example,
+        # If CUDA version is 11.2, then the PyTorch version is 1.8.0+cu112.
+        # If CUDA version is None, then the PyTorch version is 1.8.0+cpu.
+        torch_version = add_hardware_suffix_to_torch(requirement, hardware_suffix, with_available_torch_build=True)
+
+        # Get the torchvision version depending on the torch version.
+        torchvision_version = AVAILABLE_TORCH_VERSIONS[version]["torchvision"]
+        torchvision_requirement = f"torchvision{operator}{torchvision_version}"
+        if isinstance(torchvision_version, str) and not torchvision_version.startswith("0.16"):
+            torchvision_requirement += f"+{hardware_suffix}"
+
+        # Return the install arguments.
+        install_args += [
+            "--extra-index-url",
+            # "--index-url",
+            index_url,
+            torch_version,
+            torchvision_requirement,
+        ]
+    elif platform.system() in ("macos", "Darwin"):
+        torch_version = str(requirement)
+        install_args += [torch_version]
+    else:
+        msg = f"Unsupported OS: {platform.system()}"
+        raise RuntimeError(msg)
+
+    return install_args

anomalib/cli/utils/openvino.py

@@ -0,0 +1,32 @@
+"""Utils for OpenVINO parser."""
+
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+from jsonargparse import ArgumentParser
+
+from anomalib.utils.exceptions import try_import
+
+logger = logging.getLogger(__name__)
+
+
+if try_import("openvino"):
+    from openvino.tools.ovc.cli_parser import get_common_cli_parser
+else:
+    get_common_cli_parser = None
+
+
+def add_openvino_export_arguments(parser: ArgumentParser) -> None:
+    """Add OpenVINO arguments to parser under --mo key."""
+    if get_common_cli_parser is not None:
+        group = parser.add_argument_group("OpenVINO Model Optimizer arguments (optional)")
+        ov_parser = get_common_cli_parser()
+        # remove redundant keys from mo keys
+        for arg in ov_parser._actions:  # noqa: SLF001
+            if arg.dest in ("help", "input_model", "output_dir"):
+                continue
+            group.add_argument(f"--ov_args.{arg.dest}", type=arg.type, default=arg.default, help=arg.help)
+    else:
+        logger.info("OpenVINO is possibly not installed in the environment. Skipping adding it to parser.")

anomalib/data/__init__.py

@@ -0,0 +1,72 @@
+"""Anomalib Datasets."""
+
+# Copyright (C) 2022-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+import importlib
+import logging
+from enum import Enum
+from itertools import chain
+
+from omegaconf import DictConfig, ListConfig
+
+from anomalib.utils.config import to_tuple
+
+from .base import AnomalibDataModule, AnomalibDataset
+from .depth import DepthDataFormat, Folder3D, MVTec3D
+from .image import BTech, Folder, ImageDataFormat, Kolektor, MVTec, MVTecLoco, Visa
+from .predict import PredictDataset
+from .utils import LabelName
+from .video import Avenue, ShanghaiTech, UCSDped, VideoDataFormat
+
+logger = logging.getLogger(__name__)
+
+
+DataFormat = Enum(  # type: ignore[misc]
+    "DataFormat",
+    {i.name: i.value for i in chain(DepthDataFormat, ImageDataFormat, VideoDataFormat)},
+)
+
+
+def get_datamodule(config: DictConfig | ListConfig) -> AnomalibDataModule:
+    """Get Anomaly Datamodule.
+
+    Args:
+        config (DictConfig | ListConfig): Configuration of the anomaly model.
+
+    Returns:
+        PyTorch Lightning DataModule
+    """
+    logger.info("Loading the datamodule")
+
+    module = importlib.import_module(".".join(config.data.class_path.split(".")[:-1]))
+    dataclass = getattr(module, config.data.class_path.split(".")[-1])
+    init_args = {**config.data.get("init_args", {})}  # get dict
+    if "image_size" in init_args:
+        init_args["image_size"] = to_tuple(init_args["image_size"])
+
+    return dataclass(**init_args)
+
+
+__all__ = [
+    "AnomalibDataset",
+    "AnomalibDataModule",
+    "DepthDataFormat",
+    "ImageDataFormat",
+    "VideoDataFormat",
+    "get_datamodule",
+    "BTech",
+    "Folder",
+    "Folder3D",
+    "PredictDataset",
+    "Kolektor",
+    "MVTec",
+    "MVTec3D",
+    "MVTecLoco",
+    "Avenue",
+    "UCSDped",
+    "ShanghaiTech",
+    "Visa",
+    "LabelName",
+]

anomalib/data/base/__init__.py

@@ -0,0 +1,18 @@
+"""Base classes for custom dataset and datamodules."""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+from .datamodule import AnomalibDataModule
+from .dataset import AnomalibDataset
+from .depth import AnomalibDepthDataset
+from .video import AnomalibVideoDataModule, AnomalibVideoDataset
+
+__all__ = [
+    "AnomalibDataset",
+    "AnomalibDataModule",
+    "AnomalibVideoDataset",
+    "AnomalibVideoDataModule",
+    "AnomalibDepthDataset",
+]

anomalib/data/base/datamodule.py

@@ -0,0 +1,305 @@
+"""Anomalib datamodule base class."""
+
+# Copyright (C) 2022-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+import logging
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+from lightning.pytorch import LightningDataModule
+from lightning.pytorch.trainer.states import TrainerFn
+from lightning.pytorch.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
+from torch.utils.data.dataloader import DataLoader, default_collate
+from torchvision.transforms.v2 import Resize, Transform
+
+from anomalib.data.utils import TestSplitMode, ValSplitMode, random_split, split_by_label
+from anomalib.data.utils.synthetic import SyntheticAnomalyDataset
+
+if TYPE_CHECKING:
+    from pandas import DataFrame
+
+    from anomalib.data.base.dataset import AnomalibDataset
+
+logger = logging.getLogger(__name__)
+
+
+def collate_fn(batch: list) -> dict[str, Any]:
+    """Collate bounding boxes as lists.
+
+    Bounding boxes and `masks` (not `mask`) are collated as a list of tensors. If `masks` exists,
+    the `mask_path` is also collated as a list since each element in the batch could be unequal.
+    For all other entries, the default collate function is used.
+
+    Args:
+        batch (List): list of items in the batch where len(batch) is equal to the batch size.
+
+    Returns:
+        dict[str, Any]: Dictionary containing the collated batch information.
+    """
+    elem = batch[0]  # sample an element from the batch to check the type.
+    out_dict = {}
+    if isinstance(elem, dict):
+        if "boxes" in elem:
+            # collate boxes as list
+            out_dict["boxes"] = [item.pop("boxes") for item in batch]
+        if "semantic_mask" in elem:
+            # semantic masks have a variable number of channels, so we collate them as a list
+            out_dict["semantic_mask"] = [item.pop("semantic_mask") for item in batch]
+        if "mask_path" in elem and isinstance(elem["mask_path"], list):
+            # collate mask paths as list
+            out_dict["mask_path"] = [item.pop("mask_path") for item in batch]
+        # collate other data normally
+        out_dict.update({key: default_collate([item[key] for item in batch]) for key in elem})
+        return out_dict
+    return default_collate(batch)
+
+
+class AnomalibDataModule(LightningDataModule, ABC):
+    """Base Anomalib data module.
+
+    Args:
+        train_batch_size (int): Batch size used by the train dataloader.
+        eval_batch_size (int): Batch size used by the val and test dataloaders.
+        num_workers (int): Number of workers used by the train, val and test dataloaders.
+        val_split_mode (ValSplitMode): Determines how the validation split is obtained.
+            Options: [none, same_as_test, from_test, synthetic]
+        val_split_ratio (float): Fraction of the train or test images held our for validation.
+        test_split_mode (Optional[TestSplitMode], optional): Determines how the test split is obtained.
+            Options: [none, from_dir, synthetic].
+            Defaults to ``None``.
+        test_split_ratio (float): Fraction of the train images held out for testing.
+            Defaults to ``None``.
+        image_size (tuple[int, int], optional): Size to which input images should be resized.
+            Defaults to ``None``.
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        train_transform (Transform, optional): Transforms that should be applied to the input images during training.
+            Defaults to ``None``.
+        eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation.
+            Defaults to ``None``.
+        seed (int | None, optional): Seed used during random subset splitting.
+            Defaults to ``None``.
+    """
+
+    def __init__(
+        self,
+        train_batch_size: int,
+        eval_batch_size: int,
+        num_workers: int,
+        val_split_mode: ValSplitMode | str,
+        val_split_ratio: float,
+        test_split_mode: TestSplitMode | str | None = None,
+        test_split_ratio: float | None = None,
+        image_size: tuple[int, int] | None = None,
+        transform: Transform | None = None,
+        train_transform: Transform | None = None,
+        eval_transform: Transform | None = None,
+        seed: int | None = None,
+    ) -> None:
+        super().__init__()
+        self.train_batch_size = train_batch_size
+        self.eval_batch_size = eval_batch_size
+        self.num_workers = num_workers
+        self.test_split_mode = TestSplitMode(test_split_mode) if test_split_mode else TestSplitMode.NONE
+        self.test_split_ratio = test_split_ratio
+        self.val_split_mode = ValSplitMode(val_split_mode)
+        self.val_split_ratio = val_split_ratio
+        self.image_size = image_size
+        self.seed = seed
+
+        # set transforms
+        if bool(train_transform) != bool(eval_transform):
+            msg = "Only one of train_transform and eval_transform was specified. This is not recommended because \
+                    it could lead to unexpected behaviour. Please ensure training and eval transforms have the same \
+                    reshape and normalization characteristics."
+            logger.warning(msg)
+        self._train_transform = train_transform or transform
+        self._eval_transform = eval_transform or transform
+
+        self.train_data: AnomalibDataset
+        self.val_data: AnomalibDataset
+        self.test_data: AnomalibDataset
+
+        self._samples: DataFrame | None = None
+        self._category: str = ""
+
+        self._is_setup = False  # flag to track if setup has been called from the trainer
+
+    @property
+    def name(self) -> str:
+        """Name of the datamodule."""
+        return self.__class__.__name__
+
+    def setup(self, stage: str | None = None) -> None:
+        """Set up train, validation and test data.
+
+        Args:
+            stage: str | None:  Train/Val/Test stages.
+                Defaults to ``None``.
+        """
+        has_subset = any(hasattr(self, subset) for subset in ["train_data", "val_data", "test_data"])
+        if not has_subset or not self._is_setup:
+            self._setup(stage)
+            self._create_test_split()
+            self._create_val_split()
+            if isinstance(stage, TrainerFn):
+                # only set the flag if the stage is a TrainerFn, which means the setup has been called from a trainer
+                self._is_setup = True
+
+    @abstractmethod
+    def _setup(self, _stage: str | None = None) -> None:
+        """Set up the datasets and perform dynamic subset splitting.
+
+        This method may be overridden in subclass for custom splitting behaviour.
+
+        Note:
+            The stage argument is not used here. This is because, for a given instance of an AnomalibDataModule
+            subclass, all three subsets are created at the first call of setup(). This is to accommodate the subset
+            splitting behaviour of anomaly tasks, where the validation set is usually extracted from the test set, and
+            the test set must therefore be created as early as the `fit` stage.
+
+        """
+        raise NotImplementedError
+
+    @property
+    def category(self) -> str:
+        """Get the category of the datamodule."""
+        return self._category
+
+    @category.setter
+    def category(self, category: str) -> None:
+        """Set the category of the datamodule."""
+        self._category = category
+
+    def _create_test_split(self) -> None:
+        """Obtain the test set based on the settings in the config."""
+        if self.test_data.has_normal:
+            # split the test data into normal and anomalous so these can be processed separately
+            normal_test_data, self.test_data = split_by_label(self.test_data)
+        elif self.test_split_mode != TestSplitMode.NONE:
+            # when the user did not provide any normal images for testing, we sample some from the training set,
+            # except when the user explicitly requested no test splitting.
+            logger.info(
+                "No normal test images found. Sampling from training set using a split ratio of %0.2f",
+                self.test_split_ratio,
+            )
+            if self.test_split_ratio is not None:
+                self.train_data, normal_test_data = random_split(self.train_data, self.test_split_ratio, seed=self.seed)
+
+        if self.test_split_mode == TestSplitMode.FROM_DIR:
+            self.test_data += normal_test_data
+        elif self.test_split_mode == TestSplitMode.SYNTHETIC:
+            self.test_data = SyntheticAnomalyDataset.from_dataset(normal_test_data)
+        elif self.test_split_mode != TestSplitMode.NONE:
+            msg = f"Unsupported Test Split Mode: {self.test_split_mode}"
+            raise ValueError(msg)
+
+    def _create_val_split(self) -> None:
+        """Obtain the validation set based on the settings in the config."""
+        if self.val_split_mode == ValSplitMode.FROM_TRAIN:
+            # randomly sampled from train set
+            self.train_data, self.val_data = random_split(
+                self.train_data,
+                self.val_split_ratio,
+                label_aware=True,
+                seed=self.seed,
+            )
+        elif self.val_split_mode == ValSplitMode.FROM_TEST:
+            # randomly sampled from test set
+            self.test_data, self.val_data = random_split(
+                self.test_data,
+                self.val_split_ratio,
+                label_aware=True,
+                seed=self.seed,
+            )
+        elif self.val_split_mode == ValSplitMode.SAME_AS_TEST:
+            # equal to test set
+            self.val_data = self.test_data
+        elif self.val_split_mode == ValSplitMode.SYNTHETIC:
+            # converted from random training sample
+            self.train_data, normal_val_data = random_split(self.train_data, self.val_split_ratio, seed=self.seed)
+            self.val_data = SyntheticAnomalyDataset.from_dataset(normal_val_data)
+        elif self.val_split_mode == ValSplitMode.FROM_DIR:
+            # the val_data is prepared in subclass
+            assert hasattr(
+                self,
+                "val_data",
+            ), f"FROM_DIR is not supported for {self.__class__.__name__} which does not assign val_data in _setup."
+        elif self.val_split_mode != ValSplitMode.NONE:
+            msg = f"Unknown validation split mode: {self.val_split_mode}"
+            raise ValueError(msg)
+
+    def train_dataloader(self) -> TRAIN_DATALOADERS:
+        """Get train dataloader."""
+        return DataLoader(
+            dataset=self.train_data,
+            shuffle=True,
+            batch_size=self.train_batch_size,
+            num_workers=self.num_workers,
+        )
+
+    def val_dataloader(self) -> EVAL_DATALOADERS:
+        """Get validation dataloader."""
+        return DataLoader(
+            dataset=self.val_data,
+            shuffle=False,
+            batch_size=self.eval_batch_size,
+            num_workers=self.num_workers,
+            collate_fn=collate_fn,
+        )
+
+    def test_dataloader(self) -> EVAL_DATALOADERS:
+        """Get test dataloader."""
+        return DataLoader(
+            dataset=self.test_data,
+            shuffle=False,
+            batch_size=self.eval_batch_size,
+            num_workers=self.num_workers,
+            collate_fn=collate_fn,
+        )
+
+    def predict_dataloader(self) -> EVAL_DATALOADERS:
+        """Use the test dataloader for inference unless overridden."""
+        return self.test_dataloader()
+
+    @property
+    def transform(self) -> Transform:
+        """Property that returns the user-specified transform for the datamodule, if any.
+
+        This property is accessed by the engine to set the transform for the model. The eval_transform takes precedence
+        over the train_transform, because the transform that we store in the model is the one that should be used during
+        inference.
+        """
+        if self._eval_transform:
+            return self._eval_transform
+        return None
+
+    @property
+    def train_transform(self) -> Transform:
+        """Get the transforms that will be passed to the train dataset.
+
+        If the train_transform is not set, the engine will request the transform from the model.
+        """
+        if self._train_transform:
+            return self._train_transform
+        if getattr(self, "trainer", None) and self.trainer.model and self.trainer.model.transform:
+            return self.trainer.model.transform
+        if self.image_size:
+            return Resize(self.image_size, antialias=True)
+        return None
+
+    @property
+    def eval_transform(self) -> Transform:
+        """Get the transform that will be passed to the val/test/predict datasets.
+
+        If the eval_transform is not set, the engine will request the transform from the model.
+        """
+        if self._eval_transform:
+            return self._eval_transform
+        if getattr(self, "trainer", None) and self.trainer.model and self.trainer.model.transform:
+            return self.trainer.model.transform
+        if self.image_size:
+            return Resize(self.image_size, antialias=True)
+        return None

anomalib/data/base/dataset.py

@@ -0,0 +1,208 @@
+"""Anomalib dataset base class."""
+
+# Copyright (C) 2022-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import copy
+import logging
+from abc import ABC
+from collections.abc import Sequence
+from pathlib import Path
+
+import pandas as pd
+import torch
+from pandas import DataFrame
+from torch.utils.data import Dataset
+from torchvision.transforms.v2 import Transform
+from torchvision.tv_tensors import Mask
+
+from anomalib import TaskType
+from anomalib.data.utils import LabelName, masks_to_boxes, read_image, read_mask
+
+_EXPECTED_COLUMNS_CLASSIFICATION = ["image_path", "split"]
+_EXPECTED_COLUMNS_SEGMENTATION = [*_EXPECTED_COLUMNS_CLASSIFICATION, "mask_path"]
+_EXPECTED_COLUMNS_PERTASK = {
+    "classification": _EXPECTED_COLUMNS_CLASSIFICATION,
+    "segmentation": _EXPECTED_COLUMNS_SEGMENTATION,
+    "detection": _EXPECTED_COLUMNS_SEGMENTATION,
+}
+
+logger = logging.getLogger(__name__)
+
+
+class AnomalibDataset(Dataset, ABC):
+    """Anomalib dataset.
+
+    The dataset is based on a dataframe that contains the information needed by the dataloader to load each of
+    the dataset items into memory.
+
+    The samples dataframe must be set from the subclass using the setter of the `samples` property.
+
+    The DataFrame must, at least, include the following columns:
+        - `split` (str): The subset to which the dataset item is assigned (e.g., 'train', 'test').
+        - `image_path` (str): Path to the file system location where the image is stored.
+        - `label_index` (int): Index of the anomaly label, typically 0 for 'normal' and 1 for 'anomalous'.
+        - `mask_path` (str, optional): Path to the ground truth masks (for the anomalous images only).
+        Required if task is 'segmentation'.
+
+    Example DataFrame:
+        +---+-------------------+-----------+-------------+------------------+-------+
+        |   | image_path        | label     | label_index | mask_path        | split |
+        +---+-------------------+-----------+-------------+------------------+-------+
+        | 0 | path/to/image.png | anomalous | 1           | path/to/mask.png | train |
+        +---+-------------------+-----------+-------------+------------------+-------+
+
+    Note:
+        The example above is illustrative and may need to be adjusted based on the specific dataset structure.
+
+    Args:
+        task (str): Task type, either 'classification' or 'segmentation'
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+    """
+
+    def __init__(self, task: TaskType | str, transform: Transform | None = None) -> None:
+        super().__init__()
+        self.task = TaskType(task)
+        self.transform = transform
+        self._samples: DataFrame | None = None
+        self._category: str | None = None
+
+    @property
+    def name(self) -> str:
+        """Name of the dataset."""
+        class_name = self.__class__.__name__
+
+        # Remove the `_dataset` suffix from the class name
+        if class_name.endswith("Dataset"):
+            class_name = class_name[:-7]
+
+        return class_name
+
+    def __len__(self) -> int:
+        """Get length of the dataset."""
+        return len(self.samples)
+
+    def subsample(self, indices: Sequence[int], inplace: bool = False) -> "AnomalibDataset":
+        """Subsamples the dataset at the provided indices.
+
+        Args:
+            indices (Sequence[int]): Indices at which the dataset is to be subsampled.
+            inplace (bool): When true, the subsampling will be performed on the instance itself.
+                Defaults to ``False``.
+        """
+        if len(set(indices)) != len(indices):
+            msg = "No duplicates allowed in indices."
+            raise ValueError(msg)
+        dataset = self if inplace else copy.deepcopy(self)
+        dataset.samples = self.samples.iloc[indices].reset_index(drop=True)
+        return dataset
+
+    @property
+    def samples(self) -> DataFrame:
+        """Get the samples dataframe."""
+        if self._samples is None:
+            msg = (
+                "Dataset does not have a samples dataframe. Ensure that a dataframe has been assigned to "
+                "`dataset.samples`."
+            )
+            raise RuntimeError(msg)
+        return self._samples
+
+    @samples.setter
+    def samples(self, samples: DataFrame) -> None:
+        """Overwrite the samples with a new dataframe.
+
+        Args:
+            samples (DataFrame): DataFrame with new samples.
+        """
+        # validate the passed samples by checking the
+        if not isinstance(samples, DataFrame):
+            msg = f"samples must be a pandas.DataFrame, found {type(samples)}"
+            raise TypeError(msg)
+
+        expected_columns = _EXPECTED_COLUMNS_PERTASK[self.task]
+        if not all(col in samples.columns for col in expected_columns):
+            msg = f"samples must have (at least) columns {expected_columns}, found {samples.columns}"
+            raise ValueError(msg)
+
+        if not samples["image_path"].apply(lambda p: Path(p).exists()).all():
+            msg = "missing file path(s) in samples"
+            raise FileNotFoundError(msg)
+
+        self._samples = samples.sort_values(by="image_path", ignore_index=True)
+
+    @property
+    def category(self) -> str | None:
+        """Get the category of the dataset."""
+        return self._category
+
+    @category.setter
+    def category(self, category: str) -> None:
+        """Set the category of the dataset."""
+        self._category = category
+
+    @property
+    def has_normal(self) -> bool:
+        """Check if the dataset contains any normal samples."""
+        return LabelName.NORMAL in list(self.samples.label_index)
+
+    @property
+    def has_anomalous(self) -> bool:
+        """Check if the dataset contains any anomalous samples."""
+        return LabelName.ABNORMAL in list(self.samples.label_index)
+
+    def __getitem__(self, index: int) -> dict[str, str | torch.Tensor]:
+        """Get dataset item for the index ``index``.
+
+        Args:
+            index (int): Index to get the item.
+
+        Returns:
+            dict[str, str | torch.Tensor]: Dict of image tensor during training. Otherwise, Dict containing image path,
+                target path, image tensor, label and transformed bounding box.
+        """
+        image_path = self.samples.iloc[index].image_path
+        mask_path = self.samples.iloc[index].mask_path
+        label_index = self.samples.iloc[index].label_index
+
+        image = read_image(image_path, as_tensor=True)
+        item = {"image_path": image_path, "label": label_index}
+
+        if self.task == TaskType.CLASSIFICATION:
+            item["image"] = self.transform(image) if self.transform else image
+        elif self.task in (TaskType.DETECTION, TaskType.SEGMENTATION):
+            # Only Anomalous (1) images have masks in anomaly datasets
+            # Therefore, create empty mask for Normal (0) images.
+            mask = (
+                Mask(torch.zeros(image.shape[-2:])).to(torch.uint8)
+                if label_index == LabelName.NORMAL
+                else read_mask(mask_path, as_tensor=True)
+            )
+            item["image"], item["mask"] = self.transform(image, mask) if self.transform else (image, mask)
+
+            if self.task == TaskType.DETECTION:
+                # create boxes from masks for detection task
+                boxes, _ = masks_to_boxes(item["mask"])
+                item["boxes"] = boxes[0]
+        else:
+            msg = f"Unknown task type: {self.task}"
+            raise ValueError(msg)
+
+        return item
+
+    def __add__(self, other_dataset: "AnomalibDataset") -> "AnomalibDataset":
+        """Concatenate this dataset with another dataset.
+
+        Args:
+            other_dataset (AnomalibDataset): Dataset to concatenate with.
+
+        Returns:
+            AnomalibDataset: Concatenated dataset.
+        """
+        if not isinstance(other_dataset, self.__class__):
+            msg = "Cannot concatenate datasets that are not of the same type."
+            raise TypeError(msg)
+        dataset = copy.deepcopy(self)
+        dataset.samples = pd.concat([self.samples, other_dataset.samples], ignore_index=True)
+        return dataset

anomalib/data/base/depth.py

@@ -0,0 +1,76 @@
+"""Base Depth Dataset."""
+
+# Copyright (C) 2023-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC
+
+import torch
+from PIL import Image
+from torchvision.transforms.functional import to_tensor
+from torchvision.transforms.v2 import Transform
+from torchvision.tv_tensors import Mask
+
+from anomalib import TaskType
+from anomalib.data.base.dataset import AnomalibDataset
+from anomalib.data.utils import LabelName, masks_to_boxes, read_depth_image
+
+
+class AnomalibDepthDataset(AnomalibDataset, ABC):
+    """Base depth anomalib dataset class.
+
+    Args:
+        task (str): Task type, either 'classification' or 'segmentation'
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+    """
+
+    def __init__(self, task: TaskType, transform: Transform | None = None) -> None:
+        super().__init__(task, transform)
+
+        self.transform = transform
+
+    def __getitem__(self, index: int) -> dict[str, str | torch.Tensor]:
+        """Return rgb image, depth image and mask.
+
+        Args:
+            index (int): Index of the item to be returned.
+
+        Returns:
+            dict[str, str | torch.Tensor]: Dictionary containing the image, depth image and mask.
+        """
+        image_path = self.samples.iloc[index].image_path
+        mask_path = self.samples.iloc[index].mask_path
+        label_index = self.samples.iloc[index].label_index
+        depth_path = self.samples.iloc[index].depth_path
+
+        image = to_tensor(Image.open(image_path))
+        depth_image = to_tensor(read_depth_image(depth_path))
+        item = {"image_path": image_path, "depth_path": depth_path, "label": label_index}
+
+        if self.task == TaskType.CLASSIFICATION:
+            item["image"], item["depth_image"] = (
+                self.transform(image, depth_image) if self.transform else (image, depth_image)
+            )
+        elif self.task in (TaskType.DETECTION, TaskType.SEGMENTATION):
+            # Only Anomalous (1) images have masks in anomaly datasets
+            # Therefore, create empty mask for Normal (0) images.
+            mask = (
+                Mask(torch.zeros(image.shape[-2:]))
+                if label_index == LabelName.NORMAL
+                else Mask(to_tensor(Image.open(mask_path)).squeeze())
+            )
+            item["image"], item["depth_image"], item["mask"] = (
+                self.transform(image, depth_image, mask) if self.transform else (image, depth_image, mask)
+            )
+            item["mask_path"] = mask_path
+
+            if self.task == TaskType.DETECTION:
+                # create boxes from masks for detection task
+                boxes, _ = masks_to_boxes(item["mask"])
+                item["boxes"] = boxes[0]
+        else:
+            msg = f"Unknown task type: {self.task}"
+            raise ValueError(msg)
+
+        return item

anomalib/data/base/video.py

@@ -0,0 +1,213 @@
+"""Base Video Dataset."""
+
+# Copyright (C) 2023-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC
+from enum import Enum
+from typing import TYPE_CHECKING, Any
+
+import torch
+from pandas import DataFrame
+from torchvision.transforms.v2 import Transform
+from torchvision.transforms.v2.functional import to_dtype_video
+from torchvision.tv_tensors import Mask
+
+from anomalib import TaskType
+from anomalib.data.base.datamodule import AnomalibDataModule
+from anomalib.data.base.dataset import AnomalibDataset
+from anomalib.data.utils import ValSplitMode, masks_to_boxes
+from anomalib.data.utils.video import ClipsIndexer
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class VideoTargetFrame(str, Enum):
+    """Target frame for a video-clip.
+
+    Used in multi-frame models to determine which frame's ground truth information will be used.
+    """
+
+    FIRST = "first"
+    LAST = "last"
+    MID = "mid"
+    ALL = "all"
+
+
+class AnomalibVideoDataset(AnomalibDataset, ABC):
+    """Base video anomalib dataset class.
+
+    Args:
+        task (str): Task type, either 'classification' or 'segmentation'
+        clip_length_in_frames (int): Number of video frames in each clip.
+        frames_between_clips (int): Number of frames between each consecutive video clip.
+        transform (Transform, optional): Transforms that should be applied to the input clips.
+            Defaults to ``None``.
+        target_frame (VideoTargetFrame): Specifies the target frame in the video clip, used for ground truth retrieval.
+            Defaults to ``VideoTargetFrame.LAST``.
+    """
+
+    def __init__(
+        self,
+        task: TaskType,
+        clip_length_in_frames: int,
+        frames_between_clips: int,
+        transform: Transform | None = None,
+        target_frame: VideoTargetFrame = VideoTargetFrame.LAST,
+    ) -> None:
+        super().__init__(task, transform)
+
+        self.clip_length_in_frames = clip_length_in_frames
+        self.frames_between_clips = frames_between_clips
+        self.transform = transform
+
+        self.indexer: ClipsIndexer | None = None
+        self.indexer_cls: Callable | None = None
+
+        self.target_frame = target_frame
+
+    def __len__(self) -> int:
+        """Get length of the dataset."""
+        if not isinstance(self.indexer, ClipsIndexer):
+            msg = "self.indexer must be an instance of ClipsIndexer."
+            raise TypeError(msg)
+        return self.indexer.num_clips()
+
+    @property
+    def samples(self) -> DataFrame:
+        """Get the samples dataframe."""
+        return super().samples
+
+    @samples.setter
+    def samples(self, samples: DataFrame) -> None:
+        """Overwrite samples and re-index subvideos.
+
+        Args:
+            samples (DataFrame): DataFrame with new samples.
+
+        Raises:
+            ValueError: If the indexer class is not set.
+        """
+        super(AnomalibVideoDataset, self.__class__).samples.fset(self, samples)  # type: ignore[attr-defined]
+        self._setup_clips()
+
+    def _setup_clips(self) -> None:
+        """Compute the video and frame indices of the subvideos.
+
+        Should be called after each change to self._samples
+        """
+        if not callable(self.indexer_cls):
+            msg = "self.indexer_cls must be callable."
+            raise TypeError(msg)
+        self.indexer = self.indexer_cls(  # pylint: disable=not-callable
+            video_paths=list(self.samples.image_path),
+            mask_paths=list(self.samples.mask_path),
+            clip_length_in_frames=self.clip_length_in_frames,
+            frames_between_clips=self.frames_between_clips,
+        )
+
+    def _select_targets(self, item: dict[str, Any]) -> dict[str, Any]:
+        """Select the target frame from the clip.
+
+        Args:
+            item (dict[str, Any]): Item containing the clip information.
+
+        Raises:
+            ValueError: If the target frame is not one of the supported options.
+
+        Returns:
+            dict[str, Any]: Selected item from the clip.
+        """
+        if self.target_frame == VideoTargetFrame.FIRST:
+            idx = 0
+        elif self.target_frame == VideoTargetFrame.LAST:
+            idx = -1
+        elif self.target_frame == VideoTargetFrame.MID:
+            idx = int(self.clip_length_in_frames / 2)
+        else:
+            msg = f"Unknown video target frame: {self.target_frame}"
+            raise ValueError(msg)
+
+        if item.get("mask") is not None:
+            item["mask"] = item["mask"][idx, ...]
+        if item.get("boxes") is not None:
+            item["boxes"] = item["boxes"][idx]
+        if item.get("label") is not None:
+            item["label"] = item["label"][idx]
+        if item.get("original_image") is not None:
+            item["original_image"] = item["original_image"][idx]
+        if item.get("frames") is not None:
+            item["frames"] = item["frames"][idx]
+        return item
+
+    def __getitem__(self, index: int) -> dict[str, str | torch.Tensor]:
+        """Get the dataset item for the index ``index``.
+
+        Args:
+            index (int): Index of the item to be returned.
+
+        Returns:
+            dict[str, str | torch.Tensor]: Dictionary containing the mask, clip and file system information.
+        """
+        if not isinstance(self.indexer, ClipsIndexer):
+            msg = "self.indexer must be an instance of ClipsIndexer."
+            raise TypeError(msg)
+        item = self.indexer.get_item(index)
+        item["image"] = to_dtype_video(video=item["image"], scale=True)
+        # include the untransformed image for visualization
+        item["original_image"] = item["image"].to(torch.uint8)
+
+        # apply transforms
+        if item.get("mask") is not None:
+            if self.transform:
+                item["image"], item["mask"] = self.transform(item["image"], Mask(item["mask"]))
+            item["label"] = torch.Tensor([1 in frame for frame in item["mask"]]).int().squeeze(0)
+            if self.task == TaskType.DETECTION:
+                item["boxes"], _ = masks_to_boxes(item["mask"])
+                item["boxes"] = item["boxes"][0] if len(item["boxes"]) == 1 else item["boxes"]
+        elif self.transform:
+            item["image"] = self.transform(item["image"])
+
+        # squeeze temporal dimensions in case clip length is 1
+        item["image"] = item["image"].squeeze(0)
+
+        # include only target frame in gt
+        if self.clip_length_in_frames > 1 and self.target_frame != VideoTargetFrame.ALL:
+            item = self._select_targets(item)
+
+        if item["mask"] is None:
+            item.pop("mask")
+
+        return item
+
+
+class AnomalibVideoDataModule(AnomalibDataModule):
+    """Base class for video data modules."""
+
+    def _create_test_split(self) -> None:
+        """Video datamodules do not support dynamic assignment of the test split."""
+
+    def _setup(self, _stage: str | None = None) -> None:
+        """Set up the datasets and perform dynamic subset splitting.
+
+        This method may be overridden in subclass for custom splitting behaviour.
+
+        Video datamodules are not compatible with synthetic anomaly generation.
+        """
+        if self.train_data is None:
+            msg = "self.train_data cannot be None."
+            raise ValueError(msg)
+
+        if self.test_data is None:
+            msg = "self.test_data cannot be None."
+            raise ValueError(msg)
+
+        self.train_data.setup()
+        self.test_data.setup()
+
+        if self.val_split_mode == ValSplitMode.SYNTHETIC:
+            msg = f"Val split mode {self.test_split_mode} not supported for video datasets."
+            raise ValueError(msg)
+
+        self._create_val_split()

anomalib/data/depth/__init__.py

@@ -0,0 +1,20 @@
+"""Anomalib Depth Datasets."""
+
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+from enum import Enum
+
+from .folder_3d import Folder3D
+from .mvtec_3d import MVTec3D
+
+
+class DepthDataFormat(str, Enum):
+    """Supported Depth Dataset Types."""
+
+    MVTEC_3D = "mvtec_3d"
+    FOLDER_3D = "folder_3d"
+
+
+__all__ = ["Folder3D", "MVTec3D"]

anomalib/data/depth/folder_3d.py

@@ -0,0 +1,433 @@
+"""Custom Folder Dataset.
+
+This script creates a custom dataset from a folder.
+"""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+from pathlib import Path
+
+from pandas import DataFrame, isna
+from torchvision.transforms.v2 import Transform
+
+from anomalib import TaskType
+from anomalib.data.base import AnomalibDataModule, AnomalibDepthDataset
+from anomalib.data.errors import MisMatchError
+from anomalib.data.utils import (
+    DirType,
+    LabelName,
+    Split,
+    TestSplitMode,
+    ValSplitMode,
+)
+from anomalib.data.utils.path import _prepare_files_labels, validate_and_resolve_path
+
+
+def make_folder3d_dataset(  # noqa: C901
+    normal_dir: str | Path,
+    root: str | Path | None = None,
+    abnormal_dir: str | Path | None = None,
+    normal_test_dir: str | Path | None = None,
+    mask_dir: str | Path | None = None,
+    normal_depth_dir: str | Path | None = None,
+    abnormal_depth_dir: str | Path | None = None,
+    normal_test_depth_dir: str | Path | None = None,
+    split: str | Split | None = None,
+    extensions: tuple[str, ...] | None = None,
+) -> DataFrame:
+    """Make Folder Dataset.
+
+    Args:
+        normal_dir (str | Path): Path to the directory containing normal images.
+        root (str | Path | None): Path to the root directory of the dataset.
+            Defaults to ``None``.
+        abnormal_dir (str | Path | None, optional): Path to the directory containing abnormal images.
+            Defaults to ``None``.
+        normal_test_dir (str | Path | None, optional): Path to the directory containing normal images for the test
+        dataset. Normal test images will be a split of `normal_dir` if `None`.
+            Defaults to ``None``.
+        mask_dir (str | Path | None, optional): Path to the directory containing the mask annotations.
+            Defaults to ``None``.
+        normal_depth_dir (str | Path | None, optional): Path to the directory containing
+            normal depth images for the test dataset. Normal test depth images will be a split of `normal_dir`
+            Defaults to ``None``.
+        abnormal_depth_dir (str | Path | None, optional): Path to the directory containing abnormal depth images for
+            the test dataset.
+            Defaults to ``None``.
+        normal_test_depth_dir (str | Path | None, optional): Path to the directory containing normal depth images for
+            the test dataset. Normal test images will be a split of `normal_dir` if `None`.
+            Defaults to ``None``.
+        split (str | Split | None, optional): Dataset split (ie., Split.FULL, Split.TRAIN or Split.TEST).
+            Defaults to ``None``.
+        extensions (tuple[str, ...] | None, optional): Type of the image extensions to read from the directory.
+            Defaults to ``None``.
+
+    Returns:
+        DataFrame: an output dataframe containing samples for the requested split (ie., train or test)
+    """
+    normal_dir = validate_and_resolve_path(normal_dir, root)
+    abnormal_dir = validate_and_resolve_path(abnormal_dir, root) if abnormal_dir else None
+    normal_test_dir = validate_and_resolve_path(normal_test_dir, root) if normal_test_dir else None
+    mask_dir = validate_and_resolve_path(mask_dir, root) if mask_dir else None
+    normal_depth_dir = validate_and_resolve_path(normal_depth_dir, root) if normal_depth_dir else None
+    abnormal_depth_dir = validate_and_resolve_path(abnormal_depth_dir, root) if abnormal_depth_dir else None
+    normal_test_depth_dir = validate_and_resolve_path(normal_test_depth_dir, root) if normal_test_depth_dir else None
+
+    if not normal_dir.is_dir():
+        msg = "A folder location must be provided in normal_dir."
+        raise ValueError(msg)
+
+    filenames = []
+    labels = []
+    dirs = {DirType.NORMAL: normal_dir}
+
+    if abnormal_dir:
+        dirs[DirType.ABNORMAL] = abnormal_dir
+
+    if normal_test_dir:
+        dirs[DirType.NORMAL_TEST] = normal_test_dir
+
+    if normal_depth_dir:
+        dirs[DirType.NORMAL_DEPTH] = normal_depth_dir
+
+    if abnormal_depth_dir:
+        dirs[DirType.ABNORMAL_DEPTH] = abnormal_depth_dir
+
+    if normal_test_depth_dir:
+        dirs[DirType.NORMAL_TEST_DEPTH] = normal_test_depth_dir
+
+    if mask_dir:
+        dirs[DirType.MASK] = mask_dir
+
+    for dir_type, path in dirs.items():
+        filename, label = _prepare_files_labels(path, dir_type, extensions)
+        filenames += filename
+        labels += label
+
+    samples = DataFrame({"image_path": filenames, "label": labels})
+    samples = samples.sort_values(by="image_path", ignore_index=True)
+
+    # Create label index for normal (0) and abnormal (1) images.
+    samples.loc[
+        (samples.label == DirType.NORMAL) | (samples.label == DirType.NORMAL_TEST),
+        "label_index",
+    ] = LabelName.NORMAL
+    samples.loc[(samples.label == DirType.ABNORMAL), "label_index"] = LabelName.ABNORMAL
+    samples.label_index = samples.label_index.astype("Int64")
+
+    # If a path to mask is provided, add it to the sample dataframe.
+    if normal_depth_dir:
+        samples.loc[samples.label == DirType.NORMAL, "depth_path"] = samples.loc[
+            samples.label == DirType.NORMAL_DEPTH
+        ].image_path.to_numpy()
+        samples.loc[samples.label == DirType.ABNORMAL, "depth_path"] = samples.loc[
+            samples.label == DirType.ABNORMAL_DEPTH
+        ].image_path.to_numpy()
+
+        if normal_test_dir:
+            samples.loc[samples.label == DirType.NORMAL_TEST, "depth_path"] = samples.loc[
+                samples.label == DirType.NORMAL_TEST_DEPTH
+            ].image_path.to_numpy()
+
+        # make sure every rgb image has a corresponding depth image and that the file exists
+        mismatch = (
+            samples.loc[samples.label_index == LabelName.ABNORMAL]
+            .apply(lambda x: Path(x.image_path).stem in Path(x.depth_path).stem, axis=1)
+            .all()
+        )
+        if not mismatch:
+            msg = """Mismatch between anomalous images and depth images. Make sure the mask files
+            in 'xyz' folder follow the same naming convention as the anomalous images in the dataset
+            (e.g. image: '000.png', depth: '000.tiff')."""
+            raise MisMatchError(msg)
+
+        missing_depth_files = samples.depth_path.apply(
+            lambda x: Path(x).exists() if not isna(x) else True,
+        ).all()
+        if not missing_depth_files:
+            msg = "Missing depth image files."
+            raise FileNotFoundError(msg)
+
+        samples = samples.astype({"depth_path": "str"})
+
+    # If a path to mask is provided, add it to the sample dataframe.
+    if mask_dir and abnormal_dir:
+        samples.loc[samples.label == DirType.ABNORMAL, "mask_path"] = samples.loc[
+            samples.label == DirType.MASK
+        ].image_path.to_numpy()
+        samples["mask_path"] = samples["mask_path"].fillna("")
+        samples = samples.astype({"mask_path": "str"})
+
+        # make sure all the files exist
+        if not samples.mask_path.apply(
+            lambda x: Path(x).exists() if x != "" else True,
+        ).all():
+            msg = f"Missing mask files. mask_dir={mask_dir}"
+            raise FileNotFoundError(msg)
+    else:
+        samples["mask_path"] = ""
+
+    # remove all the rows with temporal image samples that have already been assigned
+    samples = samples.loc[
+        (samples.label == DirType.NORMAL) | (samples.label == DirType.ABNORMAL) | (samples.label == DirType.NORMAL_TEST)
+    ]
+
+    # Ensure the pathlib objects are converted to str.
+    # This is because torch dataloader doesn't like pathlib.
+    samples = samples.astype({"image_path": "str"})
+
+    # Create train/test split.
+    # By default, all the normal samples are assigned as train.
+    #   and all the abnormal samples are test.
+    samples.loc[(samples.label == DirType.NORMAL), "split"] = Split.TRAIN
+    samples.loc[(samples.label == DirType.ABNORMAL) | (samples.label == DirType.NORMAL_TEST), "split"] = Split.TEST
+
+    # Get the data frame for the split.
+    if split:
+        samples = samples[samples.split == split]
+        samples = samples.reset_index(drop=True)
+
+    return samples
+
+
+class Folder3DDataset(AnomalibDepthDataset):
+    """Folder dataset.
+
+    Args:
+        name (str): Name of the dataset.
+        task (TaskType): Task type. (``classification``, ``detection`` or ``segmentation``).
+        transform (Transform): Transforms that should be applied to the input images.
+        normal_dir (str | Path): Path to the directory containing normal images.
+        root (str | Path | None): Root folder of the dataset.
+            Defaults to ``None``.
+        abnormal_dir (str | Path | None, optional): Path to the directory containing abnormal images.
+            Defaults to ``None``.
+        normal_test_dir (str | Path | None, optional): Path to the directory containing
+            normal images for the test dataset.
+            Defaults to ``None``.
+        mask_dir (str | Path | None, optional): Path to the directory containing
+            the mask annotations.
+            Defaults to ``None``.
+        normal_depth_dir (str | Path | None, optional): Path to the directory containing
+            normal depth images for the test dataset. Normal test depth images will be a split of `normal_dir`
+            Defaults to ``None``.
+        abnormal_depth_dir (str | Path | None, optional): Path to the directory containing abnormal depth images for
+            the test dataset.
+            Defaults to ``None``.
+        normal_test_depth_dir (str | Path | None, optional): Path to the directory containing
+            normal depth images for the test dataset. Normal test images will be a split of `normal_dir` if `None`.
+            Defaults to ``None``.
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        split (str | Split | None): Fixed subset split that follows from folder structure on file system.
+            Choose from [Split.FULL, Split.TRAIN, Split.TEST]
+            Defaults to ``None``.
+        extensions (tuple[str, ...] | None, optional): Type of the image extensions to read from the directory.
+            Defaults to ``None``.
+
+    Raises:
+        ValueError: When task is set to classification and `mask_dir` is provided. When `mask_dir` is
+            provided, `task` should be set to `segmentation`.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        task: TaskType,
+        normal_dir: str | Path,
+        root: str | Path | None = None,
+        abnormal_dir: str | Path | None = None,
+        normal_test_dir: str | Path | None = None,
+        mask_dir: str | Path | None = None,
+        normal_depth_dir: str | Path | None = None,
+        abnormal_depth_dir: str | Path | None = None,
+        normal_test_depth_dir: str | Path | None = None,
+        transform: Transform | None = None,
+        split: str | Split | None = None,
+        extensions: tuple[str, ...] | None = None,
+    ) -> None:
+        super().__init__(task, transform)
+
+        self._name = name
+        self.split = split
+        self.root = root
+        self.normal_dir = normal_dir
+        self.abnormal_dir = abnormal_dir
+        self.normal_test_dir = normal_test_dir
+        self.mask_dir = mask_dir
+        self.normal_depth_dir = normal_depth_dir
+        self.abnormal_depth_dir = abnormal_depth_dir
+        self.normal_test_depth_dir = normal_test_depth_dir
+        self.extensions = extensions
+
+        self.samples = make_folder3d_dataset(
+            root=self.root,
+            normal_dir=self.normal_dir,
+            abnormal_dir=self.abnormal_dir,
+            normal_test_dir=self.normal_test_dir,
+            mask_dir=self.mask_dir,
+            normal_depth_dir=self.normal_depth_dir,
+            abnormal_depth_dir=self.abnormal_depth_dir,
+            normal_test_depth_dir=self.normal_test_depth_dir,
+            split=self.split,
+            extensions=self.extensions,
+        )
+
+    @property
+    def name(self) -> str:
+        """Name of the dataset.
+
+        Folder3D dataset overrides the name property to provide a custom name.
+        """
+        return self._name
+
+
+class Folder3D(AnomalibDataModule):
+    """Folder DataModule.
+
+    Args:
+        name (str): Name of the dataset. This is used to name the datamodule, especially when logging/saving.
+        normal_dir (str | Path): Name of the directory containing normal images.
+        root (str | Path | None): Path to the root folder containing normal and abnormal dirs.
+            Defaults to ``None``.
+        abnormal_dir (str | Path | None): Name of the directory containing abnormal images.
+            Defaults to ``abnormal``.
+        normal_test_dir (str | Path | None, optional): Path to the directory containing normal images for the test
+            dataset.
+            Defaults to ``None``.
+        mask_dir (str | Path | None, optional): Path to the directory containing the mask annotations.
+            Defaults to ``None``.
+        normal_depth_dir (str | Path | None, optional): Path to the directory containing
+            normal depth images for the test dataset. Normal test depth images will be a split of `normal_dir`
+        abnormal_depth_dir (str | Path | None, optional): Path to the directory containing
+            abnormal depth images for the test dataset.
+        normal_test_depth_dir (str | Path | None, optional): Path to the directory containing
+            normal depth images for the test dataset. Normal test images will be a split of `normal_dir`
+            if `None`. Defaults to None.
+        normal_split_ratio (float, optional): Ratio to split normal training images and add to the
+            test set in case test set doesn't contain any normal images.
+            Defaults to 0.2.
+        extensions (tuple[str, ...] | None, optional): Type of the image extensions to read from the
+            directory. Defaults to None.
+        train_batch_size (int, optional): Training batch size.
+            Defaults to ``32``.
+        eval_batch_size (int, optional): Test batch size.
+            Defaults to ``32``.
+        num_workers (int, optional): Number of workers.
+            Defaults to ``8``.
+        task (TaskType, optional): Task type. Could be ``classification``, ``detection`` or ``segmentation``.
+            Defaults to ``TaskType.SEGMENTATION``.
+        image_size (tuple[int, int], optional): Size to which input images should be resized.
+            Defaults to ``None``.
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        train_transform (Transform, optional): Transforms that should be applied to the input images during training.
+            Defaults to ``None``.
+        eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation.
+            Defaults to ``None``.
+        test_split_mode (TestSplitMode): Setting that determines how the testing subset is obtained.
+            Defaults to ``TestSplitMode.FROM_DIR``.
+        test_split_ratio (float): Fraction of images from the train set that will be reserved for testing.
+            Defaults to ``0.2``.
+        val_split_mode (ValSplitMode): Setting that determines how the validation subset is obtained.
+            Defaults to ``ValSplitMode.FROM_TEST``.
+        val_split_ratio (float): Fraction of train or test images that will be reserved for validation.
+            Defaults to ``0.5``.
+        seed (int | None, optional): Seed used during random subset splitting.
+            Defaults to ``None``.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        normal_dir: str | Path,
+        root: str | Path,
+        abnormal_dir: str | Path | None = None,
+        normal_test_dir: str | Path | None = None,
+        mask_dir: str | Path | None = None,
+        normal_depth_dir: str | Path | None = None,
+        abnormal_depth_dir: str | Path | None = None,
+        normal_test_depth_dir: str | Path | None = None,
+        extensions: tuple[str] | None = None,
+        train_batch_size: int = 32,
+        eval_batch_size: int = 32,
+        num_workers: int = 8,
+        task: TaskType | str = TaskType.SEGMENTATION,
+        image_size: tuple[int, int] | None = None,
+        transform: Transform | None = None,
+        train_transform: Transform | None = None,
+        eval_transform: Transform | None = None,
+        test_split_mode: TestSplitMode | str = TestSplitMode.FROM_DIR,
+        test_split_ratio: float = 0.2,
+        val_split_mode: ValSplitMode | str = ValSplitMode.FROM_TEST,
+        val_split_ratio: float = 0.5,
+        seed: int | None = None,
+    ) -> None:
+        super().__init__(
+            train_batch_size=train_batch_size,
+            eval_batch_size=eval_batch_size,
+            num_workers=num_workers,
+            image_size=image_size,
+            transform=transform,
+            train_transform=train_transform,
+            eval_transform=eval_transform,
+            test_split_mode=test_split_mode,
+            test_split_ratio=test_split_ratio,
+            val_split_mode=val_split_mode,
+            val_split_ratio=val_split_ratio,
+            seed=seed,
+        )
+        self._name = name
+        self.task = TaskType(task)
+        self.root = Path(root)
+        self.normal_dir = normal_dir
+        self.abnormal_dir = abnormal_dir
+        self.normal_test_dir = normal_test_dir
+        self.mask_dir = mask_dir
+        self.normal_depth_dir = normal_depth_dir
+        self.abnormal_depth_dir = abnormal_depth_dir
+        self.normal_test_depth_dir = normal_test_depth_dir
+        self.extensions = extensions
+
+    def _setup(self, _stage: str | None = None) -> None:
+        self.train_data = Folder3DDataset(
+            name=self.name,
+            task=self.task,
+            transform=self.train_transform,
+            split=Split.TRAIN,
+            root=self.root,
+            normal_dir=self.normal_dir,
+            abnormal_dir=self.abnormal_dir,
+            normal_test_dir=self.normal_test_dir,
+            mask_dir=self.mask_dir,
+            normal_depth_dir=self.normal_depth_dir,
+            abnormal_depth_dir=self.abnormal_depth_dir,
+            normal_test_depth_dir=self.normal_test_depth_dir,
+            extensions=self.extensions,
+        )
+
+        self.test_data = Folder3DDataset(
+            name=self.name,
+            task=self.task,
+            transform=self.eval_transform,
+            split=Split.TEST,
+            root=self.root,
+            normal_dir=self.normal_dir,
+            abnormal_dir=self.abnormal_dir,
+            normal_test_dir=self.normal_test_dir,
+            normal_depth_dir=self.normal_depth_dir,
+            abnormal_depth_dir=self.abnormal_depth_dir,
+            normal_test_depth_dir=self.normal_test_depth_dir,
+            mask_dir=self.mask_dir,
+            extensions=self.extensions,
+        )
+
+    @property
+    def name(self) -> str:
+        """Name of the datamodule.
+
+        Folder3D datamodule overrides the name property to provide a custom name.
+        """
+        return self._name

anomalib/data/depth/mvtec_3d.py

@@ -0,0 +1,302 @@
+"""MVTec 3D-AD Dataset (CC BY-NC-SA 4.0).
+
+Description:
+    This script contains PyTorch Dataset, Dataloader and PyTorch Lightning DataModule for the MVTec 3D-AD dataset.
+    If the dataset is not on the file system, the script downloads and extracts the dataset and create PyTorch data
+    objects.
+
+License:
+    MVTec 3D-AD dataset is released under the Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
+        License (CC BY-NC-SA 4.0)(https://creativecommons.org/licenses/by-nc-sa/4.0/).
+
+Reference:
+    - Paul Bergmann, Xin Jin, David Sattlegger, Carsten Steger: The MVTec 3D-AD Dataset for Unsupervised 3D Anomaly
+        Detection and Localization in: Proceedings of the 17th International Joint Conference on Computer Vision,
+        Imaging and Computer Graphics Theory and Applications - Volume 5: VISAPP, 202-213, 2022, DOI: 10.5220/
+        0010865000003124.
+"""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+import logging
+from collections.abc import Sequence
+from pathlib import Path
+
+from pandas import DataFrame
+from torchvision.transforms.v2 import Transform
+
+from anomalib import TaskType
+from anomalib.data.base import AnomalibDataModule, AnomalibDepthDataset
+from anomalib.data.errors import MisMatchError
+from anomalib.data.utils import (
+    DownloadInfo,
+    LabelName,
+    Split,
+    TestSplitMode,
+    ValSplitMode,
+    download_and_extract,
+    validate_path,
+)
+
+logger = logging.getLogger(__name__)
+
+
+IMG_EXTENSIONS = [".png", ".PNG", ".tiff"]
+
+DOWNLOAD_INFO = DownloadInfo(
+    name="mvtec_3d",
+    url="https://www.mydrive.ch/shares/45920/dd1eb345346df066c63b5c95676b961b/download/428824485-1643285832"
+    "/mvtec_3d_anomaly_detection.tar.xz",
+    hashsum="d8bb2800fbf3ac88e798da6ae10dc819",
+)
+
+CATEGORIES = ("bagel", "cable_gland", "carrot", "cookie", "dowel", "foam", "peach", "potato", "rope", "tire")
+
+
+def make_mvtec_3d_dataset(
+    root: str | Path,
+    split: str | Split | None = None,
+    extensions: Sequence[str] | None = None,
+) -> DataFrame:
+    """Create MVTec 3D-AD samples by parsing the MVTec AD data file structure.
+
+    The files are expected to follow this structure:
+    - `path/to/dataset/split/category/image_filename.png`
+    - `path/to/dataset/ground_truth/category/mask_filename.png`
+
+    This function creates a DataFrame to store the parsed information. The DataFrame follows this format:
+
+    +---+---------------+-------+---------+---------------+---------------------------------------+-------------+
+    |   | path          | split | label   | image_path    | mask_path                             | label_index |
+    +---+---------------+-------+---------+---------------+---------------------------------------+-------------+
+    | 0 | datasets/name | test  | defect  | filename.png  | ground_truth/defect/filename_mask.png | 1           |
+    +---+---------------+-------+---------+---------------+---------------------------------------+-------------+
+
+    Args:
+        root (Path): Path to the dataset.
+        split (str | Split | None, optional): Dataset split (e.g., 'train' or 'test').
+            Defaults to ``None``.
+        extensions (Sequence[str] | None, optional): List of file extensions to be included in the dataset.
+            Defaults to ``None``.
+
+    Examples:
+        The following example shows how to get training samples from the MVTec 3D-AD 'bagel' category:
+
+        >>> from pathlib import Path
+        >>> root = Path('./MVTec3D')
+        >>> category = 'bagel'
+        >>> path = root / category
+        >>> print(path)
+        PosixPath('MVTec3D/bagel')
+
+        >>> samples = create_mvtec_3d_ad_samples(path, split='train')
+        >>> print(samples.head())
+            path          split label image_path                          mask_path                        label_index
+            MVTec3D/bagel train good MVTec3D/bagel/train/good/rgb/105.png MVTec3D/bagel/ground_truth/good/gt/105.png 0
+            MVTec3D/bagel train good MVTec3D/bagel/train/good/rgb/017.png MVTec3D/bagel/ground_truth/good/gt/017.png 0
+
+    Returns:
+        DataFrame: An output DataFrame containing the samples of the dataset.
+    """
+    if extensions is None:
+        extensions = IMG_EXTENSIONS
+
+    root = validate_path(root)
+    samples_list = [(str(root),) + f.parts[-4:] for f in root.glob(r"**/*") if f.suffix in extensions]
+    if not samples_list:
+        msg = f"Found 0 images in {root}"
+        raise RuntimeError(msg)
+
+    samples = DataFrame(samples_list, columns=["path", "split", "label", "type", "file_name"])
+
+    # Modify image_path column by converting to absolute path
+    samples.loc[(samples.type == "rgb"), "image_path"] = (
+        samples.path + "/" + samples.split + "/" + samples.label + "/" + "rgb/" + samples.file_name
+    )
+    samples.loc[(samples.type == "rgb"), "depth_path"] = (
+        samples.path
+        + "/"
+        + samples.split
+        + "/"
+        + samples.label
+        + "/"
+        + "xyz/"
+        + samples.file_name.str.split(".").str[0]
+        + ".tiff"
+    )
+
+    # Create label index for normal (0) and anomalous (1) images.
+    samples.loc[(samples.label == "good"), "label_index"] = LabelName.NORMAL
+    samples.loc[(samples.label != "good"), "label_index"] = LabelName.ABNORMAL
+    samples.label_index = samples.label_index.astype(int)
+
+    # separate masks from samples
+    mask_samples = samples.loc[((samples.split == "test") & (samples.type == "rgb"))].sort_values(
+        by="image_path",
+        ignore_index=True,
+    )
+    samples = samples.sort_values(by="image_path", ignore_index=True)
+
+    # assign mask paths to all test images
+    samples.loc[((samples.split == "test") & (samples.type == "rgb")), "mask_path"] = (
+        mask_samples.path + "/" + samples.split + "/" + samples.label + "/" + "gt/" + samples.file_name
+    )
+    samples = samples.dropna(subset=["image_path"])
+    samples = samples.astype({"image_path": "str", "mask_path": "str", "depth_path": "str"})
+
+    # assert that the right mask files are associated with the right test images
+    mismatch_masks = (
+        samples.loc[samples.label_index == LabelName.ABNORMAL]
+        .apply(lambda x: Path(x.image_path).stem in Path(x.mask_path).stem, axis=1)
+        .all()
+    )
+    if not mismatch_masks:
+        msg = """Mismatch between anomalous images and ground truth masks. Make sure the mask files
+          in 'ground_truth' folder follow the same naming convention as the anomalous images in
+          the dataset (e.g. image: '000.png', mask: '000.png' or '000_mask.png')."""
+        raise MisMatchError(msg)
+
+    mismatch_depth = (
+        samples.loc[samples.label_index == LabelName.ABNORMAL]
+        .apply(lambda x: Path(x.image_path).stem in Path(x.depth_path).stem, axis=1)
+        .all()
+    )
+    if not mismatch_depth:
+        msg = """Mismatch between anomalous images and depth images. Make sure the mask files in
+          'xyz' folder follow the same naming convention as the anomalous images in the dataset
+          (e.g. image: '000.png', depth: '000.tiff')."""
+        raise MisMatchError(msg)
+
+    if split:
+        samples = samples[samples.split == split].reset_index(drop=True)
+
+    return samples
+
+
+class MVTec3DDataset(AnomalibDepthDataset):
+    """MVTec 3D dataset class.
+
+    Args:
+        task (TaskType): Task type, ``classification``, ``detection`` or ``segmentation``
+        root (Path | str): Path to the root of the dataset
+            Defaults to ``"./datasets/MVTec3D"``.
+        category (str): Sub-category of the dataset, e.g. 'bagel'
+            Defaults to ``"bagel"``.
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        split (str | Split | None): Split of the dataset, usually Split.TRAIN or Split.TEST
+            Defaults to ``None``.
+    """
+
+    def __init__(
+        self,
+        task: TaskType,
+        root: Path | str = "./datasets/MVTec3D",
+        category: str = "bagel",
+        transform: Transform | None = None,
+        split: str | Split | None = None,
+    ) -> None:
+        super().__init__(task=task, transform=transform)
+
+        self.root_category = Path(root) / Path(category)
+        self.split = split
+        self.samples = make_mvtec_3d_dataset(self.root_category, split=self.split, extensions=IMG_EXTENSIONS)
+
+
+class MVTec3D(AnomalibDataModule):
+    """MVTec Datamodule.
+
+    Args:
+        root (Path | str): Path to the root of the dataset
+            Defaults to ``"./datasets/MVTec3D"``.
+        category (str): Category of the MVTec dataset (e.g. "bottle" or "cable").
+            Defaults to ``bagel``.
+        train_batch_size (int, optional): Training batch size.
+            Defaults to ``32``.
+        eval_batch_size (int, optional): Test batch size.
+            Defaults to ``32``.
+        num_workers (int, optional): Number of workers.
+            Defaults to ``8``.
+        task (TaskType): Task type, 'classification', 'detection' or 'segmentation'
+            Defaults to ``TaskType.SEGMENTATION``.
+        image_size (tuple[int, int], optional): Size to which input images should be resized.
+            Defaults to ``None``.
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        train_transform (Transform, optional): Transforms that should be applied to the input images during training.
+            Defaults to ``None``.
+        eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation.
+            Defaults to ``None``.
+        test_split_mode (TestSplitMode): Setting that determines how the testing subset is obtained.
+            Defaults to ``TestSplitMode.FROM_DIR``.
+        test_split_ratio (float): Fraction of images from the train set that will be reserved for testing.
+            Defaults to ``0.2``.
+        val_split_mode (ValSplitMode): Setting that determines how the validation subset is obtained.
+            Defaults to ``ValSplitMode.SAME_AS_TEST``.
+        val_split_ratio (float): Fraction of train or test images that will be reserved for validation.
+            Defaults to ``0.5``.
+        seed (int | None, optional): Seed which may be set to a fixed value for reproducibility.
+            Defaults to ``None``.
+    """
+
+    def __init__(
+        self,
+        root: Path | str = "./datasets/MVTec3D",
+        category: str = "bagel",
+        train_batch_size: int = 32,
+        eval_batch_size: int = 32,
+        num_workers: int = 8,
+        task: TaskType | str = TaskType.SEGMENTATION,
+        image_size: tuple[int, int] | None = None,
+        transform: Transform | None = None,
+        train_transform: Transform | None = None,
+        eval_transform: Transform | None = None,
+        test_split_mode: TestSplitMode | str = TestSplitMode.FROM_DIR,
+        test_split_ratio: float = 0.2,
+        val_split_mode: ValSplitMode | str = ValSplitMode.SAME_AS_TEST,
+        val_split_ratio: float = 0.5,
+        seed: int | None = None,
+    ) -> None:
+        super().__init__(
+            train_batch_size=train_batch_size,
+            eval_batch_size=eval_batch_size,
+            num_workers=num_workers,
+            image_size=image_size,
+            transform=transform,
+            train_transform=train_transform,
+            eval_transform=eval_transform,
+            test_split_mode=test_split_mode,
+            test_split_ratio=test_split_ratio,
+            val_split_mode=val_split_mode,
+            val_split_ratio=val_split_ratio,
+            seed=seed,
+        )
+
+        self.task = TaskType(task)
+        self.root = Path(root)
+        self.category = category
+
+    def _setup(self, _stage: str | None = None) -> None:
+        self.train_data = MVTec3DDataset(
+            task=self.task,
+            transform=self.train_transform,
+            split=Split.TRAIN,
+            root=self.root,
+            category=self.category,
+        )
+        self.test_data = MVTec3DDataset(
+            task=self.task,
+            transform=self.eval_transform,
+            split=Split.TEST,
+            root=self.root,
+            category=self.category,
+        )
+
+    def prepare_data(self) -> None:
+        """Download the dataset if not available."""
+        if (self.root / self.category).is_dir():
+            logger.info("Found the dataset.")
+        else:
+            download_and_extract(self.root, DOWNLOAD_INFO)

anomalib/data/errors.py

@@ -0,0 +1,19 @@
+"""Custom Exception Class for Mismatch Detection (MisMatchError)."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+class MisMatchError(Exception):
+    """Exception raised when a mismatch is detected.
+
+    Attributes:
+        message (str): Explanation of the error.
+    """
+
+    def __init__(self, message: str = "") -> None:
+        if message:
+            self.message = message
+        else:
+            self.message = "Mismatch detected."
+        super().__init__(self.message)

anomalib/data/image/__init__.py

@@ -0,0 +1,33 @@
+"""Anomalib Image Datasets.
+
+This module contains the supported image datasets for Anomalib.
+"""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+from enum import Enum
+
+from .btech import BTech
+from .folder import Folder
+from .kolektor import Kolektor
+from .mvtec import MVTec
+from .mvtec_loco import MVTecLoco
+from .visa import Visa
+
+
+class ImageDataFormat(str, Enum):
+    """Supported Image Dataset Types."""
+
+    MVTEC = "mvtec"
+    MVTEC_3D = "mvtec_3d"
+    MVTEC_LOCO = "mvtec_loco"
+    BTECH = "btech"
+    KOLEKTOR = "kolektor"
+    FOLDER = "folder"
+    FOLDER_3D = "folder_3d"
+    VISA = "visa"
+
+
+__all__ = ["BTech", "Folder", "Kolektor", "MVTec", "MVTecLoco", "Visa"]

anomalib/data/image/btech.py

@@ -0,0 +1,362 @@
+"""BTech Dataset.
+
+This script contains PyTorch Lightning DataModule for the BTech dataset.
+
+If the dataset is not on the file system, the script downloads and
+extracts the dataset and create PyTorch data objects.
+"""
+
+# Copyright (C) 2022-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import shutil
+from pathlib import Path
+
+import cv2
+import pandas as pd
+from pandas.core.frame import DataFrame
+from torchvision.transforms.v2 import Transform
+from tqdm import tqdm
+
+from anomalib import TaskType
+from anomalib.data.base import AnomalibDataModule, AnomalibDataset
+from anomalib.data.utils import (
+    DownloadInfo,
+    LabelName,
+    Split,
+    TestSplitMode,
+    ValSplitMode,
+    download_and_extract,
+    validate_path,
+)
+
+logger = logging.getLogger(__name__)
+
+DOWNLOAD_INFO = DownloadInfo(
+    name="btech",
+    url="https://avires.dimi.uniud.it/papers/btad/btad.zip",
+    hashsum="461c9387e515bfed41ecaae07c50cf6b10def647b36c9e31d239ab2736b10d2a",
+)
+
+CATEGORIES = ("01", "02", "03")
+
+
+def make_btech_dataset(path: Path, split: str | Split | None = None) -> DataFrame:
+    """Create BTech samples by parsing the BTech data file structure.
+
+    The files are expected to follow the structure:
+
+        .. code-block:: bash
+
+            path/to/dataset/split/category/image_filename.png
+            path/to/dataset/ground_truth/category/mask_filename.png
+
+    Args:
+        path (Path): Path to dataset
+        split (str | Split | None, optional): Dataset split (ie., either train or test).
+            Defaults to ``None``.
+
+    Example:
+        The following example shows how to get training samples from BTech 01 category:
+
+        .. code-block:: python
+
+            >>> root = Path('./BTech')
+            >>> category = '01'
+            >>> path = root / category
+            >>> path
+            PosixPath('BTech/01')
+
+            >>> samples = make_btech_dataset(path, split='train')
+            >>> samples.head()
+            path     split label image_path                  mask_path                     label_index
+            0  BTech/01 train 01    BTech/01/train/ok/105.bmp BTech/01/ground_truth/ok/105.png      0
+            1  BTech/01 train 01    BTech/01/train/ok/017.bmp BTech/01/ground_truth/ok/017.png      0
+            ...
+
+    Returns:
+        DataFrame: an output dataframe containing samples for the requested split (ie., train or test)
+    """
+    path = validate_path(path)
+
+    samples_list = [
+        (str(path),) + filename.parts[-3:] for filename in path.glob("**/*") if filename.suffix in (".bmp", ".png")
+    ]
+    if not samples_list:
+        msg = f"Found 0 images in {path}"
+        raise RuntimeError(msg)
+
+    samples = pd.DataFrame(samples_list, columns=["path", "split", "label", "image_path"])
+    samples = samples[samples.split != "ground_truth"]
+
+    # Create mask_path column
+    # (safely handles cases where non-mask image_paths end with either .png or .bmp)
+    samples["mask_path"] = (
+        samples.path
+        + "/ground_truth/"
+        + samples.label
+        + "/"
+        + samples.image_path.str.rstrip("png").str.rstrip(".").str.rstrip("bmp").str.rstrip(".")
+        + ".png"
+    )
+
+    # Modify image_path column by converting to absolute path
+    samples["image_path"] = samples.path + "/" + samples.split + "/" + samples.label + "/" + samples.image_path
+
+    # Good images don't have mask
+    samples.loc[(samples.split == "test") & (samples.label == "ok"), "mask_path"] = ""
+
+    # Create label index for normal (0) and anomalous (1) images.
+    samples.loc[(samples.label == "ok"), "label_index"] = LabelName.NORMAL
+    samples.loc[(samples.label != "ok"), "label_index"] = LabelName.ABNORMAL
+    samples.label_index = samples.label_index.astype(int)
+
+    # Get the data frame for the split.
+    if split:
+        samples = samples[samples.split == split]
+        samples = samples.reset_index(drop=True)
+
+    return samples
+
+
+class BTechDataset(AnomalibDataset):
+    """Btech Dataset class.
+
+    Args:
+        root: Path to the BTech dataset
+        category: Name of the BTech category.
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        split: 'train', 'val' or 'test'
+        task: ``classification``, ``detection`` or ``segmentation``
+        create_validation_set: Create a validation subset in addition to the train and test subsets
+
+    Examples:
+        >>> from anomalib.data.image.btech import BTechDataset
+        >>> from anomalib.data.utils.transforms import get_transforms
+        >>> transform = get_transforms(image_size=256)
+        >>> dataset = BTechDataset(
+        ...     task="classification",
+        ...     transform=transform,
+        ...     root='./datasets/BTech',
+        ...     category='01',
+        ... )
+        >>> dataset[0].keys()
+        >>> dataset.setup()
+        dict_keys(['image'])
+
+        >>> dataset.split = "test"
+        >>> dataset[0].keys()
+        dict_keys(['image', 'image_path', 'label'])
+
+        >>> dataset.task = "segmentation"
+        >>> dataset.split = "train"
+        >>> dataset[0].keys()
+        dict_keys(['image'])
+
+        >>> dataset.split = "test"
+        >>> dataset[0].keys()
+        dict_keys(['image_path', 'label', 'mask_path', 'image', 'mask'])
+
+        >>> dataset[0]["image"].shape, dataset[0]["mask"].shape
+        (torch.Size([3, 256, 256]), torch.Size([256, 256]))
+    """
+
+    def __init__(
+        self,
+        root: str | Path,
+        category: str,
+        transform: Transform | None = None,
+        split: str | Split | None = None,
+        task: TaskType | str = TaskType.SEGMENTATION,
+    ) -> None:
+        super().__init__(task, transform)
+
+        self.root_category = Path(root) / category
+        self.split = split
+        self.samples = make_btech_dataset(path=self.root_category, split=self.split)
+
+
+class BTech(AnomalibDataModule):
+    """BTech Lightning Data Module.
+
+    Args:
+        root (Path | str): Path to the BTech dataset.
+            Defaults to ``"./datasets/BTech"``.
+        category (str): Name of the BTech category.
+            Defaults to ``"01"``.
+        train_batch_size (int, optional): Training batch size.
+            Defaults to ``32``.
+        eval_batch_size (int, optional): Eval batch size.
+            Defaults to ``32``.
+        num_workers (int, optional): Number of workers.
+            Defaults to ``8``.
+        task (TaskType, optional): Task type.
+            Defaults to ``TaskType.SEGMENTATION``.
+        image_size (tuple[int, int], optional): Size to which input images should be resized.
+            Defaults to ``None``.
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        train_transform (Transform, optional): Transforms that should be applied to the input images during training.
+            Defaults to ``None``.
+        eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation.
+            Defaults to ``None``.
+        test_split_mode (TestSplitMode, optional): Setting that determines how the testing subset is obtained.
+            Defaults to ``TestSplitMode.FROM_DIR``.
+        test_split_ratio (float, optional): Fraction of images from the train set that will be reserved for testing.
+            Defaults to ``0.2``.
+        val_split_mode (ValSplitMode, optional): Setting that determines how the validation subset is obtained.
+            Defaults to ``ValSplitMode.SAME_AS_TEST``.
+        val_split_ratio (float, optional): Fraction of train or test images that will be reserved for validation.
+            Defaults to ``0.5``.
+        seed (int | None, optional): Seed which may be set to a fixed value for reproducibility.
+            Defaults to ``None``.
+
+    Examples:
+        To create the BTech datamodule, we need to instantiate the class, and call the ``setup`` method.
+
+        >>> from anomalib.data import BTech
+        >>> datamodule = BTech(
+        ...     root="./datasets/BTech",
+        ...     category="01",
+        ...     image_size=256,
+        ...     train_batch_size=32,
+        ...     eval_batch_size=32,
+        ...     num_workers=8,
+        ...     transform_config_train=None,
+        ...     transform_config_eval=None,
+        ... )
+        >>> datamodule.setup()
+
+        To get the train dataloader and the first batch of data:
+
+        >>> i, data = next(enumerate(datamodule.train_dataloader()))
+        >>> data.keys()
+        dict_keys(['image'])
+        >>> data["image"].shape
+        torch.Size([32, 3, 256, 256])
+
+        To access the validation dataloader and the first batch of data:
+
+        >>> i, data = next(enumerate(datamodule.val_dataloader()))
+        >>> data.keys()
+        dict_keys(['image_path', 'label', 'mask_path', 'image', 'mask'])
+        >>> data["image"].shape, data["mask"].shape
+        (torch.Size([32, 3, 256, 256]), torch.Size([32, 256, 256]))
+
+        Similarly, to access the test dataloader and the first batch of data:
+
+        >>> i, data = next(enumerate(datamodule.test_dataloader()))
+        >>> data.keys()
+        dict_keys(['image_path', 'label', 'mask_path', 'image', 'mask'])
+        >>> data["image"].shape, data["mask"].shape
+        (torch.Size([32, 3, 256, 256]), torch.Size([32, 256, 256]))
+    """
+
+    def __init__(
+        self,
+        root: Path | str = "./datasets/BTech",
+        category: str = "01",
+        train_batch_size: int = 32,
+        eval_batch_size: int = 32,
+        num_workers: int = 8,
+        task: TaskType | str = TaskType.SEGMENTATION,
+        image_size: tuple[int, int] | None = None,
+        transform: Transform | None = None,
+        train_transform: Transform | None = None,
+        eval_transform: Transform | None = None,
+        test_split_mode: TestSplitMode | str = TestSplitMode.FROM_DIR,
+        test_split_ratio: float = 0.2,
+        val_split_mode: ValSplitMode | str = ValSplitMode.SAME_AS_TEST,
+        val_split_ratio: float = 0.5,
+        seed: int | None = None,
+    ) -> None:
+        super().__init__(
+            train_batch_size=train_batch_size,
+            eval_batch_size=eval_batch_size,
+            num_workers=num_workers,
+            image_size=image_size,
+            transform=transform,
+            train_transform=train_transform,
+            eval_transform=eval_transform,
+            test_split_mode=test_split_mode,
+            test_split_ratio=test_split_ratio,
+            val_split_mode=val_split_mode,
+            val_split_ratio=val_split_ratio,
+            seed=seed,
+        )
+
+        self.root = Path(root)
+        self.category = category
+        self.task = TaskType(task)
+
+    def _setup(self, _stage: str | None = None) -> None:
+        self.train_data = BTechDataset(
+            task=self.task,
+            transform=self.train_transform,
+            split=Split.TRAIN,
+            root=self.root,
+            category=self.category,
+        )
+        self.test_data = BTechDataset(
+            task=self.task,
+            transform=self.eval_transform,
+            split=Split.TEST,
+            root=self.root,
+            category=self.category,
+        )
+
+    def prepare_data(self) -> None:
+        """Download the dataset if not available.
+
+        This method checks if the specified dataset is available in the file system.
+        If not, it downloads and extracts the dataset into the appropriate directory.
+
+        Example:
+            Assume the dataset is not available on the file system.
+            Here's how the directory structure looks before and after calling the
+            `prepare_data` method:
+
+            Before:
+
+            .. code-block:: bash
+
+                $ tree datasets
+                datasets
+                ├── dataset1
+                └── dataset2
+
+            Calling the method:
+
+            .. code-block:: python
+
+                >> datamodule = BTech(root="./datasets/BTech", category="01")
+                >> datamodule.prepare_data()
+
+            After:
+
+            .. code-block:: bash
+
+                $ tree datasets
+                datasets
+                ├── dataset1
+                ├── dataset2
+                └── BTech
+                    ├── 01
+                    ├── 02
+                    └── 03
+        """
+        if (self.root / self.category).is_dir():
+            logger.info("Found the dataset.")
+        else:
+            download_and_extract(self.root.parent, DOWNLOAD_INFO)
+
+            # rename folder and convert images
+            logger.info("Renaming the dataset directory")
+            shutil.move(src=str(self.root.parent / "BTech_Dataset_transformed"), dst=str(self.root))
+            logger.info("Convert the bmp formats to png to have consistent image extensions")
+            for filename in tqdm(self.root.glob("**/*.bmp"), desc="Converting bmp to png"):
+                image = cv2.imread(str(filename))
+                cv2.imwrite(str(filename.with_suffix(".png")), image)
+                filename.unlink()

anomalib/data/image/folder.py

@@ -0,0 +1,478 @@
+"""Custom Folder Dataset.
+
+This script creates a custom dataset from a folder.
+"""
+
+# Copyright (C) 2022-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Sequence
+from pathlib import Path
+
+from pandas import DataFrame
+from torchvision.transforms.v2 import Transform
+
+from anomalib import TaskType
+from anomalib.data.base import AnomalibDataModule, AnomalibDataset
+from anomalib.data.errors import MisMatchError
+from anomalib.data.utils import (
+    DirType,
+    LabelName,
+    Split,
+    TestSplitMode,
+    ValSplitMode,
+)
+from anomalib.data.utils.path import _prepare_files_labels, validate_and_resolve_path
+
+
+def make_folder_dataset(
+    normal_dir: str | Path | Sequence[str | Path],
+    root: str | Path | None = None,
+    abnormal_dir: str | Path | Sequence[str | Path] | None = None,
+    normal_test_dir: str | Path | Sequence[str | Path] | None = None,
+    mask_dir: str | Path | Sequence[str | Path] | None = None,
+    split: str | Split | None = None,
+    extensions: tuple[str, ...] | None = None,
+) -> DataFrame:
+    """Make Folder Dataset.
+
+    Args:
+        normal_dir (str | Path | Sequence): Path to the directory containing normal images.
+        root (str | Path | None): Path to the root directory of the dataset.
+            Defaults to ``None``.
+        abnormal_dir (str | Path | Sequence | None, optional): Path to the directory containing abnormal images.
+            Defaults to ``None``.
+        normal_test_dir (str | Path | Sequence | None, optional): Path to the directory containing normal images for
+            the test dataset. Normal test images will be a split of `normal_dir` if `None`.
+            Defaults to ``None``.
+        mask_dir (str | Path | Sequence | None, optional): Path to the directory containing the mask annotations.
+            Defaults to ``None``.
+        split (str | Split | None, optional): Dataset split (ie., Split.FULL, Split.TRAIN or Split.TEST).
+            Defaults to ``None``.
+        extensions (tuple[str, ...] | None, optional): Type of the image extensions to read from the directory.
+            Defaults to ``None``.
+
+    Returns:
+        DataFrame: an output dataframe containing samples for the requested split (ie., train or test).
+
+    Examples:
+        Assume that we would like to use this ``make_folder_dataset`` to create a dataset from a folder.
+        We could then create the dataset as follows,
+
+        .. code-block:: python
+
+            folder_df = make_folder_dataset(
+                normal_dir=dataset_root / "good",
+                abnormal_dir=dataset_root / "crack",
+                split="train",
+            )
+            folder_df.head()
+
+        .. code-block:: bash
+
+                      image_path           label  label_index mask_path        split
+            0  ./toy/good/00.jpg  DirType.NORMAL            0            Split.TRAIN
+            1  ./toy/good/01.jpg  DirType.NORMAL            0            Split.TRAIN
+            2  ./toy/good/02.jpg  DirType.NORMAL            0            Split.TRAIN
+            3  ./toy/good/03.jpg  DirType.NORMAL            0            Split.TRAIN
+            4  ./toy/good/04.jpg  DirType.NORMAL            0            Split.TRAIN
+    """
+
+    def _resolve_path_and_convert_to_list(path: str | Path | Sequence[str | Path] | None) -> list[Path]:
+        """Convert path to list of paths.
+
+        Args:
+            path (str | Path | Sequence | None): Path to replace with Sequence[str | Path].
+
+        Examples:
+            >>> _resolve_path_and_convert_to_list("dir")
+            [Path("path/to/dir")]
+            >>> _resolve_path_and_convert_to_list(["dir1", "dir2"])
+            [Path("path/to/dir1"), Path("path/to/dir2")]
+
+        Returns:
+            list[Path]: The result of path replaced by Sequence[str | Path].
+        """
+        if isinstance(path, Sequence) and not isinstance(path, str):
+            return [validate_and_resolve_path(dir_path, root) for dir_path in path]
+        return [validate_and_resolve_path(path, root)] if path is not None else []
+
+    # All paths are changed to the List[Path] type and used.
+    normal_dir = _resolve_path_and_convert_to_list(normal_dir)
+    abnormal_dir = _resolve_path_and_convert_to_list(abnormal_dir)
+    normal_test_dir = _resolve_path_and_convert_to_list(normal_test_dir)
+    mask_dir = _resolve_path_and_convert_to_list(mask_dir)
+    if len(normal_dir) == 0:
+        msg = "A folder location must be provided in normal_dir."
+        raise ValueError(msg)
+
+    filenames = []
+    labels = []
+    dirs = {DirType.NORMAL: normal_dir}
+
+    if abnormal_dir:
+        dirs[DirType.ABNORMAL] = abnormal_dir
+
+    if normal_test_dir:
+        dirs[DirType.NORMAL_TEST] = normal_test_dir
+
+    if mask_dir:
+        dirs[DirType.MASK] = mask_dir
+
+    for dir_type, paths in dirs.items():
+        for path in paths:
+            filename, label = _prepare_files_labels(path, dir_type, extensions)
+            filenames += filename
+            labels += label
+
+    samples = DataFrame({"image_path": filenames, "label": labels})
+    samples = samples.sort_values(by="image_path", ignore_index=True)
+
+    # Create label index for normal (0) and abnormal (1) images.
+    samples.loc[
+        (samples.label == DirType.NORMAL) | (samples.label == DirType.NORMAL_TEST),
+        "label_index",
+    ] = LabelName.NORMAL
+    samples.loc[(samples.label == DirType.ABNORMAL), "label_index"] = LabelName.ABNORMAL
+    samples.label_index = samples.label_index.astype("Int64")
+
+    # If a path to mask is provided, add it to the sample dataframe.
+
+    if len(mask_dir) > 0 and len(abnormal_dir) > 0:
+        samples.loc[samples.label == DirType.ABNORMAL, "mask_path"] = samples.loc[
+            samples.label == DirType.MASK
+        ].image_path.to_numpy()
+        samples["mask_path"] = samples["mask_path"].fillna("")
+        samples = samples.astype({"mask_path": "str"})
+
+        # make sure all every rgb image has a corresponding mask image.
+        if not (
+            samples.loc[samples.label_index == LabelName.ABNORMAL]
+            .apply(lambda x: Path(x.image_path).stem in Path(x.mask_path).stem, axis=1)
+            .all()
+        ):
+            msg = """Mismatch between anomalous images and mask images. Make sure the mask files "
+                     "folder follow the same naming convention as the anomalous images in the dataset "
+                     "(e.g. image: '000.png', mask: '000.png')."""
+            raise MisMatchError(msg)
+
+    else:
+        samples["mask_path"] = ""
+
+    # remove all the rows with temporal image samples that have already been assigned
+    samples = samples.loc[
+        (samples.label == DirType.NORMAL) | (samples.label == DirType.ABNORMAL) | (samples.label == DirType.NORMAL_TEST)
+    ]
+
+    # Ensure the pathlib objects are converted to str.
+    # This is because torch dataloader doesn't like pathlib.
+    samples = samples.astype({"image_path": "str"})
+
+    # Create train/test split.
+    # By default, all the normal samples are assigned as train.
+    #   and all the abnormal samples are test.
+    samples.loc[(samples.label == DirType.NORMAL), "split"] = Split.TRAIN
+    samples.loc[(samples.label == DirType.ABNORMAL) | (samples.label == DirType.NORMAL_TEST), "split"] = Split.TEST
+
+    # Get the data frame for the split.
+    if split:
+        samples = samples[samples.split == split]
+        samples = samples.reset_index(drop=True)
+
+    return samples
+
+
+class FolderDataset(AnomalibDataset):
+    """Folder dataset.
+
+    This class is used to create a dataset from a folder. The class utilizes the Torch Dataset class.
+
+    Args:
+        name (str): Name of the dataset. This is used to name the datamodule, especially when logging/saving.
+        task (TaskType): Task type. (``classification``, ``detection`` or ``segmentation``).
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        normal_dir (str | Path | Sequence): Path to the directory containing normal images.
+        root (str | Path | None): Root folder of the dataset.
+            Defaults to ``None``.
+        abnormal_dir (str | Path | Sequence | None, optional): Path to the directory containing abnormal images.
+            Defaults to ``None``.
+        normal_test_dir (str | Path | Sequence | None, optional): Path to the directory containing
+            normal images for the test dataset.
+            Defaults to ``None``.
+        mask_dir (str | Path | Sequence | None, optional): Path to the directory containing
+            the mask annotations.
+            Defaults to ``None``.
+        split (str | Split | None): Fixed subset split that follows from folder structure on file system.
+            Choose from [Split.FULL, Split.TRAIN, Split.TEST]
+            Defaults to ``None``.
+        extensions (tuple[str, ...] | None, optional): Type of the image extensions to read from the directory.
+            Defaults to ``None``.
+
+    Raises:
+        ValueError: When task is set to classification and `mask_dir` is provided. When `mask_dir` is
+            provided, `task` should be set to `segmentation`.
+
+    Examples:
+        Assume that we would like to use this ``FolderDataset`` to create a dataset from a folder for a classification
+        task. We could first create the transforms,
+
+        >>> from anomalib.data.utils import InputNormalizationMethod, get_transforms
+        >>> transform = get_transforms(image_size=256, normalization=InputNormalizationMethod.NONE)
+
+        We could then create the dataset as follows,
+
+        .. code-block:: python
+
+            folder_dataset_classification_train = FolderDataset(
+                normal_dir=dataset_root / "good",
+                abnormal_dir=dataset_root / "crack",
+                split="train",
+                transform=transform,
+                task=TaskType.CLASSIFICATION,
+            )
+
+    """
+
+    def __init__(
+        self,
+        name: str,
+        task: TaskType,
+        normal_dir: str | Path | Sequence[str | Path],
+        transform: Transform | None = None,
+        root: str | Path | None = None,
+        abnormal_dir: str | Path | Sequence[str | Path] | None = None,
+        normal_test_dir: str | Path | Sequence[str | Path] | None = None,
+        mask_dir: str | Path | Sequence[str | Path] | None = None,
+        split: str | Split | None = None,
+        extensions: tuple[str, ...] | None = None,
+    ) -> None:
+        super().__init__(task, transform)
+
+        self._name = name
+        self.split = split
+        self.root = root
+        self.normal_dir = normal_dir
+        self.abnormal_dir = abnormal_dir
+        self.normal_test_dir = normal_test_dir
+        self.mask_dir = mask_dir
+        self.extensions = extensions
+
+        self.samples = make_folder_dataset(
+            root=self.root,
+            normal_dir=self.normal_dir,
+            abnormal_dir=self.abnormal_dir,
+            normal_test_dir=self.normal_test_dir,
+            mask_dir=self.mask_dir,
+            split=self.split,
+            extensions=self.extensions,
+        )
+
+    @property
+    def name(self) -> str:
+        """Name of the dataset.
+
+        Folder dataset overrides the name property to provide a custom name.
+        """
+        return self._name
+
+
+class Folder(AnomalibDataModule):
+    """Folder DataModule.
+
+    Args:
+        name (str): Name of the dataset. This is used to name the datamodule, especially when logging/saving.
+        normal_dir (str | Path | Sequence): Name of the directory containing normal images.
+        root (str | Path | None): Path to the root folder containing normal and abnormal dirs.
+            Defaults to ``None``.
+        abnormal_dir (str | Path | None | Sequence): Name of the directory containing abnormal images.
+            Defaults to ``None``.
+        normal_test_dir (str | Path | Sequence | None, optional): Path to the directory containing
+            normal images for the test dataset.
+            Defaults to ``None``.
+        mask_dir (str | Path | Sequence | None, optional): Path to the directory containing
+            the mask annotations.
+            Defaults to ``None``.
+        normal_split_ratio (float, optional): Ratio to split normal training images and add to the
+            test set in case test set doesn't contain any normal images.
+            Defaults to 0.2.
+        extensions (tuple[str, ...] | None, optional): Type of the image extensions to read from the
+            directory.
+            Defaults to ``None``.
+        train_batch_size (int, optional): Training batch size.
+            Defaults to ``32``.
+        eval_batch_size (int, optional): Validation, test and predict batch size.
+            Defaults to ``32``.
+        num_workers (int, optional): Number of workers.
+            Defaults to ``8``.
+        task (TaskType, optional): Task type. Could be ``classification``, ``detection`` or ``segmentation``.
+            Defaults to ``segmentation``.
+        image_size (tuple[int, int], optional): Size to which input images should be resized.
+            Defaults to ``None``.
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        train_transform (Transform, optional): Transforms that should be applied to the input images during training.
+            Defaults to ``None``.
+        eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation.
+            Defaults to ``None``.
+        test_split_mode (TestSplitMode): Setting that determines how the testing subset is obtained.
+            Defaults to ``TestSplitMode.FROM_DIR``.
+        test_split_ratio (float): Fraction of images from the train set that will be reserved for testing.
+            Defaults to ``0.2``.
+        val_split_mode (ValSplitMode): Setting that determines how the validation subset is obtained.
+            Defaults to ``ValSplitMode.FROM_TEST``.
+        val_split_ratio (float): Fraction of train or test images that will be reserved for validation.
+            Defaults to ``0.5``.
+        seed (int | None, optional): Seed used during random subset splitting.
+            Defaults to ``None``.
+
+    Examples:
+        The following code demonstrates how to use the ``Folder`` datamodule. Assume that the dataset is structured
+        as follows:
+
+        .. code-block:: bash
+
+            $ tree sample_dataset
+            sample_dataset
+            ├── colour
+            │   ├── 00.jpg
+            │   ├── ...
+            │   └── x.jpg
+            ├── crack
+            │   ├── 00.jpg
+            │   ├── ...
+            │   └── y.jpg
+            ├── good
+            │   ├── ...
+            │   └── z.jpg
+            ├── LICENSE
+            └── mask
+                ├── colour
+                │   ├── ...
+                │   └── x.jpg
+                └── crack
+                    ├── ...
+                    └── y.jpg
+
+        .. code-block:: python
+
+            folder_datamodule = Folder(
+                root=dataset_root,
+                normal_dir="good",
+                abnormal_dir="crack",
+                task=TaskType.SEGMENTATION,
+                mask_dir=dataset_root / "mask" / "crack",
+                image_size=256,
+                normalization=InputNormalizationMethod.NONE,
+            )
+            folder_datamodule.setup()
+
+        To access the training images,
+
+        .. code-block:: python
+
+            >> i, data = next(enumerate(folder_datamodule.train_dataloader()))
+            >> print(data.keys(), data["image"].shape)
+
+        To access the test images,
+
+        .. code-block:: python
+
+            >> i, data = next(enumerate(folder_datamodule.test_dataloader()))
+            >> print(data.keys(), data["image"].shape)
+    """
+
+    def __init__(
+        self,
+        name: str,
+        normal_dir: str | Path | Sequence[str | Path],
+        root: str | Path | None = None,
+        abnormal_dir: str | Path | Sequence[str | Path] | None = None,
+        normal_test_dir: str | Path | Sequence[str | Path] | None = None,
+        mask_dir: str | Path | Sequence[str | Path] | None = None,
+        normal_split_ratio: float = 0.2,
+        extensions: tuple[str] | None = None,
+        train_batch_size: int = 32,
+        eval_batch_size: int = 32,
+        num_workers: int = 8,
+        task: TaskType | str = TaskType.SEGMENTATION,
+        image_size: tuple[int, int] | None = None,
+        transform: Transform | None = None,
+        train_transform: Transform | None = None,
+        eval_transform: Transform | None = None,
+        test_split_mode: TestSplitMode | str = TestSplitMode.FROM_DIR,
+        test_split_ratio: float = 0.2,
+        val_split_mode: ValSplitMode | str = ValSplitMode.FROM_TEST,
+        val_split_ratio: float = 0.5,
+        seed: int | None = None,
+    ) -> None:
+        self._name = name
+        self.root = root
+        self.normal_dir = normal_dir
+        self.abnormal_dir = abnormal_dir
+        self.normal_test_dir = normal_test_dir
+        self.mask_dir = mask_dir
+        self.task = TaskType(task)
+        self.extensions = extensions
+        test_split_mode = TestSplitMode(test_split_mode)
+        val_split_mode = ValSplitMode(val_split_mode)
+        super().__init__(
+            train_batch_size=train_batch_size,
+            eval_batch_size=eval_batch_size,
+            num_workers=num_workers,
+            test_split_mode=test_split_mode,
+            test_split_ratio=test_split_ratio,
+            val_split_mode=val_split_mode,
+            val_split_ratio=val_split_ratio,
+            image_size=image_size,
+            transform=transform,
+            train_transform=train_transform,
+            eval_transform=eval_transform,
+            seed=seed,
+        )
+
+        if task == TaskType.SEGMENTATION and test_split_mode == TestSplitMode.FROM_DIR and mask_dir is None:
+            msg = (
+                f"Segmentation task requires mask directory if test_split_mode is {test_split_mode}. "
+                "You could set test_split_mode to {TestSplitMode.NONE} or provide a mask directory."
+            )
+            raise ValueError(
+                msg,
+            )
+
+        self.normal_split_ratio = normal_split_ratio
+
+    def _setup(self, _stage: str | None = None) -> None:
+        self.train_data = FolderDataset(
+            name=self.name,
+            task=self.task,
+            transform=self.train_transform,
+            split=Split.TRAIN,
+            root=self.root,
+            normal_dir=self.normal_dir,
+            abnormal_dir=self.abnormal_dir,
+            normal_test_dir=self.normal_test_dir,
+            mask_dir=self.mask_dir,
+            extensions=self.extensions,
+        )
+
+        self.test_data = FolderDataset(
+            name=self.name,
+            task=self.task,
+            transform=self.eval_transform,
+            split=Split.TEST,
+            root=self.root,
+            normal_dir=self.normal_dir,
+            abnormal_dir=self.abnormal_dir,
+            normal_test_dir=self.normal_test_dir,
+            mask_dir=self.mask_dir,
+            extensions=self.extensions,
+        )
+
+    @property
+    def name(self) -> str:
+        """Name of the datamodule.
+
+        Folder datamodule overrides the name property to provide a custom name.
+        """
+        return self._name

anomalib/data/image/kolektor.py

@@ -0,0 +1,342 @@
+"""Kolektor Surface-Defect Dataset (CC BY-NC-SA 4.0).
+
+Description:
+    This script provides a PyTorch Dataset, DataLoader, and PyTorch Lightning DataModule for the Kolektor
+    Surface-Defect dataset. The dataset can be accessed at `Kolektor Surface-Defect Dataset <https://www.vicos.si/resources/kolektorsdd/>`_.
+
+License:
+    The Kolektor Surface-Defect dataset is released under the Creative Commons Attribution-NonCommercial-ShareAlike
+    4.0 International License (CC BY-NC-SA 4.0). For more details, visit
+    `Creative Commons License <https://creativecommons.org/licenses/by-nc-sa/4.0/>`_.
+
+Reference:
+    Tabernik, Domen, Samo Šela, Jure Skvarč, and Danijel Skočaj. "Segmentation-based deep-learning approach
+    for surface-defect detection." Journal of Intelligent Manufacturing 31, no. 3 (2020): 759-776.
+"""
+
+# Copyright (C) 2023-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from pathlib import Path
+
+import numpy as np
+from cv2 import imread
+from pandas import DataFrame
+from sklearn.model_selection import train_test_split
+from torchvision.transforms.v2 import Transform
+
+from anomalib import TaskType
+from anomalib.data.base import AnomalibDataModule, AnomalibDataset
+from anomalib.data.errors import MisMatchError
+from anomalib.data.utils import (
+    DownloadInfo,
+    Split,
+    TestSplitMode,
+    ValSplitMode,
+    download_and_extract,
+    validate_path,
+)
+
+__all__ = ["Kolektor", "KolektorDataset", "make_kolektor_dataset"]
+
+logger = logging.getLogger(__name__)
+
+DOWNLOAD_INFO = DownloadInfo(
+    name="kolektor",
+    url="https://go.vicos.si/kolektorsdd",
+    hashsum="65dc621693418585de9c4467d1340ea7958a6181816f0dc2883a1e8b61f9d4dc",
+    filename="KolektorSDD.zip",
+)
+
+
+def is_mask_anomalous(path: str) -> int:
+    """Check if a mask shows defects.
+
+    Args:
+        path (str): Path to the mask file.
+
+    Returns:
+        int: 1 if the mask shows defects, 0 otherwise.
+
+    Example:
+        Assume that the following image is a mask for a defective image.
+        Then the function will return 1.
+
+        >>> from anomalib.data.image.kolektor import is_mask_anomalous
+        >>> path = './KolektorSDD/kos01/Part0_label.bmp'
+        >>> is_mask_anomalous(path)
+        1
+    """
+    img_arr = imread(path)
+    if np.all(img_arr == 0):
+        return 0
+    return 1
+
+
+def make_kolektor_dataset(
+    root: str | Path,
+    train_split_ratio: float = 0.8,
+    split: str | Split | None = None,
+) -> DataFrame:
+    """Create Kolektor samples by parsing the Kolektor data file structure.
+
+    The files are expected to follow this structure:
+    - Image files: `path/to/dataset/item/image_filename.jpg`, `path/to/dataset/kos01/Part0.jpg`
+    - Mask files: `path/to/dataset/item/mask_filename.bmp`, `path/to/dataset/kos01/Part0_label.bmp`
+
+    This function creates a DataFrame to store the parsed information in the following format:
+
+    +---+-------------------+--------+-------+---------+-----------------------+------------------------+-------------+
+    |   | path              | item   | split | label   | image_path            | mask_path              | label_index |
+    +---+-------------------+--------+-------+---------+-----------------------+------------------------+-------------+
+    | 0 | KolektorSDD       | kos01  | test  | Bad     | /path/to/image_file   | /path/to/mask_file     | 1           |
+    +---+-------------------+--------+-------+---------+-----------------------+------------------------+-------------+
+
+    Args:
+        root (Path): Path to the dataset.
+        train_split_ratio (float, optional): Ratio for splitting good images into train/test sets.
+            Defaults to ``0.8``.
+        split (str | Split | None, optional): Dataset split (either 'train' or 'test').
+            Defaults to ``None``.
+
+    Returns:
+        pandas.DataFrame: An output DataFrame containing the samples of the dataset.
+
+    Example:
+        The following example shows how to get training samples from the Kolektor Dataset:
+
+        >>> from pathlib import Path
+        >>> root = Path('./KolektorSDD/')
+        >>> samples = create_kolektor_samples(root, train_split_ratio=0.8)
+        >>> samples.head()
+               path       item  split label   image_path                    mask_path                   label_index
+           0   KolektorSDD   kos01  train Good  KolektorSDD/kos01/Part0.jpg  KolektorSDD/kos01/Part0_label.bmp  0
+           1   KolektorSDD   kos01  train Good  KolektorSDD/kos01/Part1.jpg  KolektorSDD/kos01/Part1_label.bmp  0
+           2   KolektorSDD   kos01  train Good  KolektorSDD/kos01/Part2.jpg  KolektorSDD/kos01/Part2_label.bmp  0
+           3   KolektorSDD   kos01  test  Good  KolektorSDD/kos01/Part3.jpg  KolektorSDD/kos01/Part3_label.bmp  0
+           4   KolektorSDD   kos01  train Good  KolektorSDD/kos01/Part4.jpg  KolektorSDD/kos01/Part4_label.bmp  0
+    """
+    root = validate_path(root)
+
+    # Get list of images and masks
+    samples_list = [(str(root),) + f.parts[-2:] for f in root.glob(r"**/*") if f.suffix == ".jpg"]
+    masks_list = [(str(root),) + f.parts[-2:] for f in root.glob(r"**/*") if f.suffix == ".bmp"]
+
+    if not samples_list:
+        msg = f"Found 0 images in {root}"
+        raise RuntimeError(msg)
+
+    # Create dataframes
+    samples = DataFrame(samples_list, columns=["path", "item", "image_path"])
+    masks = DataFrame(masks_list, columns=["path", "item", "image_path"])
+
+    # Modify image_path column by converting to absolute path
+    samples["image_path"] = samples.path + "/" + samples.item + "/" + samples.image_path
+    masks["image_path"] = masks.path + "/" + masks.item + "/" + masks.image_path
+
+    # Sort samples by image path
+    samples = samples.sort_values(by="image_path", ignore_index=True)
+    masks = masks.sort_values(by="image_path", ignore_index=True)
+
+    # Add mask paths for sample images
+    samples["mask_path"] = masks.image_path.to_numpy()
+
+    # Use is_good func to configure the label_index
+    samples["label_index"] = samples["mask_path"].apply(is_mask_anomalous)
+    samples.label_index = samples.label_index.astype(int)
+
+    # Use label indexes to label data
+    samples.loc[(samples.label_index == 0), "label"] = "Good"
+    samples.loc[(samples.label_index == 1), "label"] = "Bad"
+
+    # Add all 'Bad' samples to test set
+    samples.loc[(samples.label == "Bad"), "split"] = "test"
+
+    # Divide 'good' images to train/test on 0.8/0.2 ratio
+    train_samples, test_samples = train_test_split(
+        samples[samples.label == "Good"],
+        train_size=train_split_ratio,
+        random_state=42,
+    )
+    samples.loc[train_samples.index, "split"] = "train"
+    samples.loc[test_samples.index, "split"] = "test"
+
+    # Reorder columns
+    samples = samples[["path", "item", "split", "label", "image_path", "mask_path", "label_index"]]
+
+    # assert that the right mask files are associated with the right test images
+    if not (
+        samples.loc[samples.label_index == 1]
+        .apply(lambda x: Path(x.image_path).stem in Path(x.mask_path).stem, axis=1)
+        .all()
+    ):
+        msg = """Mismatch between anomalous images and ground truth masks. Make sure the mask files
+        follow the same naming convention as the anomalous images in the dataset
+        (e.g. image: 'Part0.jpg', mask: 'Part0_label.bmp')."""
+        raise MisMatchError(msg)
+
+    # Get the dataframe for the required split
+    if split:
+        samples = samples[samples.split == split].reset_index(drop=True)
+
+    return samples
+
+
+class KolektorDataset(AnomalibDataset):
+    """Kolektor dataset class.
+
+    Args:
+        task (TaskType): Task type, ``classification``, ``detection`` or ``segmentation``
+        root (Path | str): Path to the root of the dataset
+            Defaults to ``./datasets/kolektor``.
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        split (str | Split | None): Split of the dataset, usually Split.TRAIN or Split.TEST
+            Defaults to ``None``.
+    """
+
+    def __init__(
+        self,
+        task: TaskType,
+        root: Path | str = "./datasets/kolektor",
+        transform: Transform | None = None,
+        split: str | Split | None = None,
+    ) -> None:
+        super().__init__(task=task, transform=transform)
+
+        self.root = root
+        self.split = split
+        self.samples = make_kolektor_dataset(self.root, train_split_ratio=0.8, split=self.split)
+
+
+class Kolektor(AnomalibDataModule):
+    """Kolektor Datamodule.
+
+    Args:
+        root (Path | str): Path to the root of the dataset
+        train_batch_size (int, optional): Training batch size.
+            Defaults to ``32``.
+        eval_batch_size (int, optional): Test batch size.
+            Defaults to ``32``.
+        num_workers (int, optional): Number of workers.
+            Defaults to ``8``.
+        task TaskType): Task type, 'classification', 'detection' or 'segmentation'
+            Defaults to ``TaskType.SEGMENTATION``.
+        image_size (tuple[int, int], optional): Size to which input images should be resized.
+            Defaults to ``None``.
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        train_transform (Transform, optional): Transforms that should be applied to the input images during training.
+            Defaults to ``None``.
+        eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation.
+            Defaults to ``None``.
+        test_split_mode (TestSplitMode): Setting that determines how the testing subset is obtained.
+            Defaults to ``TestSplitMode.FROM_DIR``
+        test_split_ratio (float): Fraction of images from the train set that will be reserved for testing.
+            Defaults to ``0.2``
+        val_split_mode (ValSplitMode): Setting that determines how the validation subset is obtained.
+            Defaults to ``ValSplitMode.SAME_AS_TEST``
+        val_split_ratio (float): Fraction of train or test images that will be reserved for validation.
+            Defaults to ``0.5``
+        seed (int | None, optional): Seed which may be set to a fixed value for reproducibility.
+            Defaults to ``None``.
+    """
+
+    def __init__(
+        self,
+        root: Path | str = "./datasets/kolektor",
+        train_batch_size: int = 32,
+        eval_batch_size: int = 32,
+        num_workers: int = 8,
+        task: TaskType | str = TaskType.SEGMENTATION,
+        image_size: tuple[int, int] | None = None,
+        transform: Transform | None = None,
+        train_transform: Transform | None = None,
+        eval_transform: Transform | None = None,
+        test_split_mode: TestSplitMode | str = TestSplitMode.FROM_DIR,
+        test_split_ratio: float = 0.2,
+        val_split_mode: ValSplitMode | str = ValSplitMode.SAME_AS_TEST,
+        val_split_ratio: float = 0.5,
+        seed: int | None = None,
+    ) -> None:
+        super().__init__(
+            train_batch_size=train_batch_size,
+            eval_batch_size=eval_batch_size,
+            num_workers=num_workers,
+            image_size=image_size,
+            transform=transform,
+            train_transform=train_transform,
+            eval_transform=eval_transform,
+            test_split_mode=test_split_mode,
+            test_split_ratio=test_split_ratio,
+            val_split_mode=val_split_mode,
+            val_split_ratio=val_split_ratio,
+            seed=seed,
+        )
+
+        self.task = TaskType(task)
+        self.root = Path(root)
+
+    def _setup(self, _stage: str | None = None) -> None:
+        self.train_data = KolektorDataset(
+            task=self.task,
+            transform=self.train_transform,
+            split=Split.TRAIN,
+            root=self.root,
+        )
+        self.test_data = KolektorDataset(
+            task=self.task,
+            transform=self.eval_transform,
+            split=Split.TEST,
+            root=self.root,
+        )
+
+    def prepare_data(self) -> None:
+        """Download the dataset if not available.
+
+        This method checks if the specified dataset is available in the file system.
+        If not, it downloads and extracts the dataset into the appropriate directory.
+
+        Example:
+            Assume the dataset is not available on the file system.
+            Here's how the directory structure looks before and after calling the
+            `prepare_data` method:
+
+            Before:
+
+            .. code-block:: bash
+
+                $ tree datasets
+                datasets
+                ├── dataset1
+                └── dataset2
+
+            Calling the method:
+
+            .. code-block:: python
+
+                >> datamodule = Kolektor(root="./datasets/kolektor")
+                >> datamodule.prepare_data()
+
+            After:
+
+            .. code-block:: bash
+
+                $ tree datasets
+                datasets
+                ├── dataset1
+                ├── dataset2
+                └── kolektor
+                    ├── kolektorsdd
+                    ├── kos01
+                    ├── ...
+                    └── kos50
+                        ├── Part0.jpg
+                        ├── Part0_label.bmp
+                        └── ...
+        """
+        if (self.root).is_dir():
+            logger.info("Found the dataset.")
+        else:
+            download_and_extract(self.root, DOWNLOAD_INFO)

anomalib/data/image/mvtec.py

@@ -0,0 +1,414 @@
+"""MVTec AD Dataset (CC BY-NC-SA 4.0).
+
+Description:
+    This script contains PyTorch Dataset, Dataloader and PyTorch Lightning
+    DataModule for the MVTec AD dataset. If the dataset is not on the file system,
+    the script downloads and extracts the dataset and create PyTorch data objects.
+
+License:
+    MVTec AD dataset is released under the Creative Commons
+    Attribution-NonCommercial-ShareAlike 4.0 International License
+    (CC BY-NC-SA 4.0)(https://creativecommons.org/licenses/by-nc-sa/4.0/).
+
+References:
+    - Paul Bergmann, Kilian Batzner, Michael Fauser, David Sattlegger, Carsten Steger:
+      The MVTec Anomaly Detection Dataset: A Comprehensive Real-World Dataset for
+      Unsupervised Anomaly Detection; in: International Journal of Computer Vision
+      129(4):1038-1059, 2021, DOI: 10.1007/s11263-020-01400-4.
+
+    - Paul Bergmann, Michael Fauser, David Sattlegger, Carsten Steger: MVTec AD —
+      A Comprehensive Real-World Dataset for Unsupervised Anomaly Detection;
+      in: IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR),
+      9584-9592, 2019, DOI: 10.1109/CVPR.2019.00982.
+"""
+
+# Copyright (C) 2022-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from collections.abc import Sequence
+from pathlib import Path
+
+from pandas import DataFrame
+from torchvision.transforms.v2 import Transform
+
+from anomalib import TaskType
+from anomalib.data.base import AnomalibDataModule, AnomalibDataset
+from anomalib.data.errors import MisMatchError
+from anomalib.data.utils import (
+    DownloadInfo,
+    LabelName,
+    Split,
+    TestSplitMode,
+    ValSplitMode,
+    download_and_extract,
+    validate_path,
+)
+
+logger = logging.getLogger(__name__)
+
+
+IMG_EXTENSIONS = (".png", ".PNG")
+
+DOWNLOAD_INFO = DownloadInfo(
+    name="mvtec",
+    url="https://www.mydrive.ch/shares/38536/3830184030e49fe74747669442f0f282/download/420938113-1629952094"
+    "/mvtec_anomaly_detection.tar.xz",
+    hashsum="cf4313b13603bec67abb49ca959488f7eedce2a9f7795ec54446c649ac98cd3d",
+)
+
+CATEGORIES = (
+    "bottle",
+    "cable",
+    "capsule",
+    "carpet",
+    "grid",
+    "hazelnut",
+    "leather",
+    "metal_nut",
+    "pill",
+    "screw",
+    "tile",
+    "toothbrush",
+    "transistor",
+    "wood",
+    "zipper",
+)
+
+
+def make_mvtec_dataset(
+    root: str | Path,
+    split: str | Split | None = None,
+    extensions: Sequence[str] | None = None,
+) -> DataFrame:
+    """Create MVTec AD samples by parsing the MVTec AD data file structure.
+
+    The files are expected to follow the structure:
+        path/to/dataset/split/category/image_filename.png
+        path/to/dataset/ground_truth/category/mask_filename.png
+
+    This function creates a dataframe to store the parsed information based on the following format:
+
+    +---+---------------+-------+---------+---------------+---------------------------------------+-------------+
+    |   | path          | split | label   | image_path    | mask_path                             | label_index |
+    +===+===============+=======+=========+===============+=======================================+=============+
+    | 0 | datasets/name | test  | defect  | filename.png  | ground_truth/defect/filename_mask.png | 1           |
+    +---+---------------+-------+---------+---------------+---------------------------------------+-------------+
+
+    Args:
+        root (Path): Path to dataset
+        split (str | Split | None, optional): Dataset split (ie., either train or test).
+            Defaults to ``None``.
+        extensions (Sequence[str] | None, optional): List of file extensions to be included in the dataset.
+            Defaults to ``None``.
+
+    Examples:
+        The following example shows how to get training samples from MVTec AD bottle category:
+
+        >>> root = Path('./MVTec')
+        >>> category = 'bottle'
+        >>> path = root / category
+        >>> path
+        PosixPath('MVTec/bottle')
+
+        >>> samples = make_mvtec_dataset(path, split='train', split_ratio=0.1, seed=0)
+        >>> samples.head()
+           path         split label image_path                           mask_path                   label_index
+        0  MVTec/bottle train good MVTec/bottle/train/good/105.png MVTec/bottle/ground_truth/good/105_mask.png 0
+        1  MVTec/bottle train good MVTec/bottle/train/good/017.png MVTec/bottle/ground_truth/good/017_mask.png 0
+        2  MVTec/bottle train good MVTec/bottle/train/good/137.png MVTec/bottle/ground_truth/good/137_mask.png 0
+        3  MVTec/bottle train good MVTec/bottle/train/good/152.png MVTec/bottle/ground_truth/good/152_mask.png 0
+        4  MVTec/bottle train good MVTec/bottle/train/good/109.png MVTec/bottle/ground_truth/good/109_mask.png 0
+
+    Returns:
+        DataFrame: an output dataframe containing the samples of the dataset.
+    """
+    if extensions is None:
+        extensions = IMG_EXTENSIONS
+
+    root = validate_path(root)
+    samples_list = [(str(root),) + f.parts[-3:] for f in root.glob(r"**/*") if f.suffix in extensions]
+    if not samples_list:
+        msg = f"Found 0 images in {root}"
+        raise RuntimeError(msg)
+
+    samples = DataFrame(samples_list, columns=["path", "split", "label", "image_path"])
+
+    # Modify image_path column by converting to absolute path
+    samples["image_path"] = samples.path + "/" + samples.split + "/" + samples.label + "/" + samples.image_path
+
+    # Create label index for normal (0) and anomalous (1) images.
+    samples.loc[(samples.label == "good"), "label_index"] = LabelName.NORMAL
+    samples.loc[(samples.label != "good"), "label_index"] = LabelName.ABNORMAL
+    samples.label_index = samples.label_index.astype(int)
+
+    # separate masks from samples
+    mask_samples = samples.loc[samples.split == "ground_truth"].sort_values(by="image_path", ignore_index=True)
+    samples = samples[samples.split != "ground_truth"].sort_values(by="image_path", ignore_index=True)
+
+    # assign mask paths to anomalous test images
+    samples["mask_path"] = ""
+    samples.loc[
+        (samples.split == "test") & (samples.label_index == LabelName.ABNORMAL),
+        "mask_path",
+    ] = mask_samples.image_path.to_numpy()
+
+    # assert that the right mask files are associated with the right test images
+    abnormal_samples = samples.loc[samples.label_index == LabelName.ABNORMAL]
+    if (
+        len(abnormal_samples)
+        and not abnormal_samples.apply(lambda x: Path(x.image_path).stem in Path(x.mask_path).stem, axis=1).all()
+    ):
+        msg = """Mismatch between anomalous images and ground truth masks. Make sure t
+        he mask files in 'ground_truth' folder follow the same naming convention as the
+        anomalous images in the dataset (e.g. image: '000.png', mask: '000.png' or '000_mask.png')."""
+        raise MisMatchError(msg)
+
+    if split:
+        samples = samples[samples.split == split].reset_index(drop=True)
+
+    return samples
+
+
+class MVTecDataset(AnomalibDataset):
+    """MVTec dataset class.
+
+    Args:
+        task (TaskType): Task type, ``classification``, ``detection`` or ``segmentation``.
+        root (Path | str): Path to the root of the dataset.
+            Defaults to ``./datasets/MVTec``.
+        category (str): Sub-category of the dataset, e.g. 'bottle'
+            Defaults to ``bottle``.
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        split (str | Split | None): Split of the dataset, usually Split.TRAIN or Split.TEST
+            Defaults to ``None``.
+
+    Examples:
+        .. code-block:: python
+
+            from anomalib.data.image.mvtec import MVTecDataset
+            from anomalib.data.utils.transforms import get_transforms
+
+            transform = get_transforms(image_size=256)
+            dataset = MVTecDataset(
+                task="classification",
+                transform=transform,
+                root='./datasets/MVTec',
+                category='zipper',
+            )
+            dataset.setup()
+            print(dataset[0].keys())
+            # Output: dict_keys(['image_path', 'label', 'image'])
+
+        When the task is segmentation, the dataset will also contain the mask:
+
+        .. code-block:: python
+
+            dataset.task = "segmentation"
+            dataset.setup()
+            print(dataset[0].keys())
+            # Output: dict_keys(['image_path', 'label', 'image', 'mask_path', 'mask'])
+
+        The image is a torch tensor of shape (C, H, W) and the mask is a torch tensor of shape (H, W).
+
+        .. code-block:: python
+
+            print(dataset[0]["image"].shape, dataset[0]["mask"].shape)
+            # Output: (torch.Size([3, 256, 256]), torch.Size([256, 256]))
+
+    """
+
+    def __init__(
+        self,
+        task: TaskType,
+        root: Path | str = "./datasets/MVTec",
+        category: str = "bottle",
+        transform: Transform | None = None,
+        split: str | Split | None = None,
+    ) -> None:
+        super().__init__(task=task, transform=transform)
+
+        self.root_category = Path(root) / Path(category)
+        self.category = category
+        self.split = split
+        self.samples = make_mvtec_dataset(self.root_category, split=self.split, extensions=IMG_EXTENSIONS)
+
+
+class MVTec(AnomalibDataModule):
+    """MVTec Datamodule.
+
+    Args:
+        root (Path | str): Path to the root of the dataset.
+            Defaults to ``"./datasets/MVTec"``.
+        category (str): Category of the MVTec dataset (e.g. "bottle" or "cable").
+            Defaults to ``"bottle"``.
+        train_batch_size (int, optional): Training batch size.
+            Defaults to ``32``.
+        eval_batch_size (int, optional): Test batch size.
+            Defaults to ``32``.
+        num_workers (int, optional): Number of workers.
+            Defaults to ``8``.
+        task TaskType): Task type, 'classification', 'detection' or 'segmentation'
+            Defaults to ``TaskType.SEGMENTATION``.
+        image_size (tuple[int, int], optional): Size to which input images should be resized.
+            Defaults to ``None``.
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        train_transform (Transform, optional): Transforms that should be applied to the input images during training.
+            Defaults to ``None``.
+        eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation.
+            Defaults to ``None``.
+        test_split_mode (TestSplitMode): Setting that determines how the testing subset is obtained.
+            Defaults to ``TestSplitMode.FROM_DIR``.
+        test_split_ratio (float): Fraction of images from the train set that will be reserved for testing.
+            Defaults to ``0.2``.
+        val_split_mode (ValSplitMode): Setting that determines how the validation subset is obtained.
+            Defaults to ``ValSplitMode.SAME_AS_TEST``.
+        val_split_ratio (float): Fraction of train or test images that will be reserved for validation.
+            Defaults to ``0.5``.
+        seed (int | None, optional): Seed which may be set to a fixed value for reproducibility.
+            Defualts to ``None``.
+
+    Examples:
+        To create an MVTec AD datamodule with default settings:
+
+        >>> datamodule = MVTec()
+        >>> datamodule.setup()
+        >>> i, data = next(enumerate(datamodule.train_dataloader()))
+        >>> data.keys()
+        dict_keys(['image_path', 'label', 'image', 'mask_path', 'mask'])
+
+        >>> data["image"].shape
+        torch.Size([32, 3, 256, 256])
+
+        To change the category of the dataset:
+
+        >>> datamodule = MVTec(category="cable")
+
+        To change the image and batch size:
+
+        >>> datamodule = MVTec(image_size=(512, 512), train_batch_size=16, eval_batch_size=8)
+
+        MVTec AD dataset does not provide a validation set. If you would like
+        to use a separate validation set, you can use the ``val_split_mode`` and
+        ``val_split_ratio`` arguments to create a validation set.
+
+        >>> datamodule = MVTec(val_split_mode=ValSplitMode.FROM_TEST, val_split_ratio=0.1)
+
+        This will subsample the test set by 10% and use it as the validation set.
+        If you would like to create a validation set synthetically that would
+        not change the test set, you can use the ``ValSplitMode.SYNTHETIC`` option.
+
+        >>> datamodule = MVTec(val_split_mode=ValSplitMode.SYNTHETIC, val_split_ratio=0.2)
+
+    """
+
+    def __init__(
+        self,
+        root: Path | str = "./datasets/MVTec",
+        category: str = "bottle",
+        train_batch_size: int = 32,
+        eval_batch_size: int = 32,
+        num_workers: int = 8,
+        task: TaskType | str = TaskType.SEGMENTATION,
+        image_size: tuple[int, int] | None = None,
+        transform: Transform | None = None,
+        train_transform: Transform | None = None,
+        eval_transform: Transform | None = None,
+        test_split_mode: TestSplitMode | str = TestSplitMode.FROM_DIR,
+        test_split_ratio: float = 0.2,
+        val_split_mode: ValSplitMode | str = ValSplitMode.SAME_AS_TEST,
+        val_split_ratio: float = 0.5,
+        seed: int | None = None,
+    ) -> None:
+        super().__init__(
+            train_batch_size=train_batch_size,
+            eval_batch_size=eval_batch_size,
+            image_size=image_size,
+            transform=transform,
+            train_transform=train_transform,
+            eval_transform=eval_transform,
+            num_workers=num_workers,
+            test_split_mode=test_split_mode,
+            test_split_ratio=test_split_ratio,
+            val_split_mode=val_split_mode,
+            val_split_ratio=val_split_ratio,
+            seed=seed,
+        )
+
+        self.task = TaskType(task)
+        self.root = Path(root)
+        self.category = category
+
+    def _setup(self, _stage: str | None = None) -> None:
+        """Set up the datasets and perform dynamic subset splitting.
+
+        This method may be overridden in subclass for custom splitting behaviour.
+
+        Note:
+            The stage argument is not used here. This is because, for a given instance of an AnomalibDataModule
+            subclass, all three subsets are created at the first call of setup(). This is to accommodate the subset
+            splitting behaviour of anomaly tasks, where the validation set is usually extracted from the test set, and
+            the test set must therefore be created as early as the `fit` stage.
+
+        """
+        self.train_data = MVTecDataset(
+            task=self.task,
+            transform=self.train_transform,
+            split=Split.TRAIN,
+            root=self.root,
+            category=self.category,
+        )
+        self.test_data = MVTecDataset(
+            task=self.task,
+            transform=self.eval_transform,
+            split=Split.TEST,
+            root=self.root,
+            category=self.category,
+        )
+
+    def prepare_data(self) -> None:
+        """Download the dataset if not available.
+
+        This method checks if the specified dataset is available in the file system.
+        If not, it downloads and extracts the dataset into the appropriate directory.
+
+        Example:
+            Assume the dataset is not available on the file system.
+            Here's how the directory structure looks before and after calling the
+            `prepare_data` method:
+
+            Before:
+
+            .. code-block:: bash
+
+                $ tree datasets
+                datasets
+                ├── dataset1
+                └── dataset2
+
+            Calling the method:
+
+            .. code-block:: python
+
+                >> datamodule = MVTec(root="./datasets/MVTec", category="bottle")
+                >> datamodule.prepare_data()
+
+            After:
+
+            .. code-block:: bash
+
+                $ tree datasets
+                datasets
+                ├── dataset1
+                ├── dataset2
+                └── MVTec
+                    ├── bottle
+                    ├── ...
+                    └── zipper
+        """
+        if (self.root / self.category).is_dir():
+            logger.info("Found the dataset.")
+        else:
+            download_and_extract(self.root, DOWNLOAD_INFO)

anomalib/data/image/mvtec_loco.py

@@ -0,0 +1,480 @@
+"""MVTec LOCO AD Dataset (CC BY-NC-SA 4.0).
+
+Description:
+    This script contains PyTorch Dataset, Dataloader and PyTorch Lightning
+    DataModule for the MVTec LOCO AD dataset. If the dataset is not on the file system,
+    the script downloads and extracts the dataset and create PyTorch data objects.
+
+License:
+    MVTec LOCO AD dataset is released under the Creative Commons
+    Attribution-NonCommercial-ShareAlike 4.0 International License
+    (CC BY-NC-SA 4.0)(https://creativecommons.org/licenses/by-nc-sa/4.0/).
+
+References:
+    - Paul Bergmann, Kilian Batzner, Michael Fauser, David Sattlegger, and Carsten Steger:
+      Beyond Dents and Scratches: Logical Constraints in Unsupervised Anomaly Detection and Localization;
+      in: International Journal of Computer Vision (IJCV) 130, 947-969, 2022, DOI: 10.1007/s11263-022-01578-9
+"""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from collections.abc import Sequence
+from pathlib import Path
+
+import torch
+from pandas import DataFrame
+from PIL import Image
+from torchvision.transforms.v2 import Transform
+from torchvision.transforms.v2.functional import to_image
+from torchvision.tv_tensors import Mask
+
+from anomalib import TaskType
+from anomalib.data.base import AnomalibDataModule, AnomalibDataset
+from anomalib.data.utils import (
+    DownloadInfo,
+    LabelName,
+    Split,
+    TestSplitMode,
+    ValSplitMode,
+    download_and_extract,
+    masks_to_boxes,
+    read_image,
+    validate_path,
+)
+
+logger = logging.getLogger(__name__)
+
+
+IMG_EXTENSIONS = (".png", ".PNG")
+
+DOWNLOAD_INFO = DownloadInfo(
+    name="mvtec_loco",
+    url="https://www.mydrive.ch/shares/48237/1b9106ccdfbb09a0c414bd49fe44a14a/download/430647091-1646842701"
+    "/mvtec_loco_anomaly_detection.tar.xz",
+    hashsum="9e7c84dba550fd2e59d8e9e231c929c45ba737b6b6a6d3814100f54d63aae687",
+)
+
+CATEGORIES = (
+    "breakfast_box",
+    "juice_bottle",
+    "pushpins",
+    "screw_bag",
+    "splicing_connectors",
+)
+
+
+def make_mvtec_loco_dataset(
+    root: str | Path,
+    split: str | Split | None = None,
+    extensions: Sequence[str] = IMG_EXTENSIONS,
+) -> DataFrame:
+    """Create MVTec LOCO AD samples by parsing the original MVTec LOCO AD data file structure.
+
+    The files are expected to follow the structure:
+        path/to/dataset/split/category/image_filename.png
+        path/to/dataset/ground_truth/category/image_filename/000.png
+
+    where there can be multiple ground-truth masks for the corresponding anomalous images.
+
+    This function creates a dataframe to store the parsed information based on the following format:
+
+    +---+---------------+-------+---------+-------------------------+-----------------------------+-------------+
+    |   | path          | split | label   | image_path              | mask_path                  | label_index |
+    +===+===============+=======+=========+===============+=======================================+=============+
+    | 0 | datasets/name | test  | defect  | path/to/image/file.png  | [path/to/masks/file.png]    | 1           |
+    +---+---------------+-------+---------+-------------------------+-----------------------------+-------------+
+
+    Args:
+        root (str | Path): Path to dataset
+        split (str | Split | None): Dataset split (ie., either train or test).
+            Defaults to ``None``.
+        extensions (Sequence[str]): List of file extensions to be included in the dataset.
+            Defaults to ``None``.
+
+    Returns:
+        DataFrame: an output dataframe containing the samples of the dataset.
+
+    Examples:
+        The following example shows how to get test samples from MVTec LOCO AD pushpins category:
+
+        >>> root = Path('./MVTec_LOCO')
+        >>> category = 'pushpins'
+        >>> path = root / category
+        >>> samples = make_mvtec_loco_dataset(path, split='test')
+    """
+    root = validate_path(root)
+
+    # Retrieve the image and mask files
+    samples_list = []
+    for f in root.glob("**/*"):
+        if f.suffix in extensions:
+            parts = f.parts
+            # 'ground_truth' and non 'ground_truth' path have a different structure
+            if "ground_truth" not in parts:
+                split_folder, label_folder, image_file = parts[-3:]
+                image_path = f"{root}/{split_folder}/{label_folder}/{image_file}"
+                samples_list.append((str(root), split_folder, label_folder, "", image_path))
+            else:
+                split_folder, label_folder, image_folder, image_file = parts[-4:]
+                image_path = f"{root}/{split_folder}/{label_folder}/{image_folder}/{image_file}"
+                samples_list.append((str(root), split_folder, label_folder, image_folder, image_path))
+
+    if not samples_list:
+        msg = f"Found 0 images in {root}"
+        raise RuntimeError(msg)
+
+    samples = DataFrame(samples_list, columns=["path", "split", "label", "image_folder", "image_path"])
+
+    # Replace validation to Split.VAL.value in the split column
+    samples["split"] = samples["split"].replace("validation", Split.VAL.value)
+
+    # Create label index for normal (0) and anomalous (1) images.
+    samples.loc[(samples.label == "good"), "label_index"] = LabelName.NORMAL
+    samples.loc[(samples.label != "good"), "label_index"] = LabelName.ABNORMAL
+    samples.label_index = samples.label_index.astype(int)
+
+    # separate ground-truth masks from samples
+    mask_samples = samples.loc[samples.split == "ground_truth"].sort_values(by="image_path", ignore_index=True)
+    samples = samples[samples.split != "ground_truth"].sort_values(by="image_path", ignore_index=True)
+
+    # Group masks and aggregate the path into a list
+    mask_samples = (
+        mask_samples.groupby(["path", "split", "label", "image_folder"])["image_path"]
+        .agg(list)
+        .reset_index()
+        .rename(columns={"image_path": "mask_path"})
+    )
+
+    # assign mask paths to anomalous test images
+    samples["mask_path"] = ""
+    samples.loc[
+        (samples.split == "test") & (samples.label_index == LabelName.ABNORMAL),
+        "mask_path",
+    ] = mask_samples.mask_path.to_numpy()
+
+    # validate that the right mask files are associated with the right test images
+    if len(samples.loc[samples.label_index == LabelName.ABNORMAL]):
+        image_stems = samples.loc[samples.label_index == LabelName.ABNORMAL]["image_path"].apply(lambda x: Path(x).stem)
+        mask_parent_stems = samples.loc[samples.label_index == LabelName.ABNORMAL]["mask_path"].apply(
+            lambda x: {Path(mask_path).parent.stem for mask_path in x},
+        )
+
+        if not all(
+            next(iter(mask_stems)) == image_stem
+            for image_stem, mask_stems in zip(image_stems, mask_parent_stems, strict=True)
+        ):
+            error_message = (
+                "Mismatch between anomalous images and ground truth masks. "
+                "Make sure the parent folder of the mask files in 'ground_truth' folder "
+                "follows the same naming convention as the anomalous images in the dataset "
+                "(e.g., image: '005.png', mask: '005/000.png')."
+            )
+            raise ValueError(error_message)
+
+    if split:
+        samples = samples[samples.split == split].reset_index(drop=True)
+
+    return samples
+
+
+class MVTecLocoDataset(AnomalibDataset):
+    """MVTec LOCO dataset class.
+
+    Args:
+        task (TaskType): Task type, ``classification``, ``detection`` or ``segmentation``.
+        root (Path | str): Path to the root of the dataset.
+            Defaults to ``./datasets/MVTec_LOCO``.
+        category (str): Sub-category of the dataset, e.g. 'breakfast_box'
+            Defaults to ``breakfast_box``.
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        split (str | Split | None): Split of the dataset, Split.TRAIN, Split.VAL, or Split.TEST
+            Defaults to ``None``.
+
+    Examples:
+        .. code-block:: python
+
+            from anomalib.data.image.mvtec_loco import MVTecLocoDataset
+            from anomalib.data.utils.transforms import get_transforms
+            from torchvision.transforms.v2 import Resize
+
+            transform = Resize((256, 256))
+            dataset = MVTecLocoDataset(
+                task="classification",
+                transform=transform,
+                root='./datasets/MVTec_LOCO',
+                category='breakfast_box',
+            )
+            dataset.setup()
+            print(dataset[0].keys())
+            # Output: dict_keys(['image_path', 'label', 'image'])
+
+        When the task is segmentation, the dataset will also contain the mask:
+
+        .. code-block:: python
+
+            dataset.task = "segmentation"
+            dataset.setup()
+            print(dataset[0].keys())
+            # Output: dict_keys(['image_path', 'label', 'image', 'mask_path', 'mask'])
+
+        The image is a torch tensor of shape (C, H, W) and the mask is a torch tensor of shape (H, W).
+
+        .. code-block:: python
+
+            print(dataset[0]["image"].shape, dataset[0]["mask"].shape)
+            # Output: (torch.Size([3, 256, 256]), torch.Size([256, 256]))
+    """
+
+    def __init__(
+        self,
+        task: TaskType,
+        root: Path | str = "./datasets/MVTec_LOCO",
+        category: str = "breakfast_box",
+        transform: Transform | None = None,
+        split: str | Split | None = None,
+    ) -> None:
+        super().__init__(task=task, transform=transform)
+
+        self.root_category = Path(root) / category
+        self.split = split
+        self.samples = make_mvtec_loco_dataset(
+            self.root_category,
+            split=self.split,
+            extensions=IMG_EXTENSIONS,
+        )
+
+    @staticmethod
+    def _read_mask(mask_path: str | Path) -> Mask:
+        image = Image.open(mask_path).convert("L")
+        return Mask(to_image(image).squeeze(), dtype=torch.uint8)
+
+    def __getitem__(self, index: int) -> dict[str, str | torch.Tensor]:
+        """Get dataset item for the index ``index``.
+
+        This method is mostly based on the super class implementation, with some different as follows:
+            - Using 'torch.where' to make sure the 'mask' in the return item is binarized
+            - An additional 'masks' is added, the non-binary masks with original size for the SPRO metric calculation
+        Args:
+            index (int): Index to get the item.
+
+        Returns:
+            dict[str, str | torch.Tensor]: Dict of image tensor during training. Otherwise, Dict containing image path,
+                target path, image tensor, label and transformed bounding box.
+        """
+        image_path = self.samples.iloc[index].image_path
+        mask_path = self.samples.iloc[index].mask_path
+        label_index = self.samples.iloc[index].label_index
+
+        image = read_image(image_path, as_tensor=True)
+        item = {"image_path": image_path, "label": label_index}
+
+        if self.task == TaskType.CLASSIFICATION:
+            item["image"] = self.transform(image) if self.transform else image
+        elif self.task in (TaskType.DETECTION, TaskType.SEGMENTATION):
+            # Only Anomalous (1) images have masks in anomaly datasets
+            # Therefore, create empty mask for Normal (0) images.
+            if isinstance(mask_path, str):
+                mask_path = [mask_path]
+            semantic_mask = (
+                Mask(torch.zeros(image.shape[-2:])).to(torch.uint8)
+                if label_index == LabelName.NORMAL
+                else Mask(torch.stack([self._read_mask(path) for path in mask_path]))
+            )
+
+            binary_mask = Mask(semantic_mask.view(-1, *semantic_mask.shape[-2:]).int().any(dim=0).to(torch.uint8))
+            item["image"], item["mask"] = self.transform(image, binary_mask) if self.transform else (image, binary_mask)
+
+            item["mask_path"] = mask_path
+            # List of masks with the original size for saturation based metrics calculation
+            item["semantic_mask"] = semantic_mask
+
+            if self.task == TaskType.DETECTION:
+                # create boxes from masks for detection task
+                boxes, _ = masks_to_boxes(item["mask"])
+                item["boxes"] = boxes[0]
+        else:
+            msg = f"Unknown task type: {self.task}"
+            raise ValueError(msg)
+
+        return item
+
+
+class MVTecLoco(AnomalibDataModule):
+    """MVTec LOCO Datamodule.
+
+    Args:
+        root (Path | str): Path to the root of the dataset.
+            Defaults to ``"./datasets/MVTec_LOCO"``.
+        category (str): Category of the MVTec LOCO dataset (e.g. "breakfast_box").
+            Defaults to ``"breakfast_box"``.
+        train_batch_size (int, optional): Training batch size.
+            Defaults to ``32``.
+        eval_batch_size (int, optional): Test batch size.
+            Defaults to ``32``.
+        num_workers (int, optional): Number of workers.
+            Defaults to ``8``.
+        task TaskType): Task type, 'classification', 'detection' or 'segmentation'
+            Defaults to ``TaskType.SEGMENTATION``.
+        image_size (tuple[int, int], optional): Size to which input images should be resized.
+            Defaults to ``None``.
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        train_transform (Transform, optional): Transforms that should be applied to the input images during training.
+            Defaults to ``None``.
+        eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation.
+            Defaults to ``None``.
+        test_split_mode (TestSplitMode): Setting that determines how the testing subset is obtained.
+            Defaults to ``TestSplitMode.FROM_DIR``.
+        test_split_ratio (float): Fraction of images from the train set that will be reserved for testing.
+            Defaults to ``0.2``.
+        val_split_mode (ValSplitMode): Setting that determines how the validation subset is obtained.
+            Defaults to ``ValSplitMode.FROM_DIR``.
+        val_split_ratio (float): Fraction of train or test images that will be reserved for validation.
+            Defaults to ``0.5``.
+        seed (int | None, optional): Seed which may be set to a fixed value for reproducibility.
+            Defaults to ``None``.
+
+    Examples:
+        To create an MVTec LOCO AD datamodule with default settings:
+
+        >>> datamodule = MVTecLoco(root="anomalib/datasets/MVTec_LOCO")
+        >>> datamodule.setup()
+        >>> i, data = next(enumerate(datamodule.train_dataloader()))
+        >>> data.keys()
+        dict_keys(['image_path', 'label', 'image', 'mask_path', 'mask'])
+
+        >>> data["image"].shape
+        torch.Size([32, 3, 256, 256])
+
+        To change the category of the dataset:
+
+        >>> datamodule = MVTecLoco(category="pushpins")
+
+        To change the image and batch size:
+
+        >>> datamodule = MVTecLoco(image_size=(512, 512), train_batch_size=16, eval_batch_size=8)
+
+        MVTec LOCO AD dataset provide an independent validation set with normal images only in the 'validation' folder.
+        If you would like to use a different validation set splitted from train or test set,
+        you can use the ``val_split_mode`` and ``val_split_ratio`` arguments to create a new validation set.
+
+        >>> datamodule = MVTecLoco(val_split_mode=ValSplitMode.FROM_TEST, val_split_ratio=0.1)
+
+        This will subsample the test set by 10% and use it as the validation set.
+        If you would like to create a validation set synthetically that would
+        not change the test set, you can use the ``ValSplitMode.SYNTHETIC`` option.
+
+        >>> datamodule = MVTecLoco(val_split_mode=ValSplitMode.SYNTHETIC, val_split_ratio=0.2)
+    """
+
+    def __init__(
+        self,
+        root: Path | str = "./datasets/MVTec_LOCO",
+        category: str = "breakfast_box",
+        train_batch_size: int = 32,
+        eval_batch_size: int = 32,
+        num_workers: int = 8,
+        task: TaskType = TaskType.SEGMENTATION,
+        image_size: tuple[int, int] | None = None,
+        transform: Transform | None = None,
+        train_transform: Transform | None = None,
+        eval_transform: Transform | None = None,
+        test_split_mode: TestSplitMode = TestSplitMode.FROM_DIR,
+        test_split_ratio: float = 0.2,
+        val_split_mode: ValSplitMode = ValSplitMode.FROM_DIR,
+        val_split_ratio: float = 0.5,
+        seed: int | None = None,
+    ) -> None:
+        super().__init__(
+            train_batch_size=train_batch_size,
+            eval_batch_size=eval_batch_size,
+            image_size=image_size,
+            transform=transform,
+            train_transform=train_transform,
+            eval_transform=eval_transform,
+            num_workers=num_workers,
+            test_split_mode=test_split_mode,
+            test_split_ratio=test_split_ratio,
+            val_split_mode=val_split_mode,
+            val_split_ratio=val_split_ratio,
+            seed=seed,
+        )
+        self.task = task
+        self.root = Path(root)
+        self.category = category
+
+    def _setup(self, _stage: str | None = None) -> None:
+        """Set up the datasets, configs, and perform dynamic subset splitting.
+
+        This method overrides the parent class's method to also setup the val dataset.
+        The MVTec LOCO dataset provides an independent validation subset.
+        """
+        self.train_data = MVTecLocoDataset(
+            task=self.task,
+            transform=self.train_transform,
+            split=Split.TRAIN,
+            root=self.root,
+            category=self.category,
+        )
+        self.val_data = MVTecLocoDataset(
+            task=self.task,
+            transform=self.eval_transform,
+            split=Split.VAL,
+            root=self.root,
+            category=self.category,
+        )
+        self.test_data = MVTecLocoDataset(
+            task=self.task,
+            transform=self.eval_transform,
+            split=Split.TEST,
+            root=self.root,
+            category=self.category,
+        )
+
+    def prepare_data(self) -> None:
+        """Download the dataset if not available.
+
+        This method checks if the specified dataset is available in the file system.
+        If not, it downloads and extracts the dataset into the appropriate directory.
+
+        Example:
+            Assume the dataset is not available on the file system.
+            Here's how the directory structure looks before and after calling the
+            `prepare_data` method:
+
+            Before:
+
+            .. code-block:: bash
+
+                $ tree datasets
+                datasets
+                ├── dataset1
+                └── dataset2
+
+            Calling the method:
+
+            .. code-block:: python
+
+                >> datamodule = MVTecLoco(root="./datasets/MVTec_LOCO", category="breakfast_box")
+                >> datamodule.prepare_data()
+
+            After:
+
+            .. code-block:: bash
+
+                $ tree datasets
+                datasets
+                ├── dataset1
+                ├── dataset2
+                └── MVTec_LOCO
+                    ├── breakfast_box
+                    ├── ...
+                    └── splicing_connectors
+        """
+        if (self.root / self.category).is_dir():
+            logger.info("Found the dataset.")
+        else:
+            download_and_extract(self.root, DOWNLOAD_INFO)

anomalib/data/image/visa.py

@@ -0,0 +1,364 @@
+"""Visual Anomaly (VisA) Dataset (CC BY-NC-SA 4.0).
+
+Description:
+    This script contains PyTorch Dataset, Dataloader and PyTorch
+        Lightning DataModule for the Visual Anomal (VisA) dataset.
+    If the dataset is not on the file system, the script downloads and
+        extracts the dataset and create PyTorch data objects.
+License:
+    The VisA dataset is released under the Creative Commons
+    Attribution-NonCommercial-ShareAlike 4.0 International License
+    (CC BY-NC-SA 4.0)(https://creativecommons.org/licenses/by-nc-sa/4.0/).
+Reference:
+    - Zou, Y., Jeong, J., Pemula, L., Zhang, D., & Dabeer, O. (2022). SPot-the-Difference
+      Self-supervised Pre-training for Anomaly Detection and Segmentation. In European
+      Conference on Computer Vision (pp. 392-408). Springer, Cham.
+"""
+
+# Copyright (C) 2022-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# Subset splitting code adapted from https://github.com/amazon-science/spot-diff
+# Original licence: Apache-2.0
+
+
+import csv
+import logging
+import shutil
+from pathlib import Path
+
+import cv2
+from torchvision.transforms.v2 import Transform
+
+from anomalib import TaskType
+from anomalib.data.base import AnomalibDataModule, AnomalibDataset
+from anomalib.data.utils import (
+    DownloadInfo,
+    Split,
+    TestSplitMode,
+    ValSplitMode,
+    download_and_extract,
+)
+
+from .mvtec import make_mvtec_dataset
+
+logger = logging.getLogger(__name__)
+
+EXTENSIONS = (".png", ".jpg", ".JPG")
+
+DOWNLOAD_INFO = DownloadInfo(
+    name="VisA",
+    url="https://amazon-visual-anomaly.s3.us-west-2.amazonaws.com/VisA_20220922.tar",
+    hashsum="2eb8690c803ab37de0324772964100169ec8ba1fa3f7e94291c9ca673f40f362",
+)
+
+CATEGORIES = (
+    "candle",
+    "capsules",
+    "cashew",
+    "chewinggum",
+    "fryum",
+    "macaroni1",
+    "macaroni2",
+    "pcb1",
+    "pcb2",
+    "pcb3",
+    "pcb4",
+    "pipe_fryum",
+)
+
+
+class VisaDataset(AnomalibDataset):
+    """VisA dataset class.
+
+    Args:
+        task (TaskType): Task type, ``classification``, ``detection`` or ``segmentation``
+        root (str | Path): Path to the root of the dataset
+        category (str): Sub-category of the dataset, e.g. 'candle'
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        split (str | Split | None): Split of the dataset, usually Split.TRAIN or Split.TEST
+            Defaults to ``None``.
+
+    Examples:
+        To create a Visa dataset for classification:
+
+        .. code-block:: python
+
+            from anomalib.data.image.visa import VisaDataset
+            from anomalib.data.utils.transforms import get_transforms
+
+            transform = get_transforms(image_size=256)
+            dataset = VisaDataset(
+                task="classification",
+                transform=transform,
+                split="train",
+                root="./datasets/visa/visa_pytorch/",
+                category="candle",
+            )
+            dataset.setup()
+            dataset[0].keys()
+
+            # Output
+            dict_keys(['image_path', 'label', 'image'])
+
+        If you want to use the dataset for segmentation, you can use the same
+        code as above, with the task set to ``segmentation``. The dataset will
+        then have a ``mask`` key in the output dictionary.
+
+        .. code-block:: python
+
+            from anomalib.data.image.visa import VisaDataset
+            from anomalib.data.utils.transforms import get_transforms
+
+            transform = get_transforms(image_size=256)
+            dataset = VisaDataset(
+                task="segmentation",
+                transform=transform,
+                split="train",
+                root="./datasets/visa/visa_pytorch/",
+                category="candle",
+            )
+            dataset.setup()
+            dataset[0].keys()
+
+            # Output
+            dict_keys(['image_path', 'label', 'image', 'mask_path', 'mask'])
+
+    """
+
+    def __init__(
+        self,
+        task: TaskType,
+        root: str | Path,
+        category: str,
+        transform: Transform | None = None,
+        split: str | Split | None = None,
+    ) -> None:
+        super().__init__(task=task, transform=transform)
+
+        self.root_category = Path(root) / category
+        self.split = split
+        self.samples = make_mvtec_dataset(self.root_category, split=self.split, extensions=EXTENSIONS)
+
+
+class Visa(AnomalibDataModule):
+    """VisA Datamodule.
+
+    Args:
+        root (Path | str): Path to the root of the dataset
+            Defaults to ``"./datasets/visa"``.
+        category (str): Category of the Visa dataset such as ``candle``.
+            Defaults to ``"candle"``.
+        train_batch_size (int, optional): Training batch size.
+            Defaults to ``32``.
+        eval_batch_size (int, optional): Test batch size.
+            Defaults to ``32``.
+        num_workers (int, optional): Number of workers.
+            Defaults to ``8``.
+        task (TaskType): Task type, 'classification', 'detection' or 'segmentation'
+            Defaults to ``TaskType.SEGMENTATION``.
+        image_size (tuple[int, int], optional): Size to which input images should be resized.
+            Defaults to ``None``.
+        transform (Transform, optional): Transforms that should be applied to the input images.
+            Defaults to ``None``.
+        train_transform (Transform, optional): Transforms that should be applied to the input images during training.
+            Defaults to ``None``.
+        eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation.
+            Defaults to ``None``.
+        test_split_mode (TestSplitMode): Setting that determines how the testing subset is obtained.
+            Defaults to ``TestSplitMode.FROM_DIR``.
+        test_split_ratio (float): Fraction of images from the train set that will be reserved for testing.
+            Defaults to ``0.2``.
+        val_split_mode (ValSplitMode): Setting that determines how the validation subset is obtained.
+            Defaults to ``ValSplitMode.SAME_AS_TEST``.
+        val_split_ratio (float): Fraction of train or test images that will be reserved for validation.
+            Defatuls to ``0.5``.
+        seed (int | None, optional): Seed which may be set to a fixed value for reproducibility.
+            Defaults to ``None``.
+    """
+
+    def __init__(
+        self,
+        root: Path | str = "./datasets/visa",
+        category: str = "capsules",
+        train_batch_size: int = 32,
+        eval_batch_size: int = 32,
+        num_workers: int = 8,
+        task: TaskType | str = TaskType.SEGMENTATION,
+        image_size: tuple[int, int] | None = None,
+        transform: Transform | None = None,
+        train_transform: Transform | None = None,
+        eval_transform: Transform | None = None,
+        test_split_mode: TestSplitMode | str = TestSplitMode.FROM_DIR,
+        test_split_ratio: float = 0.2,
+        val_split_mode: ValSplitMode | str = ValSplitMode.SAME_AS_TEST,
+        val_split_ratio: float = 0.5,
+        seed: int | None = None,
+    ) -> None:
+        super().__init__(
+            train_batch_size=train_batch_size,
+            eval_batch_size=eval_batch_size,
+            num_workers=num_workers,
+            image_size=image_size,
+            transform=transform,
+            train_transform=train_transform,
+            eval_transform=eval_transform,
+            test_split_mode=test_split_mode,
+            test_split_ratio=test_split_ratio,
+            val_split_mode=val_split_mode,
+            val_split_ratio=val_split_ratio,
+            seed=seed,
+        )
+
+        self.task = TaskType(task)
+        self.root = Path(root)
+        self.split_root = self.root / "visa_pytorch"
+        self.category = category
+
+    def _setup(self, _stage: str | None = None) -> None:
+        self.train_data = VisaDataset(
+            task=self.task,
+            transform=self.train_transform,
+            split=Split.TRAIN,
+            root=self.split_root,
+            category=self.category,
+        )
+        self.test_data = VisaDataset(
+            task=self.task,
+            transform=self.eval_transform,
+            split=Split.TEST,
+            root=self.split_root,
+            category=self.category,
+        )
+
+    def prepare_data(self) -> None:
+        """Download the dataset if not available.
+
+        This method checks if the specified dataset is available in the file system.
+        If not, it downloads and extracts the dataset into the appropriate directory.
+
+        Example:
+            Assume the dataset is not available on the file system.
+            Here's how the directory structure looks before and after calling the
+            `prepare_data` method:
+
+            Before:
+
+            .. code-block:: bash
+
+                $ tree datasets
+                datasets
+                ├── dataset1
+                └── dataset2
+
+            Calling the method:
+
+            .. code-block:: python
+
+                >> datamodule = Visa()
+                >> datamodule.prepare_data()
+
+            After:
+
+            .. code-block:: bash
+
+                $ tree datasets
+                datasets
+                ├── dataset1
+                ├── dataset2
+                └── visa
+                    ├── candle
+                    ├── ...
+                    ├── pipe_fryum
+                    │   ├── Data
+                    │   └── image_anno.csv
+                    ├── split_csv
+                    │   ├── 1cls.csv
+                    │   ├── 2cls_fewshot.csv
+                    │   └── 2cls_highshot.csv
+                    ├── VisA_20220922.tar
+                    └── visa_pytorch
+                        ├── candle
+                        ├── ...
+                        ├── pcb4
+                        └── pipe_fryum
+
+            ``prepare_data`` ensures that the dataset is converted to MVTec
+            format. ``visa_pytorch`` is the directory that contains the dataset
+            in the MVTec format. ``visa`` is the directory that contains the
+            original dataset.
+        """
+        if (self.split_root / self.category).is_dir():
+            # dataset is available, and split has been applied
+            logger.info("Found the dataset and train/test split.")
+        elif (self.root / self.category).is_dir():
+            # dataset is available, but split has not yet been applied
+            logger.info("Found the dataset. Applying train/test split.")
+            self.apply_cls1_split()
+        else:
+            # dataset is not available
+            download_and_extract(self.root, DOWNLOAD_INFO)
+            logger.info("Downloaded the dataset. Applying train/test split.")
+            self.apply_cls1_split()
+
+    def apply_cls1_split(self) -> None:
+        """Apply the 1-class subset splitting using the fixed split in the csv file.
+
+        adapted from https://github.com/amazon-science/spot-diff
+        """
+        logger.info("preparing data")
+        categories = [
+            "candle",
+            "capsules",
+            "cashew",
+            "chewinggum",
+            "fryum",
+            "macaroni1",
+            "macaroni2",
+            "pcb1",
+            "pcb2",
+            "pcb3",
+            "pcb4",
+            "pipe_fryum",
+        ]
+
+        split_file = self.root / "split_csv" / "1cls.csv"
+
+        for category in categories:
+            train_folder = self.split_root / category / "train"
+            test_folder = self.split_root / category / "test"
+            mask_folder = self.split_root / category / "ground_truth"
+
+            train_img_good_folder = train_folder / "good"
+            test_img_good_folder = test_folder / "good"
+            test_img_bad_folder = test_folder / "bad"
+            test_mask_bad_folder = mask_folder / "bad"
+
+            train_img_good_folder.mkdir(parents=True, exist_ok=True)
+            test_img_good_folder.mkdir(parents=True, exist_ok=True)
+            test_img_bad_folder.mkdir(parents=True, exist_ok=True)
+            test_mask_bad_folder.mkdir(parents=True, exist_ok=True)
+
+        with split_file.open(encoding="utf-8") as file:
+            csvreader = csv.reader(file)
+            next(csvreader)
+            for row in csvreader:
+                category, split, label, image_path, mask_path = row
+                label = "good" if label == "normal" else "bad"
+                image_name = image_path.split("/")[-1]
+                mask_name = mask_path.split("/")[-1]
+
+                img_src_path = self.root / image_path
+                msk_src_path = self.root / mask_path
+                img_dst_path = self.split_root / category / split / label / image_name
+                msk_dst_path = self.split_root / category / "ground_truth" / label / mask_name
+
+                shutil.copyfile(img_src_path, img_dst_path)
+                if split == "test" and label == "bad":
+                    mask = cv2.imread(str(msk_src_path))
+
+                    # binarize mask
+                    mask[mask != 0] = 255
+
+                    cv2.imwrite(str(msk_dst_path), mask)

anomalib/data/predict.py

@@ -0,0 +1,52 @@
+"""Inference Dataset."""
+
+# Copyright (C) 2022-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+from pathlib import Path
+from typing import Any
+
+from torch.utils.data.dataset import Dataset
+from torchvision.transforms.v2 import Transform
+
+from anomalib.data.utils import get_image_filenames, read_image
+
+
+class PredictDataset(Dataset):
+    """Inference Dataset to perform prediction.
+
+    Args:
+        path (str | Path): Path to an image or image-folder.
+        transform (A.Compose | None, optional): Transform object describing the transforms that are
+            applied to the inputs.
+        image_size (int | tuple[int, int] | None, optional): Target image size
+            to resize the original image. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        path: str | Path,
+        transform: Transform | None = None,
+        image_size: int | tuple[int, int] = (256, 256),
+    ) -> None:
+        super().__init__()
+
+        self.image_filenames = get_image_filenames(path)
+        self.transform = transform
+        self.image_size = image_size
+
+    def __len__(self) -> int:
+        """Get the number of images in the given path."""
+        return len(self.image_filenames)
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        """Get the image based on the `index`."""
+        image_filename = self.image_filenames[index]
+        image = read_image(image_filename, as_tensor=True)
+        if self.transform:
+            image = self.transform(image)
+        pre_processed = {"image": image}
+        pre_processed["image_path"] = str(image_filename)
+
+        return pre_processed

anomalib/data/transforms/__init__.py

@@ -0,0 +1,8 @@
+"""Custom input transforms for Anomalib."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .center_crop import ExportableCenterCrop
+
+__all__ = ["ExportableCenterCrop"]

anomalib/data/transforms/center_crop.py

@@ -0,0 +1,87 @@
+"""Custom Torchvision transforms for Anomalib."""
+
+# Original Code
+# Copyright (c) Soumith Chintala 2016
+# https://github.com/pytorch/vision/blob/v0.16.1/torchvision/transforms/v2/functional/_geometry.py
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Modified
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any
+
+import torch
+from torch.nn.functional import pad
+from torchvision.transforms.v2 import Transform
+from torchvision.transforms.v2.functional._geometry import (
+    _center_crop_compute_padding,
+    _center_crop_parse_output_size,
+    _parse_pad_padding,
+)
+
+
+def _center_crop_compute_crop_anchor(
+    crop_height: int,
+    crop_width: int,
+    image_height: int,
+    image_width: int,
+) -> tuple[int, int]:
+    """Compute the anchor point for center-cropping.
+
+    This function is a modified version of the torchvision.transforms.functional._center_crop_compute_crop_anchor
+    function. The original function uses `round` to compute the anchor point, which is not compatible with ONNX.
+
+    Args:
+        crop_height (int): Desired height of the crop.
+        crop_width (int): Desired width of the crop.
+        image_height (int): Height of the input image.
+        image_width (int): Width of the input image.
+    """
+    crop_top = torch.tensor((image_height - crop_height) / 2.0).round().int().item()
+    crop_left = torch.tensor((image_width - crop_width) / 2.0).round().int().item()
+    return crop_top, crop_left
+
+
+def center_crop_image(image: torch.Tensor, output_size: list[int]) -> torch.Tensor:
+    """Apply center-cropping to an input image.
+
+    Uses the modified anchor point computation function to compute the anchor point for center-cropping.
+
+    Args:
+        image (torch.Tensor): Input image to be center-cropped.
+        output_size (list[int]): Desired output size of the crop.
+    """
+    crop_height, crop_width = _center_crop_parse_output_size(output_size)
+    shape = image.shape
+    if image.numel() == 0:
+        return image.reshape(shape[:-2] + (crop_height, crop_width))
+    image_height, image_width = shape[-2:]
+
+    if crop_height > image_height or crop_width > image_width:
+        padding_ltrb = _center_crop_compute_padding(crop_height, crop_width, image_height, image_width)
+        image = pad(image, _parse_pad_padding(padding_ltrb), value=0.0)
+
+        image_height, image_width = image.shape[-2:]
+        if crop_width == image_width and crop_height == image_height:
+            return image
+
+    crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, image_height, image_width)
+    return image[..., crop_top : (crop_top + crop_height), crop_left : (crop_left + crop_width)]
+
+
+class ExportableCenterCrop(Transform):
+    """Transform that applies center-cropping to an input image and allows to be exported to ONNX.
+
+    Args:
+        size (int | tuple[int, int]): Desired output size of the crop.
+    """
+
+    def __init__(self, size: int | tuple[int, int]) -> None:
+        super().__init__()
+        self.size = list(size) if isinstance(size, tuple) else [size, size]
+
+    def _transform(self, inpt: torch.Tensor, params: dict[str, Any]) -> torch.Tensor:
+        """Apply the transform."""
+        del params
+        return center_crop_image(inpt, output_size=self.size)

anomalib/data/utils/__init__.py

@@ -0,0 +1,56 @@
+"""Helper utilities for data."""
+
+# Copyright (C) 2022-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .augmenter import Augmenter
+from .boxes import boxes_to_anomaly_maps, boxes_to_masks, masks_to_boxes
+from .download import DownloadInfo, download_and_extract
+from .generators import random_2d_perlin
+from .image import (
+    generate_output_image_filename,
+    get_image_filenames,
+    get_image_height_and_width,
+    read_depth_image,
+    read_image,
+    read_mask,
+)
+from .label import LabelName
+from .path import (
+    DirType,
+    _check_and_convert_path,
+    _prepare_files_labels,
+    resolve_path,
+    validate_and_resolve_path,
+    validate_path,
+)
+from .split import Split, TestSplitMode, ValSplitMode, concatenate_datasets, random_split, split_by_label
+
+__all__ = [
+    "generate_output_image_filename",
+    "get_image_filenames",
+    "get_image_height_and_width",
+    "random_2d_perlin",
+    "read_image",
+    "read_mask",
+    "read_depth_image",
+    "random_split",
+    "split_by_label",
+    "concatenate_datasets",
+    "Split",
+    "ValSplitMode",
+    "TestSplitMode",
+    "LabelName",
+    "DirType",
+    "Augmenter",
+    "masks_to_boxes",
+    "boxes_to_masks",
+    "boxes_to_anomaly_maps",
+    "download_and_extract",
+    "DownloadInfo",
+    "_check_and_convert_path",
+    "_prepare_files_labels",
+    "resolve_path",
+    "validate_path",
+    "validate_and_resolve_path",
+]

anomalib/data/utils/augmenter.py

@@ -0,0 +1,172 @@
+"""Augmenter module to generates out-of-distribution samples for the DRAEM implementation."""
+
+# Original Code
+# Copyright (c) 2021 VitjanZ
+# https://github.com/VitjanZ/DRAEM.
+# SPDX-License-Identifier: MIT
+#
+# Modified
+# Copyright (C) 2022-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+import math
+import random
+from pathlib import Path
+
+import cv2
+import imgaug.augmenters as iaa
+import numpy as np
+import torch
+from PIL import Image
+from torchvision.datasets.folder import IMG_EXTENSIONS
+
+from anomalib.data.utils.generators.perlin import random_2d_perlin
+
+
+def nextpow2(value: int) -> int:
+    """Return the smallest power of 2 greater than or equal to the input value."""
+    return 2 ** (math.ceil(math.log(value, 2)))
+
+
+class Augmenter:
+    """Class that generates noisy augmentations of input images.
+
+    Args:
+        anomaly_source_path (str | None): Path to a folder of images that will be used as source of the anomalous
+        noise. If not specified, random noise will be used instead.
+        p_anomalous (float): Probability that the anomalous perturbation will be applied to a given image.
+        beta (float): Parameter that determines the opacity of the noise mask.
+    """
+
+    def __init__(
+        self,
+        anomaly_source_path: str | None = None,
+        p_anomalous: float = 0.5,
+        beta: float | tuple[float, float] = (0.2, 1.0),
+    ) -> None:
+        self.p_anomalous = p_anomalous
+        self.beta = beta
+
+        self.anomaly_source_paths: list[Path] = []
+        if anomaly_source_path is not None:
+            for img_ext in IMG_EXTENSIONS:
+                self.anomaly_source_paths.extend(Path(anomaly_source_path).rglob("*" + img_ext))
+
+        self.augmenters = [
+            iaa.GammaContrast((0.5, 2.0), per_channel=True),
+            iaa.MultiplyAndAddToBrightness(mul=(0.8, 1.2), add=(-30, 30)),
+            iaa.pillike.EnhanceSharpness(),
+            iaa.AddToHueAndSaturation((-50, 50), per_channel=True),
+            iaa.Solarize(0.5, threshold=(32, 128)),
+            iaa.Posterize(),
+            iaa.Invert(),
+            iaa.pillike.Autocontrast(),
+            iaa.pillike.Equalize(),
+            iaa.Affine(rotate=(-45, 45)),
+        ]
+        self.rot = iaa.Sequential([iaa.Affine(rotate=(-90, 90))])
+
+    def rand_augmenter(self) -> iaa.Sequential:
+        """Select 3 random transforms that will be applied to the anomaly source images.
+
+        Returns:
+            A selection of 3 transforms.
+        """
+        aug_ind = np.random.default_rng().choice(np.arange(len(self.augmenters)), 3, replace=False)
+        return iaa.Sequential([self.augmenters[aug_ind[0]], self.augmenters[aug_ind[1]], self.augmenters[aug_ind[2]]])
+
+    def generate_perturbation(
+        self,
+        height: int,
+        width: int,
+        anomaly_source_path: Path | str | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Generate an image containing a random anomalous perturbation using a source image.
+
+        Args:
+            height (int): height of the generated image.
+            width: (int): width of the generated image.
+            anomaly_source_path (Path | str | None): Path to an image file. If not provided, random noise will be used
+            instead.
+
+        Returns:
+            Image containing a random anomalous perturbation, and the corresponding ground truth anomaly mask.
+        """
+        # Generate random perlin noise
+        perlin_scale = 6
+        min_perlin_scale = 0
+
+        perlin_scalex = 2 ** np.random.default_rng().integers(min_perlin_scale, perlin_scale)
+        perlin_scaley = 2 ** np.random.default_rng().integers(min_perlin_scale, perlin_scale)
+
+        perlin_noise = random_2d_perlin((nextpow2(height), nextpow2(width)), (perlin_scalex, perlin_scaley))[
+            :height,
+            :width,
+        ]
+        perlin_noise = self.rot(image=perlin_noise)
+
+        # Create mask from perlin noise
+        mask = np.where(perlin_noise > 0.5, np.ones_like(perlin_noise), np.zeros_like(perlin_noise))
+        mask = np.expand_dims(mask, axis=2).astype(np.float32)
+
+        # Load anomaly source image
+        if anomaly_source_path:
+            anomaly_source_img = np.array(Image.open(anomaly_source_path))
+            anomaly_source_img = cv2.resize(anomaly_source_img, dsize=(width, height))
+        else:  # if no anomaly source is specified, we use the perlin noise as anomalous source
+            anomaly_source_img = np.expand_dims(perlin_noise, 2).repeat(3, 2)
+            anomaly_source_img = (anomaly_source_img * 255).astype(np.uint8)
+
+        # Augment anomaly source image
+        aug = self.rand_augmenter()
+        anomaly_img_augmented = aug(image=anomaly_source_img)
+
+        # Create anomalous perturbation that we will apply to the image
+        perturbation = anomaly_img_augmented.astype(np.float32) * mask / 255.0
+
+        return perturbation, mask
+
+    def augment_batch(self, batch: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Generate anomalous augmentations for a batch of input images.
+
+        Args:
+            batch (torch.Tensor): Batch of input images
+
+        Returns:
+            - Augmented image to which anomalous perturbations have been added.
+            - Ground truth masks corresponding to the anomalous perturbations.
+        """
+        batch_size, channels, height, width = batch.shape
+
+        # Collect perturbations
+        perturbations_list = []
+        masks_list = []
+        for _ in range(batch_size):
+            if torch.rand(1) > self.p_anomalous:  # include normal samples
+                perturbations_list.append(torch.zeros((channels, height, width)))
+                masks_list.append(torch.zeros((1, height, width)))
+            else:
+                anomaly_source_path = (
+                    random.sample(self.anomaly_source_paths, 1)[0] if len(self.anomaly_source_paths) > 0 else None
+                )
+                perturbation, mask = self.generate_perturbation(height, width, anomaly_source_path)
+                perturbations_list.append(torch.Tensor(perturbation).permute((2, 0, 1)))
+                masks_list.append(torch.Tensor(mask).permute((2, 0, 1)))
+
+        perturbations = torch.stack(perturbations_list).to(batch.device)
+        masks = torch.stack(masks_list).to(batch.device)
+
+        # Apply perturbations batch wise
+        if isinstance(self.beta, float):
+            beta = self.beta
+        elif isinstance(self.beta, tuple):
+            beta = torch.rand(batch_size) * (self.beta[1] - self.beta[0]) + self.beta[0]
+            beta = beta.view(batch_size, 1, 1, 1).expand_as(batch).to(batch.device)  # type: ignore[attr-defined]
+        else:
+            msg = "Beta must be either float or tuple of floats"
+            raise TypeError(msg)
+
+        augmented_batch = batch * (1 - masks) + (beta) * perturbations + (1 - beta) * batch * (masks)
+
+        return augmented_batch, masks

anomalib/data/utils/boxes.py

@@ -0,0 +1,117 @@
+"""Helper functions for processing bounding box detections and annotations."""
+
+# Copyright (C) 2022-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+import torch
+
+from anomalib.utils.cv import connected_components_cpu, connected_components_gpu
+
+
+def masks_to_boxes(
+    masks: torch.Tensor,
+    anomaly_maps: torch.Tensor | None = None,
+) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+    """Convert a batch of segmentation masks to bounding box coordinates.
+
+    Args:
+        masks (torch.Tensor): Input tensor of shape (B, 1, H, W), (B, H, W) or (H, W)
+        anomaly_maps (Tensor | None, optional): Anomaly maps of shape (B, 1, H, W), (B, H, W) or (H, W) which are
+            used to determine an anomaly score for the converted bounding boxes.
+
+    Returns:
+        list[torch.Tensor]: A list of length B where each element is a tensor of shape (N, 4)
+            containing the bounding box coordinates of the objects in the masks in xyxy format.
+        list[torch.Tensor]: A list of length B where each element is a tensor of length (N)
+            containing an anomaly score for each of the converted boxes.
+    """
+    height, width = masks.shape[-2:]
+    masks = masks.view((-1, 1, height, width)).float()  # reshape to (B, 1, H, W) and cast to float
+    if anomaly_maps is not None:
+        anomaly_maps = anomaly_maps.view((-1,) + masks.shape[-2:])
+
+    if masks.is_cpu:
+        batch_comps = connected_components_cpu(masks).squeeze(1)
+    else:
+        batch_comps = connected_components_gpu(masks).squeeze(1)
+
+    batch_boxes = []
+    batch_scores = []
+    for im_idx, im_comps in enumerate(batch_comps):
+        labels = torch.unique(im_comps)
+        im_boxes = []
+        im_scores = []
+        for label in labels[labels != 0]:
+            y_loc, x_loc = torch.where(im_comps == label)
+            # add box
+            box = torch.Tensor([torch.min(x_loc), torch.min(y_loc), torch.max(x_loc), torch.max(y_loc)]).to(
+                masks.device,
+            )
+            im_boxes.append(box)
+            if anomaly_maps is not None:
+                im_scores.append(torch.max(anomaly_maps[im_idx, y_loc, x_loc]))
+        batch_boxes.append(torch.stack(im_boxes) if im_boxes else torch.empty((0, 4), device=masks.device))
+        batch_scores.append(torch.stack(im_scores) if im_scores else torch.empty(0, device=masks.device))
+
+    return batch_boxes, batch_scores
+
+
+def boxes_to_masks(boxes: list[torch.Tensor], image_size: tuple[int, int]) -> torch.Tensor:
+    """Convert bounding boxes to segmentations masks.
+
+    Args:
+        boxes (list[torch.Tensor]): A list of length B where each element is a tensor of shape (N, 4)
+            containing the bounding box coordinates of the regions of interest in xyxy format.
+        image_size (tuple[int, int]): Image size of the output masks in (H, W) format.
+
+    Returns:
+        Tensor: torch.Tensor of shape (B, H, W) in which each slice is a binary mask showing the pixels contained by a
+            bounding box.
+    """
+    masks = torch.zeros((len(boxes), *image_size)).to(boxes[0].device)
+    for im_idx, im_boxes in enumerate(boxes):
+        for box in im_boxes:
+            x_1, y_1, x_2, y_2 = box.int()
+            masks[im_idx, y_1 : y_2 + 1, x_1 : x_2 + 1] = 1
+    return masks
+
+
+def boxes_to_anomaly_maps(boxes: torch.Tensor, scores: torch.Tensor, image_size: tuple[int, int]) -> torch.Tensor:
+    """Convert bounding box coordinates to anomaly heatmaps.
+
+    Args:
+        boxes (list[torch.Tensor]): A list of length B where each element is a tensor of shape (N, 4)
+            containing the bounding box coordinates of the regions of interest in xyxy format.
+        scores (list[torch.Tensor]): A list of length B where each element is a 1D tensor of length N
+            containing the anomaly scores for each region of interest.
+        image_size (tuple[int, int]): Image size of the output masks in (H, W) format.
+
+    Returns:
+        Tensor: torch.Tensor of shape (B, H, W). The pixel locations within each bounding box are collectively
+            assigned the anomaly score of the bounding box. In the case of overlapping bounding boxes,
+            the highest score is used.
+    """
+    anomaly_maps = torch.zeros((len(boxes), *image_size)).to(boxes[0].device)
+    for im_idx, (im_boxes, im_scores) in enumerate(zip(boxes, scores, strict=False)):
+        im_map = torch.zeros((im_boxes.shape[0], *image_size))
+        for box_idx, (box, score) in enumerate(zip(im_boxes, im_scores, strict=True)):
+            x_1, y_1, x_2, y_2 = box.int()
+            im_map[box_idx, y_1 : y_2 + 1, x_1 : x_2 + 1] = score
+            anomaly_maps[im_idx], _ = im_map.max(dim=0)
+    return anomaly_maps
+
+
+def scale_boxes(boxes: torch.Tensor, image_size: torch.Size, new_size: torch.Size) -> torch.Tensor:
+    """Scale bbox coordinates to a new image size.
+
+    Args:
+        boxes (torch.Tensor): Boxes of shape (N, 4) - (x1, y1, x2, y2).
+        image_size (Size): Size of the original image in which the bbox coordinates were retrieved.
+        new_size (Size): New image size to which the bbox coordinates will be scaled.
+
+    Returns:
+        Tensor: Updated boxes of shape (N, 4) - (x1, y1, x2, y2).
+    """
+    scale = torch.Tensor([*new_size]) / torch.Tensor([*image_size])
+    return boxes * scale.repeat(2).to(boxes.device)

anomalib/data/utils/download.py

@@ -0,0 +1,364 @@
+"""Helper to show progress bars with `urlretrieve`, check hash of file."""
+
+# Copyright (C) 2022-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+import hashlib
+import io
+import logging
+import os
+import re
+import tarfile
+from collections.abc import Iterable
+from dataclasses import dataclass
+from pathlib import Path
+from tarfile import TarFile, TarInfo
+from urllib.request import urlretrieve
+from zipfile import ZipFile
+
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class DownloadInfo:
+    """Info needed to download a dataset from a url."""
+
+    name: str
+    url: str
+    hashsum: str
+    filename: str | None = None
+
+
+class DownloadProgressBar(tqdm):
+    """Create progress bar for urlretrieve. Subclasses `tqdm`.
+
+    For information about the parameters in constructor, refer to `tqdm`'s documentation.
+
+    Args:
+        iterable (Iterable | None): Iterable to decorate with a progressbar.
+                            Leave blank to manually manage the updates.
+        desc (str | None): Prefix for the progressbar.
+        total (int | float | None): The number of expected iterations. If unspecified,
+                                            len(iterable) is used if possible. If float("inf") or as a last
+                                            resort, only basic progress statistics are displayed
+                                            (no ETA, no progressbar).
+                                            If `gui` is True and this parameter needs subsequent updating,
+                                            specify an initial arbitrary large positive number,
+                                            e.g. 9e9.
+        leave (bool | None): upon termination of iteration. If `None`, will leave only if `position` is `0`.
+        file (io.TextIOWrapper |  io.StringIO | None): Specifies where to output the progress messages
+                                                            (default: sys.stderr). Uses `file.write(str)` and
+                                                            `file.flush()` methods.  For encoding, see
+                                                            `write_bytes`.
+        ncols (int | None): The width of the entire output message. If specified,
+                            dynamically resizes the progressbar to stay within this bound.
+                            If unspecified, attempts to use environment width. The
+                            fallback is a meter width of 10 and no limit for the counter and
+                            statistics. If 0, will not print any meter (only stats).
+        mininterval (float | None): Minimum progress display update interval [default: 0.1] seconds.
+        maxinterval (float | None): Maximum progress display update interval [default: 10] seconds.
+                                    Automatically adjusts `miniters` to correspond to `mininterval`
+                                    after long display update lag. Only works if `dynamic_miniters`
+                                    or monitor thread is enabled.
+        miniters (int | float | None): Minimum progress display update interval, in iterations.
+                                            If 0 and `dynamic_miniters`, will automatically adjust to equal
+                                            `mininterval` (more CPU efficient, good for tight loops).
+                                            If > 0, will skip display of specified number of iterations.
+                                            Tweak this and `mininterval` to get very efficient loops.
+                                            If your progress is erratic with both fast and slow iterations
+                                            (network, skipping items, etc) you should set miniters=1.
+        use_ascii (str | bool | None): If unspecified or False, use unicode (smooth blocks) to fill
+                                        the meter. The fallback is to use ASCII characters " 123456789#".
+        disable (bool | None): Whether to disable the entire progressbar wrapper
+                                    [default: False]. If set to None, disable on non-TTY.
+        unit (str | None): String that will be used to define the unit of each iteration
+                            [default: it].
+        unit_scale (int | float | bool): If 1 or True, the number of iterations will be reduced/scaled
+                            automatically and a metric prefix following the
+                            International System of Units standard will be added
+                            (kilo, mega, etc.) [default: False]. If any other non-zero
+                            number, will scale `total` and `n`.
+        dynamic_ncols (bool | None): If set, constantly alters `ncols` and `nrows` to the
+                                        environment (allowing for window resizes) [default: False].
+        smoothing (float | None): Exponential moving average smoothing factor for speed estimates
+                                    (ignored in GUI mode). Ranges from 0 (average speed) to 1
+                                    (current/instantaneous speed) [default: 0.3].
+        bar_format (str | None):  Specify a custom bar string formatting. May impact performance.
+                                    [default: '{l_bar}{bar}{r_bar}'], where
+                                    l_bar='{desc}: {percentage:3.0f}%|' and
+                                    r_bar='| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, '
+                                    '{rate_fmt}{postfix}]'
+                                    Possible vars: l_bar, bar, r_bar, n, n_fmt, total, total_fmt,
+                                    percentage, elapsed, elapsed_s, ncols, nrows, desc, unit,
+                                    rate, rate_fmt, rate_noinv, rate_noinv_fmt,
+                                    rate_inv, rate_inv_fmt, postfix, unit_divisor,
+                                    remaining, remaining_s, eta.
+                                    Note that a trailing ": " is automatically removed after {desc}
+                                    if the latter is empty.
+        initial (int | float | None): The initial counter value. Useful when restarting a progress
+                                            bar [default: 0]. If using float, consider specifying `{n:.3f}`
+                                            or similar in `bar_format`, or specifying `unit_scale`.
+        position (int | None): Specify the line offset to print this bar (starting from 0)
+                                    Automatic if unspecified.
+                                    Useful to manage multiple bars at once (eg, from threads).
+        postfix (dict | None): Specify additional stats to display at the end of the bar.
+                                    Calls `set_postfix(**postfix)` if possible (dict).
+        unit_divisor (float | None): [default: 1000], ignored unless `unit_scale` is True.
+        write_bytes (bool | None): If (default: None) and `file` is unspecified,
+                                    bytes will be written in Python 2. If `True` will also write
+                                    bytes. In all other cases will default to unicode.
+        lock_args (tuple | None): Passed to `refresh` for intermediate output
+                                    (initialisation, iterating, and updating).
+                                    nrows (int | None): The screen height. If specified, hides nested bars
+                                    outside this bound. If unspecified, attempts to use environment height.
+                                    The fallback is 20.
+        colour (str | None): Bar colour (e.g. 'green', '#00ff00').
+        delay (float | None): Don't display until [default: 0] seconds have elapsed.
+        gui (bool | None): WARNING: internal parameter - do not use.
+                                Use tqdm.gui.tqdm(...) instead. If set, will attempt to use
+                                matplotlib animations for a graphical output [default: False].
+
+
+    Example:
+        >>> with DownloadProgressBar(unit='B', unit_scale=True, miniters=1, desc=url.split('/')[-1]) as p_bar:
+        >>>         urllib.request.urlretrieve(url, filename=output_path, reporthook=p_bar.update_to)
+    """
+
+    def __init__(
+        self,
+        iterable: Iterable | None = None,
+        desc: str | None = None,
+        total: int | float | None = None,
+        leave: bool | None = True,
+        file: io.TextIOWrapper | io.StringIO | None = None,
+        ncols: int | None = None,
+        mininterval: float | None = 0.1,
+        maxinterval: float | None = 10.0,
+        miniters: int | float | None = None,
+        use_ascii: bool | str | None = None,
+        disable: bool | None = False,
+        unit: str | None = "it",
+        unit_scale: bool | int | float | None = False,
+        dynamic_ncols: bool | None = False,
+        smoothing: float | None = 0.3,
+        bar_format: str | None = None,
+        initial: int | float | None = 0,
+        position: int | None = None,
+        postfix: dict | None = None,
+        unit_divisor: float | None = 1000,
+        write_bytes: bool | None = None,
+        lock_args: tuple | None = None,
+        nrows: int | None = None,
+        colour: str | None = None,
+        delay: float | None = 0,
+        gui: bool | None = False,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            iterable=iterable,
+            desc=desc,
+            total=total,
+            leave=leave,
+            file=file,
+            ncols=ncols,
+            mininterval=mininterval,
+            maxinterval=maxinterval,
+            miniters=miniters,
+            ascii=use_ascii,
+            disable=disable,
+            unit=unit,
+            unit_scale=unit_scale,
+            dynamic_ncols=dynamic_ncols,
+            smoothing=smoothing,
+            bar_format=bar_format,
+            initial=initial,
+            position=position,
+            postfix=postfix,
+            unit_divisor=unit_divisor,
+            write_bytes=write_bytes,
+            lock_args=lock_args,
+            nrows=nrows,
+            colour=colour,
+            delay=delay,
+            gui=gui,
+            **kwargs,
+        )
+        self.total: int | float | None
+
+    def update_to(self, chunk_number: int = 1, max_chunk_size: int = 1, total_size: int | None = None) -> None:
+        """Progress bar hook for tqdm.
+
+        Based on https://stackoverflow.com/a/53877507
+        The implementor does not have to bother about passing parameters to this as it gets them from urlretrieve.
+        However the context needs a few parameters. Refer to the example.
+
+        Args:
+            chunk_number (int, optional): The current chunk being processed. Defaults to 1.
+            max_chunk_size (int, optional): Maximum size of each chunk. Defaults to 1.
+            total_size (int, optional): Total download size. Defaults to None.
+        """
+        if total_size is not None:
+            self.total = total_size
+        self.update(chunk_number * max_chunk_size - self.n)
+
+
+def is_file_potentially_dangerous(file_name: str) -> bool:
+    """Check if a file is potentially dangerous.
+
+    Args:
+        file_name (str): Filename.
+
+    Returns:
+        bool: True if the member is potentially dangerous, False otherwise.
+
+    """
+    # Some example criteria. We could expand this.
+    unsafe_patterns = ["/etc/", "/root/"]
+    return any(re.search(pattern, file_name) for pattern in unsafe_patterns)
+
+
+def safe_extract(tar_file: TarFile, root: Path, members: list[TarInfo]) -> None:
+    """Extract safe members from a tar archive.
+
+    Args:
+        tar_file (TarFile): TarFile object.
+        root (Path): Root directory where the dataset will be stored.
+        members (List[TarInfo]): List of safe members to be extracted.
+
+    """
+    for member in members:
+        tar_file.extract(member, root)
+
+
+def generate_hash(file_path: str | Path, algorithm: str = "sha256") -> str:
+    """Generate a hash of a file using the specified algorithm.
+
+    Args:
+        file_path (str | Path): Path to the file to hash.
+        algorithm (str): The hashing algorithm to use (e.g., 'sha256', 'sha3_512').
+
+    Returns:
+        str: The hexadecimal hash string of the file.
+
+    Raises:
+        ValueError: If the specified hashing algorithm is not supported.
+    """
+    # Get the hashing algorithm.
+    try:
+        hasher = getattr(hashlib, algorithm)()
+    except AttributeError as err:
+        msg = f"Unsupported hashing algorithm: {algorithm}"
+        raise ValueError(msg) from err
+
+    # Read the file in chunks to avoid loading it all into memory
+    with Path(file_path).open("rb") as file:
+        for chunk in iter(lambda: file.read(4096), b""):
+            hasher.update(chunk)
+
+    # Return the computed hash value in hexadecimal format
+    return hasher.hexdigest()
+
+
+def check_hash(file_path: Path, expected_hash: str, algorithm: str = "sha256") -> None:
+    """Raise value error if hash does not match the calculated hash of the file.
+
+    Args:
+        file_path (Path): Path to file.
+        expected_hash (str): Expected hash of the file.
+        algorithm (str): Hashing algorithm to use ('sha256', 'sha3_512', etc.).
+    """
+    # Compare the calculated hash with the expected hash
+    calculated_hash = generate_hash(file_path, algorithm)
+    if calculated_hash != expected_hash:
+        msg = (
+            f"Calculated hash {calculated_hash} of downloaded file {file_path} does not match the required hash "
+            f"{expected_hash}."
+        )
+        raise ValueError(msg)
+
+
+def extract(file_name: Path, root: Path) -> None:
+    """Extract a dataset.
+
+    Args:
+        file_name (Path): Path of the file to be extracted.
+        root (Path): Root directory where the dataset will be stored.
+
+    """
+    logger.info("Extracting dataset into root folder.")
+
+    # Safely extract zip files
+    if file_name.suffix == ".zip":
+        with ZipFile(file_name, "r") as zip_file:
+            for file_info in zip_file.infolist():
+                if not is_file_potentially_dangerous(file_info.filename):
+                    zip_file.extract(file_info, root)
+
+    # Safely extract tar files.
+    elif file_name.suffix in (".tar", ".gz", ".xz", ".tgz"):
+        with tarfile.open(file_name) as tar_file:
+            members = tar_file.getmembers()
+            safe_members = [member for member in members if not is_file_potentially_dangerous(member.name)]
+            safe_extract(tar_file, root, safe_members)
+
+    else:
+        msg = f"Unrecognized file format: {file_name}"
+        raise ValueError(msg)
+
+    logger.info("Cleaning up files.")
+    file_name.unlink()
+
+
+def download_and_extract(root: Path, info: DownloadInfo) -> None:
+    """Download and extract a dataset.
+
+    Args:
+        root (Path): Root directory where the dataset will be stored.
+        info (DownloadInfo): Info needed to download the dataset.
+    """
+    root.mkdir(parents=True, exist_ok=True)
+
+    # save the compressed file in the specified root directory, using the same file name as on the server
+    downloaded_file_path = root / info.filename if info.filename else root / info.url.split("/")[-1]
+
+    if downloaded_file_path.exists():
+        logger.info("Existing dataset archive found. Skipping download stage.")
+    else:
+        logger.info("Downloading the %s dataset.", info.name)
+        # audit url. allowing only http:// or https://
+        if info.url.startswith("http://") or info.url.startswith("https://"):
+            with DownloadProgressBar(unit="B", unit_scale=True, miniters=1, desc=info.name) as progress_bar:
+                urlretrieve(  # noqa: S310  # nosec B310
+                    url=f"{info.url}",
+                    filename=downloaded_file_path,
+                    reporthook=progress_bar.update_to,
+                )
+            logger.info("Checking the hash of the downloaded file.")
+            check_hash(downloaded_file_path, info.hashsum)
+        else:
+            msg = f"Invalid URL to download dataset. Supported 'http://' or 'https://' but '{info.url}' is requested"
+            raise RuntimeError(msg)
+
+    extract(downloaded_file_path, root)
+
+
+def is_within_directory(directory: Path, target: Path) -> bool:
+    """Check if a target path is located within a given directory.
+
+    Args:
+        directory (Path): path of the parent directory
+        target (Path): path of the target
+
+    Returns:
+        (bool): True if the target is within the directory, False otherwise
+    """
+    abs_directory = directory.resolve()
+    abs_target = target.resolve()
+
+    # TODO(djdameln): Replace with pathlib is_relative_to after switching to Python 3.10
+    # CVS-122655
+    prefix = os.path.commonprefix([abs_directory, abs_target])
+    return prefix == str(abs_directory)

anomalib/data/utils/generators/__init__.py

@@ -0,0 +1,8 @@
+"""Utilities to generate synthetic data."""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .perlin import random_2d_perlin
+
+__all__ = ["random_2d_perlin"]