[NEW MODEL] Region-based Anomaly Detection

CHANGELOG.md

@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Add RKDE model implementation (https://github.com/openvinotoolkit/anomalib/pull/821)
 - Add Visual Anomaly (VisA) dataset adapter (<https://github.com/openvinotoolkit/anomalib/pull/824>)
 - Add Synthetic anomalous dataset for validation and testing (https://github.com/openvinotoolkit/anomalib/pull/822)
 - Add Detection task type support (https://github.com/openvinotoolkit/anomalib/pull/822)

anomalib/data/utils/boxes.py

@@ -95,3 +95,18 @@ def boxes_to_anomaly_maps(boxes: Tensor, scores: Tensor, image_size: Tuple[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: Tensor, image_size: torch.Size, new_size: torch.Size) -> Tensor:
+    """Scale bbox coordinates to a new image size.
+
+    Args:
+        boxes (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 = Tensor([*new_size]) / Tensor([*image_size])
+    return boxes * scale.repeat(2).to(boxes.device)

anomalib/models/__init__.py

@@ -22,6 +22,7 @@
 from anomalib.models.padim import Padim
 from anomalib.models.patchcore import Patchcore
 from anomalib.models.reverse_distillation import ReverseDistillation
+from anomalib.models.rkde import Rkde
 from anomalib.models.stfpm import Stfpm
 
 __all__ = [
@@ -35,6 +36,7 @@
     "Padim",
     "Patchcore",
     "ReverseDistillation",
+    "Rkde",
     "Stfpm",
 ]
 
@@ -84,6 +86,7 @@ def get_model(config: Union[DictConfig, ListConfig]) -> AnomalyModule:
         "padim",
         "patchcore",
         "reverse_distillation",
+        "rkde",
         "stfpm",
     ]
     model: AnomalyModule

anomalib/models/components/classification/__init__.py

@@ -0,0 +1,5 @@
+"""Classification modules."""
+
+from .kde_classifier import FeatureScalingMethod, KDEClassifier
+
+__all__ = ["KDEClassifier", "FeatureScalingMethod"]

anomalib/models/components/classification/kde_classifier.py

@@ -0,0 +1,162 @@
+"""Kernel Density Estimation Classifier."""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import random
+from enum import Enum
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor, nn
+
+from anomalib.models.components import PCA, GaussianKDE
+
+logger = logging.getLogger(__name__)
+
+
+class FeatureScalingMethod(str, Enum):
+    """Determines how the feature embeddings are scaled."""
+
+    NORM = "norm"  # scale to unit vector length
+    SCALE = "scale"  # scale to max length observed in training (preserve relative magnitude)
+
+
+class KDEClassifier(nn.Module):
+    """Classification module for KDE-based anomaly detection.
+
+    Args:
+        n_pca_components (int, optional): Number of PCA components. Defaults to 16.
+        feature_scaling_method (FeatureScalingMethod, optional): Scaling method applied to features before passing to
+            KDE. Options are `norm` (normalize to unit vector length) and `scale` (scale to max length observed in
+            training).
+        max_training_points (int, optional): Maximum number of training points to fit the KDE model. Defaults to 40000.
+    """
+
+    def __init__(
+        self,
+        n_pca_components: int = 16,
+        feature_scaling_method: FeatureScalingMethod = FeatureScalingMethod.SCALE,
+        max_training_points: int = 40000,
+    ) -> None:
+        super().__init__()
+
+        self.n_pca_components = n_pca_components
+        self.feature_scaling_method = feature_scaling_method
+        self.max_training_points = max_training_points
+
+        self.pca_model = PCA(n_components=self.n_pca_components)
+        self.kde_model = GaussianKDE()
+
+        self.register_buffer("max_length", torch.empty([]))
+        self.max_length = torch.empty([])
+
+    def pre_process(self, feature_stack: Tensor, max_length: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]:
+        """Pre-process the CNN features.
+
+        Args:
+          feature_stack (Tensor): Features extracted from CNN
+          max_length (Optional[Tensor]): Used to unit normalize the feature_stack vector. If ``max_len`` is not
+            provided, the length is calculated from the ``feature_stack``. Defaults to None.
+
+        Returns:
+            (Tuple): Stacked features and length
+        """
+
+        if max_length is None:
+            max_length = torch.max(torch.linalg.norm(feature_stack, ord=2, dim=1))
+
+        if self.feature_scaling_method == FeatureScalingMethod.NORM:
+            feature_stack /= torch.linalg.norm(feature_stack, ord=2, dim=1)[:, None]
+        elif self.feature_scaling_method == FeatureScalingMethod.SCALE:
+            feature_stack /= max_length
+        else:
+            raise RuntimeError("Unknown pre-processing mode. Available modes are: Normalized and Scale.")
+        return feature_stack, max_length
+
+    def fit(self, embeddings: Tensor) -> bool:
+        """Fit a kde model to embeddings.
+
+        Args:
+            embeddings (Tensor): Input embeddings to fit the model.
+
+        Returns:
+            Boolean confirming whether the training is successful.
+        """
+
+        if embeddings.shape[0] < self.n_pca_components:
+            logger.info("Not enough features to commit. Not making a model.")
+            return False
+
+        # if max training points is non-zero and smaller than number of staged features, select random subset
+        if embeddings.shape[0] > self.max_training_points:
+            selected_idx = torch.tensor(random.sample(range(embeddings.shape[0]), self.max_training_points))
+            selected_features = embeddings[selected_idx]
+        else:
+            selected_features = embeddings
+
+        feature_stack = self.pca_model.fit_transform(selected_features)
+        feature_stack, max_length = self.pre_process(feature_stack)
+        self.max_length = max_length
+        self.kde_model.fit(feature_stack)
+
+        return True
+
+    def compute_kde_scores(self, features: Tensor, as_log_likelihood: Optional[bool] = False) -> Tensor:
+        """Compute the KDE scores.
+
+        The scores calculated from the KDE model are converted to densities. If `as_log_likelihood` is set to true then
+            the log of the scores are calculated.
+
+        Args:
+            features (Tensor): Features to which the PCA model is fit.
+            as_log_likelihood (Optional[bool], optional): If true, gets log likelihood scores. Defaults to False.
+
+        Returns:
+            (Tensor): Score
+        """
+
+        features = self.pca_model.transform(features)
+        features, _ = self.pre_process(features, self.max_length)
+        # Scores are always assumed to be passed as a density
+        kde_scores = self.kde_model(features)
+
+        # add small constant to avoid zero division in log computation
+        kde_scores += 1e-300
+
+        if as_log_likelihood:
+            kde_scores = torch.log(kde_scores)
+
+        return kde_scores
+
+    @staticmethod
+    def compute_probabilities(scores: Tensor) -> Tensor:
+        """Converts density scores to anomaly probabilities (see https://www.desmos.com/calculator/ifju7eesg7).
+
+        Args:
+          scores (Tensor): density of an image.
+
+        Returns:
+          probability that image with {density} is anomalous
+        """
+        return 1 / (1 + torch.exp(0.05 * (scores - 12)))
+
+    def predict(self, features: Tensor) -> Tensor:
+        """Predicts the probability that the features belong to the anomalous class.
+
+        Args:
+          features (Tensor): Feature from which the output probabilities are detected.
+
+        Returns:
+          Detection probabilities
+        """
+
+        scores = self.compute_kde_scores(features, as_log_likelihood=True)
+        probabilities = self.compute_probabilities(scores)
+
+        return probabilities
+
+    def forward(self, features: Tensor) -> Tensor:
+        """Make predictions on extracted features."""
+        return self.predict(features)

anomalib/models/dfkde/config.yaml

@@ -20,16 +20,17 @@ dataset:
 
 model:
   name: dfkde
+  # feature extraction params
   backbone: resnet18
+  layers:
+    - layer4
   pre_trained: true
+  # density estimation params
+  n_pca_components: 16
   max_training_points: 40000
-  pre_processing: scale
-  n_components: 16
-  threshold_steepness: 0.05
-  threshold_offset: 12
+  feature_scaling_method: scale # Determines how the feature embeddings are scaled. Options: [scale, norm]
+  # generic params
   normalization_method: min_max # options: [null, min_max, cdf]
-  layers:
-    - layer4
 
 metrics:
   image:

anomalib/models/dfkde/lightning_model.py

@@ -6,11 +6,13 @@
 import logging
 from typing import List, Union
 
+import torch
 from omegaconf import DictConfig, ListConfig
 from pytorch_lightning.utilities.cli import MODEL_REGISTRY
 from torch import Tensor
 
 from anomalib.models.components import AnomalyModule
+from anomalib.models.components.classification import FeatureScalingMethod
 
 from .torch_model import DfkdeModel
 
@@ -39,23 +41,19 @@ def __init__(
         layers: List[str],
         backbone: str,
         pre_trained: bool = True,
+        n_pca_components: int = 16,
+        feature_scaling_method: FeatureScalingMethod = FeatureScalingMethod.SCALE,
         max_training_points: int = 40000,
-        pre_processing: str = "scale",
-        n_components: int = 16,
-        threshold_steepness: float = 0.05,
-        threshold_offset: int = 12,
     ):
         super().__init__()
 
         self.model = DfkdeModel(
             layers=layers,
             backbone=backbone,
             pre_trained=pre_trained,
-            n_comps=n_components,
-            pre_processing=pre_processing,
-            filter_count=max_training_points,
-            threshold_steepness=threshold_steepness,
-            threshold_offset=threshold_offset,
+            n_pca_components=n_pca_components,
+            feature_scaling_method=feature_scaling_method,
+            max_training_points=max_training_points,
         )
 
         self.embeddings: List[Tensor] = []
@@ -76,7 +74,7 @@ def training_step(self, batch, _batch_idx):  # pylint: disable=arguments-differ
           Deep CNN features.
         """
 
-        embedding = self.model.get_features(batch["image"]).squeeze()
+        embedding = self.model(batch["image"])
 
         # NOTE: `self.embedding` appends each batch embedding to
         #   store the training set embedding. We manually append these
@@ -89,8 +87,10 @@ def on_validation_start(self) -> None:
         # NOTE: Previous anomalib versions fit Gaussian at the end of the epoch.
         #   This is not possible anymore with PyTorch Lightning v1.4.0 since validation
         #   is run within train epoch.
+        embeddings = torch.vstack(self.embeddings)
+
         logger.info("Fitting a KDE model to the embedding collected from the training set.")
-        self.model.fit(self.embeddings)
+        self.model.classifier.fit(embeddings)
 
     def validation_step(self, batch, _):  # pylint: disable=arguments-differ
         """Validation Step of DFKDE.
@@ -120,11 +120,9 @@ def __init__(self, hparams: Union[DictConfig, ListConfig]) -> None:
             layers=hparams.model.layers,
             backbone=hparams.model.backbone,
             pre_trained=hparams.model.pre_trained,
+            n_pca_components=hparams.model.n_pca_components,
+            feature_scaling_method=FeatureScalingMethod(hparams.model.feature_scaling_method),
             max_training_points=hparams.model.max_training_points,
-            pre_processing=hparams.model.pre_processing,
-            n_components=hparams.model.n_components,
-            threshold_steepness=hparams.model.threshold_steepness,
-            threshold_offset=hparams.model.threshold_offset,
         )
         self.hparams: Union[DictConfig, ListConfig]  # type: ignore
         self.save_hyperparameters(hparams)