New datamodules design

anomalib/config/config.py

@@ -136,6 +136,27 @@ def get_configurable_parameters(
     if "format" not in config.dataset.keys():
         config.dataset.format = "mvtec"
 
+    if "create_validation_set" in config.dataset.keys():
+        warn(
+            "The 'create_validation_set' parameter is deprecated and will be removed in v0.4.0. Please use "
+            "'validation_split_mode' instead."
+        )
+        config.dataset.validation_split_mode = "from_test" if config.dataset.create_validation_set else "same_as_test"
+
+    if "test_batch_size" in config.dataset.keys():
+        warn(
+            "The 'test_batch_size' parameter is deprecated and will be removed in v0.4.0. Please use "
+            "'eval_batch_size' instead."
+        )
+        config.dataset.eval_batch_size = config.dataset.test_batch_size
+
+    if "transform_config" in config.dataset.keys() and "val" in config.dataset.transform_config.keys():
+        warn(
+            "The 'transform_config.val' parameter is deprecated and will be removed in v0.4.0. Please use "
+            "'transform_config.eval' instead."
+        )
+        config.dataset.transform_config.eval = config.dataset.transform_config.val
+
     config = update_input_size_config(config)
 
     # Project Configs

anomalib/data/__init__.py

@@ -7,8 +7,8 @@
 from typing import Union
 
 from omegaconf import DictConfig, ListConfig
-from pytorch_lightning import LightningDataModule
 
+from .base import AnomalibDataModule, AnomalibDataset
 from .btech import BTech
 from .folder import Folder
 from .inference import InferenceDataset
@@ -17,7 +17,7 @@
 logger = logging.getLogger(__name__)
 
 
-def get_datamodule(config: Union[DictConfig, ListConfig]) -> LightningDataModule:
+def get_datamodule(config: Union[DictConfig, ListConfig]) -> AnomalibDataModule:
     """Get Anomaly Datamodule.
 
     Args:
@@ -28,37 +28,33 @@ def get_datamodule(config: Union[DictConfig, ListConfig]) -> LightningDataModule
     """
     logger.info("Loading the datamodule")
 
-    datamodule: LightningDataModule
+    datamodule: AnomalibDataModule
 
     if config.dataset.format.lower() == "mvtec":
         datamodule = MVTec(
-            # TODO: Remove config values. IAAALD-211
             root=config.dataset.path,
             category=config.dataset.category,
             image_size=(config.dataset.image_size[0], config.dataset.image_size[1]),
             train_batch_size=config.dataset.train_batch_size,
-            test_batch_size=config.dataset.test_batch_size,
+            eval_batch_size=config.dataset.eval_batch_size,
             num_workers=config.dataset.num_workers,
-            seed=config.project.seed,
             task=config.dataset.task,
             transform_config_train=config.dataset.transform_config.train,
-            transform_config_val=config.dataset.transform_config.val,
-            create_validation_set=config.dataset.create_validation_set,
+            transform_config_eval=config.dataset.transform_config.eval,
+            val_split_mode=config.dataset.val_split_mode,
         )
     elif config.dataset.format.lower() == "btech":
         datamodule = BTech(
-            # TODO: Remove config values. IAAALD-211
             root=config.dataset.path,
             category=config.dataset.category,
             image_size=(config.dataset.image_size[0], config.dataset.image_size[1]),
             train_batch_size=config.dataset.train_batch_size,
-            test_batch_size=config.dataset.test_batch_size,
+            eval_batch_size=config.dataset.eval_batch_size,
             num_workers=config.dataset.num_workers,
-            seed=config.project.seed,
             task=config.dataset.task,
             transform_config_train=config.dataset.transform_config.train,
-            transform_config_val=config.dataset.transform_config.val,
-            create_validation_set=config.dataset.create_validation_set,
+            transform_config_eval=config.dataset.transform_config.eval,
+            val_split_mode=config.dataset.val_split_mode,
         )
     elif config.dataset.format.lower() == "folder":
         datamodule = Folder(
@@ -70,14 +66,13 @@ def get_datamodule(config: Union[DictConfig, ListConfig]) -> LightningDataModule
             mask_dir=config.dataset.mask,
             extensions=config.dataset.extensions,
             split_ratio=config.dataset.split_ratio,
-            seed=config.project.seed,
             image_size=(config.dataset.image_size[0], config.dataset.image_size[1]),
             train_batch_size=config.dataset.train_batch_size,
-            test_batch_size=config.dataset.test_batch_size,
+            eval_batch_size=config.dataset.eval_batch_size,
             num_workers=config.dataset.num_workers,
             transform_config_train=config.dataset.transform_config.train,
-            transform_config_val=config.dataset.transform_config.val,
-            create_validation_set=config.dataset.create_validation_set,
+            transform_config_eval=config.dataset.transform_config.eval,
+            val_split_mode=config.dataset.val_split_mode,
         )
     else:
         raise ValueError(
@@ -90,6 +85,8 @@ def get_datamodule(config: Union[DictConfig, ListConfig]) -> LightningDataModule
 
 
 __all__ = [
+    "AnomalibDataset",
+    "AnomalibDataModule",
     "get_datamodule",
     "BTech",
     "Folder",

anomalib/data/base/__init__.py

@@ -0,0 +1,10 @@
+"""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
+
+__all__ = ["AnomalibDataset", "AnomalibDataModule"]

anomalib/data/base/datamodule.py

@@ -0,0 +1,89 @@
+"""Anomalib datamodule base class."""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+from abc import ABC
+from typing import Optional
+
+from pandas import DataFrame
+from pytorch_lightning import LightningDataModule
+from pytorch_lightning.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
+from torch.utils.data import DataLoader
+
+from anomalib.data.base.dataset import AnomalibDataset
+from anomalib.data.utils import ValSplitMode, random_split
+
+logger = logging.getLogger(__name__)
+
+
+class AnomalibDataModule(LightningDataModule, ABC):
+    """Base Anomalib data module.
+
+    Args:
+        train_batch_size (int): Batch size used by the train dataloader.
+        test_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.
+    """
+
+    def __init__(self, train_batch_size: int, eval_batch_size: int, num_workers: int, val_split_mode: ValSplitMode):
+        super().__init__()
+        self.train_batch_size = train_batch_size
+        self.eval_batch_size = eval_batch_size
+        self.num_workers = num_workers
+        self.val_split_mode = val_split_mode
+
+        self.train_data: Optional[AnomalibDataset] = None
+        self.val_data: Optional[AnomalibDataset] = None
+        self.test_data: Optional[AnomalibDataset] = None
+
+        self._samples: Optional[DataFrame] = None
+
+    def setup(self, stage: Optional[str] = None):
+        """Setup train, validation and test data.
+
+        Args:
+          stage: Optional[str]:  Train/Val/Test stages. (Default value = None)
+        """
+        if not self.is_setup:
+            self._setup(stage)
+        assert self.is_setup
+
+    def _setup(self, _stage: Optional[str] = None) -> None:
+        """Set up the datasets and perform dynamic subset splitting.
+
+        May be overridden in subclass for custom splitting behaviour.
+        """
+        assert self.train_data is not None
+        assert self.test_data is not None
+
+        self.train_data.setup()
+        self.test_data.setup()
+        if self.val_split_mode == ValSplitMode.FROM_TEST:
+            self.val_data, self.test_data = random_split(self.test_data, [0.5, 0.5], label_aware=True)
+        elif self.val_split_mode == ValSplitMode.SAME_AS_TEST:
+            self.val_data = self.test_data
+        else:
+            raise ValueError(f"Unknown validation split mode: {self.val_split_mode}")
+
+    @property
+    def is_setup(self):
+        """Checks if setup() has been called."""
+        if self.train_data is None or self.val_data is None or self.test_data is None:
+            return False
+        return self.train_data.is_setup and self.val_data.is_setup and self.test_data.is_setup
+
+    def train_dataloader(self) -> TRAIN_DATALOADERS:
+        """Get train dataloader."""
+        return DataLoader(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(self.val_data, shuffle=False, batch_size=self.eval_batch_size, num_workers=self.num_workers)
+
+    def test_dataloader(self) -> EVAL_DATALOADERS:
+        """Get test dataloader."""
+        return DataLoader(self.test_data, shuffle=False, batch_size=self.eval_batch_size, num_workers=self.num_workers)

anomalib/data/base/dataset.py

@@ -0,0 +1,155 @@
+"""Anomalib dataset base class."""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import copy
+import logging
+from abc import ABC, abstractmethod
+from typing import Dict, Sequence, Union
+
+import cv2
+import numpy as np
+import pandas as pd
+from pandas import DataFrame
+from torch import Tensor
+from torch.utils.data import Dataset
+
+from anomalib.data.utils import read_image
+from anomalib.pre_processing import PreProcessor
+
+logger = logging.getLogger(__name__)
+
+
+class AnomalibDataset(Dataset, ABC):
+    """Anomalib dataset."""
+
+    def __init__(self, task: str, pre_process: PreProcessor):
+        super().__init__()
+        self.task = task
+        self.pre_process = pre_process
+        self._samples = None
+
+    def __len__(self) -> int:
+        """Get length of the dataset."""
+        assert isinstance(self._samples, DataFrame)
+        return len(self._samples)
+
+    def subsample(self, indices: Sequence[int], inplace=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.
+        """
+        dataset = self if inplace else copy.deepcopy(self)
+        dataset.samples = self.samples.iloc[indices].reset_index(drop=True)
+        return dataset
+
+    @property
+    def is_setup(self) -> bool:
+        """Checks if setup() been called."""
+        return isinstance(self._samples, DataFrame)
+
+    @property
+    def samples(self) -> DataFrame:
+        """Get the samples dataframe."""
+        if not self.is_setup:
+            raise RuntimeError("Dataset is not setup yet. Call setup() first.")
+        return self._samples
+
+    @samples.setter
+    def samples(self, samples: DataFrame):
+        """Overwrite the samples with a new dataframe.
+
+        Args:
+            samples (DataFrame): DataFrame with new samples.
+        """
+        self._samples = samples.sort_values(by="image_path", ignore_index=True)
+
+    @property
+    def has_normal(self) -> bool:
+        """Check if the dataset contains any normal samples."""
+        return 0 in list(self.samples.label_index)
+
+    @property
+    def has_anomalous(self) -> bool:
+        """Check if the dataset contains any anomalous samples."""
+        return 1 in list(self.samples.label_index)
+
+    def __getitem__(self, index: int) -> Dict[str, Union[str, Tensor]]:
+        """Get dataset item for the index ``index``.
+
+        Args:
+            index (int): Index to get the item.
+
+        Returns:
+            Union[Dict[str, Tensor], Dict[str, Union[str, Tensor]]]: Dict of image tensor during training.
+                Otherwise, Dict containing image path, target path, image tensor, label and transformed bounding box.
+        """
+        assert isinstance(self._samples, DataFrame)
+
+        image_path = self._samples.iloc[index].image_path
+        image = read_image(image_path)
+        label_index = self._samples.iloc[index].label_index
+
+        item = dict(image_path=image_path, label=label_index)
+
+        if self.task == "classification":
+            pre_processed = self.pre_process(image=image)
+        elif self.task == "segmentation":
+            mask_path = self._samples.iloc[index].mask_path
+
+            # Only Anomalous (1) images have masks in anomaly datasets
+            # Therefore, create empty mask for Normal (0) images.
+            if label_index == 0:
+                mask = np.zeros(shape=image.shape[:2])
+            else:
+                mask = cv2.imread(mask_path, flags=0) / 255.0
+
+            pre_processed = self.pre_process(image=image, mask=mask)
+
+            item["mask_path"] = mask_path
+            item["mask"] = pre_processed["mask"]
+        else:
+            raise ValueError(f"Unknown task type: {self.task}")
+        item["image"] = pre_processed["image"]
+
+        return item
+
+    def __add__(self, other_dataset: AnomalibDataset) -> AnomalibDataset:
+        """Concatenate this dataset with another dataset."""
+        assert isinstance(other_dataset, self.__class__), "Cannot concatenate datasets that are not of the same type."
+        assert self.is_setup and other_dataset.is_setup, "Cannot concatenate uninitialized datasets. Call setup first."
+        dataset = copy.deepcopy(self)
+        dataset.samples = pd.concat([self.samples, other_dataset.samples], ignore_index=True)
+        return dataset
+
+    def setup(self) -> None:
+        """Load data/metadata into memory."""
+        if not self.is_setup:
+            self._setup()
+        assert self.is_setup, "setup() should set self._samples"
+
+    @abstractmethod
+    def _setup(self) -> DataFrame:
+        """Set up the data module.
+
+        This method should return a dataframe that contains the information needed by the dataloader to load each of
+        the dataset items into memory.
+        The dataframe must at least contain the following columns:
+            split: the subset to which the dataset item is assigned.
+            image_path: path to file system location where the image is stored.
+            label_index: index of the anomaly label, typically 0 for "normal" and 1 for "anomalous".
+            mask_path (if task == "segmentation"): path to the ground truth masks (for the anomalous images only).
+
+        Example:
+        |---|-------------------|-----------|-------------|------------------|-------|
+        |   | image_path        | label     | label_index | mask_path        | split |
+        |---|-------------------|-----------|-------------|------------------|-------|
+        | 0 | path/to/image.png | anomalous | 1           | path/to/mask.png | train |
+        |---|-------------------|-----------|-------------|------------------|-------|
+        """
+        raise NotImplementedError