Add min-max normalization

anomalib/core/callbacks/__init__.py

@@ -8,6 +8,7 @@
 from pytorch_lightning.callbacks import Callback, ModelCheckpoint
 
 from .compress import CompressModelCallback
+from .min_max_normalization import MinMaxNormalizationCallback
 from .model_loader import LoadModelCallback
 from .normalization import AnomalyScoreNormalizationCallback
 from .save_to_csv import SaveToCSVCallback
@@ -51,17 +52,22 @@ def get_callbacks(config: Union[ListConfig, DictConfig]) -> List[Callback]:
         load_model = LoadModelCallback(os.path.join(config.project.path, config.model.weight_file))
         callbacks.append(load_model)
 
-    if "normalize_scores" in config.model.keys() and config.model.normalize_scores:
-        if config.model.name in ["padim", "stfpm"]:
-            if not config.optimization.nncf.apply:
-                callbacks.append(AnomalyScoreNormalizationCallback())
+    if "normalization_method" in config.model.keys() and not config.model.normalization_method == "none":
+        if config.model.normalization_method == "cdf":
+            if config.model.name in ["padim", "stfpm"]:
+                if not config.optimization.nncf.apply:
+                    callbacks.append(AnomalyScoreNormalizationCallback())
+                else:
+                    raise NotImplementedError("CDF Score Normalization is currently not compatible with NNCF.")
             else:
-                raise NotImplementedError("Score Normalization is currently not compatible with NNCF.")
+                raise NotImplementedError("Score Normalization is currently supported for PADIM and STFPM only.")
+        elif config.model.normalization_method == "min_max":
+            callbacks.append(MinMaxNormalizationCallback())
         else:
-            raise NotImplementedError("Score Normalization is currently supported for PADIM and STFPM only.")
+            raise ValueError(f"Normalization method not recognized: {config.model.normalization_method}")
 
     if not config.project.log_images_to == []:
-        callbacks.append(VisualizerCallback(inputs_are_normalized=config.model.normalize_scores))
+        callbacks.append(VisualizerCallback(inputs_are_normalized=config.model.normalization_method == "cdf"))
 
     if "optimization" in config.keys():
         if config.optimization.nncf.apply:

anomalib/core/callbacks/min_max_normalization.py

@@ -0,0 +1,65 @@
+"""Anomaly Score Normalization Callback that uses min-max normalization."""
+from typing import Any, Dict
+
+import pytorch_lightning as pl
+from pytorch_lightning import Callback
+from pytorch_lightning.utilities.types import STEP_OUTPUT
+
+
+class MinMaxNormalizationCallback(Callback):
+    """Callback that normalizes the image-level and pixel-level anomaly scores using min-max normalization."""
+
+    def on_test_start(self, _trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        """Called when the test begins."""
+        normalized_image_threshold = (pl_module.image_threshold.value.item() - pl_module.min_max.min) / (
+            pl_module.min_max.max - pl_module.min_max.min
+        )
+        normalized_pixel_threshold = (pl_module.pixel_threshold.value.item() - pl_module.min_max.min) / (
+            pl_module.min_max.max - pl_module.min_max.min
+        )
+        pl_module.image_metrics.F1.threshold = normalized_image_threshold.item()
+        pl_module.pixel_metrics.F1.threshold = normalized_pixel_threshold.item()
+
+    def on_validation_batch_end(
+        self,
+        _trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        outputs: STEP_OUTPUT,
+        _batch: Any,
+        _batch_idx: int,
+        _dataloader_idx: int,
+    ) -> None:
+        """Called when the validation batch ends, update the min and max observed values."""
+        if "anomaly_maps" in outputs.keys():
+            pl_module.min_max(outputs["anomaly_maps"])
+        else:
+            pl_module.min_max(outputs["pred_scores"])
+
+    def on_test_batch_end(
+        self,
+        _trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        outputs: STEP_OUTPUT,
+        _batch: Any,
+        _batch_idx: int,
+        _dataloader_idx: int,
+    ) -> None:
+        """Called when the test batch ends, normalizes the predicted scores and anomaly maps."""
+        self._normalize(outputs, pl_module.min_max)
+
+    def on_predict_batch_end(
+        self,
+        _trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        outputs: Dict,
+        _batch: Any,
+        _batch_idx: int,
+        _dataloader_idx: int,
+    ) -> None:
+        """Called when the predict batch ends, normalizes the predicted scores and anomaly maps."""
+        self._normalize(outputs, pl_module.min_max)
+
+    def _normalize(self, outputs, stats):
+        outputs["pred_scores"] = (outputs["pred_scores"] - stats.min) / (stats.max - stats.min)
+        if "anomaly_maps" in outputs.keys():
+            outputs["anomaly_maps"] = (outputs["anomaly_maps"] - stats.min) / (stats.max - stats.min)

anomalib/core/callbacks/visualizer_callback.py

@@ -95,7 +95,7 @@ def on_test_batch_end(
             threshold = 0.5
             normalize = False  # anomaly maps are already normalized
         else:
-            threshold = pl_module.pixel_threshold.value.item()
+            threshold = pl_module.pixel_metrics.F1.threshold
             normalize = True  # raw anomaly maps. Still need to normalize
 
         for (filename, image, true_mask, anomaly_map) in zip(

anomalib/core/metrics/__init__.py

@@ -2,6 +2,7 @@
 from .adaptive_threshold import AdaptiveThreshold
 from .anomaly_score_distribution import AnomalyScoreDistribution
 from .auroc import AUROC
+from .min_max import MinMax
 from .optimal_f1 import OptimalF1
 
-__all__ = ["AUROC", "OptimalF1", "AdaptiveThreshold", "AnomalyScoreDistribution"]
+__all__ = ["AUROC", "OptimalF1", "AdaptiveThreshold", "AnomalyScoreDistribution", "MinMax"]

anomalib/core/metrics/min_max.py

@@ -0,0 +1,28 @@
+"""Module that tracks the min and max values of the observations in each batch."""
+from typing import Tuple
+
+import torch
+from torch import Tensor
+from torchmetrics import Metric
+
+
+class MinMax(Metric):
+    """Track the min and max values of the observations in each batch."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.add_state("min", torch.tensor(float("inf")), persistent=True)  # pylint: disable=not-callable
+        self.add_state("max", torch.tensor(float("-inf")), persistent=True)  # pylint: disable=not-callable
+
+        self.min = torch.tensor(float("inf"))  # pylint: disable=not-callable
+        self.max = torch.tensor(float("-inf"))  # pylint: disable=not-callable
+
+    # pylint: disable=arguments-differ
+    def update(self, predictions: Tensor) -> None:  # type: ignore
+        """Update the min and max values."""
+        self.max = torch.max(self.max, torch.max(predictions))
+        self.min = torch.min(self.min, torch.min(predictions))
+
+    def compute(self) -> Tuple[Tensor, Tensor]:
+        """Return min and max values."""
+        return self.min, self.max

anomalib/core/model/anomaly_module.py

@@ -22,7 +22,12 @@
 from torch import Tensor, nn
 from torchmetrics import F1, MetricCollection
 
-from anomalib.core.metrics import AUROC, AdaptiveThreshold, AnomalyScoreDistribution
+from anomalib.core.metrics import (
+    AUROC,
+    AdaptiveThreshold,
+    AnomalyScoreDistribution,
+    MinMax,
+)
 
 
 class AnomalyModule(pl.LightningModule):
@@ -47,6 +52,7 @@ def __init__(self, params: Union[DictConfig, ListConfig]):
         self.pixel_threshold = AdaptiveThreshold(self.hparams.model.threshold.pixel_default)
 
         self.training_distribution = AnomalyScoreDistribution()
+        self.min_max = MinMax()
 
         self.model: nn.Module
 
@@ -141,8 +147,8 @@ def _compute_adaptive_threshold(self, outputs):
         else:
             self.pixel_threshold.value = self.image_threshold.value
 
-        self.image_metrics.F1.threshold = self.image_threshold.value
-        self.pixel_metrics.F1.threshold = self.pixel_threshold.value
+        self.image_metrics.F1.threshold = self.image_threshold.value.item()
+        self.pixel_metrics.F1.threshold = self.pixel_threshold.value.item()
 
     def _collect_outputs(self, image_metric, pixel_metric, outputs):
         for output in outputs:
@@ -159,6 +165,7 @@ def _post_process(self, outputs):
 
     def _log_metrics(self):
         """Log computed performance metrics."""
+        print(self.image_metrics.compute())
         self.log_dict(self.image_metrics)
         if self.hparams.dataset.task == "segmentation":
             self.log_dict(self.pixel_metrics)

anomalib/models/dfkde/config.yaml

@@ -18,7 +18,7 @@ model:
   confidence_threshold: 0.5
   pre_processing: scale
   n_components: 16
-  normalize_scores: false # currently not supported for this model
+  normalization_method: min_max # options: [null, min_max, cdf]
   threshold:
     image_default: 0
     adaptive: true

anomalib/models/dfm/config.yaml

@@ -17,7 +17,7 @@ model:
   pca_level: 0.97
   score_type: fre # nll: for Gaussian modeling, fre: pca feature reconstruction error
   project_path: ./results
-  normalize_scores: false # currently not supported for this model
+  normalization_method: min_max # options: [null, min_max, cdf]
   threshold:
     image_default: 0
     adaptive: true

anomalib/models/padim/config.yaml

@@ -3,7 +3,7 @@ dataset:
   format: mvtec
   path: ./datasets/MVTec
   url: ftp://guest:[email protected]/mvtec_anomaly_detection/mvtec_anomaly_detection.tar.xz
-  category: leather
+  category: bottle
   task: segmentation
   label_format: None
   tiling:
@@ -26,7 +26,7 @@ model:
     - layer2
     - layer3
   metric: auc
-  normalize_scores: true
+  normalization_method: min_max # options: [none, min_max, cdf]
   threshold:
     image_default: 3
     pixel_default: 3

anomalib/models/patchcore/config.yaml

@@ -28,7 +28,7 @@ model:
   num_neighbors: 9
   metric: auc
   weight_file: weights/model.ckpt
-  normalize_scores: false # currently not supported for this model
+  normalization_method: min_max # options: [null, min_max, cdf]
   threshold:
     image_default: 0
     pixel_default: 0

anomalib/models/stfpm/config.yaml

@@ -33,7 +33,7 @@ model:
     patience: 3
     metric: pixel_AUROC
     mode: max
-  normalize_scores: false
+  normalization_method: min_max # options: [null, min_max, cdf]
   threshold:
     image_default: 0
     pixel_default: 0

tests/core/callbacks/normalization_callback/test_normalization_callback.py

@@ -22,16 +22,22 @@ def test_normalizer():
     config.dataset.path = get_dataset_path(config.dataset.path)
     config.model.threshold.adaptive = True
 
-    # run with normalization
-    config.model.normalize_scores = True
+    # run without normalization
+    config.model.normalization_method = "none"
     seed_everything(42)
-    results_with_normalization = run_train_test(config)
+    results_without_normalization = run_train_test(config)
+
+    # run with cdf normalization
+    config.model.normalization_method = "cdf"
+    seed_everything(42)
+    results_with_cdf_normalization = run_train_test(config)
 
     # run without normalization
-    config.model.normalize_scores = False
+    config.model.normalization_method = "min_max"
     seed_everything(42)
-    results_without_normalization = run_train_test(config)
+    results_with_minmax_normalization = run_train_test(config)
 
     # performance should be the same
     for metric in ["image_AUROC", "image_F1"]:
-        assert results_without_normalization[0][metric] == results_with_normalization[0][metric]
+        assert results_without_normalization[0][metric] == results_with_cdf_normalization[0][metric]
+        assert results_without_normalization[0][metric] == results_with_minmax_normalization[0][metric]