TorchMetrics

anomalib/core/callbacks/visualizer_callback.py

@@ -2,17 +2,16 @@
 Visualizer Callback
 """
 from pathlib import Path
+from typing import Any, Dict
 from warnings import warn
 
-from pytorch_lightning import Callback, LightningModule, Trainer
+import pytorch_lightning as pl
+from pytorch_lightning import Callback
 from skimage.segmentation import mark_boundaries
-from tqdm import tqdm
 
 from anomalib import loggers
 from anomalib.core.model import AnomalyModule
-from anomalib.core.results import SegmentationResults
 from anomalib.datasets.utils import Denormalize
-from anomalib.utils.metrics import compute_threshold_and_f1_score
 from anomalib.utils.post_process import compute_mask, superimpose_anomaly_map
 from anomalib.utils.visualizer import Visualizer
 
@@ -22,8 +21,8 @@ class VisualizerCallback(Callback):
     Callback that visualizes the inference results of a model. The callback generates a figure showing the original
     image, the ground truth segmentation mask, the predicted error heat map, and the predicted segmentation mask.
 
-    To save the images to the filesystem, add the 'local' keyword to the project.log_images_to parameter in the
-    config.yaml file. .
+    To write the images to the Sigopt logger, add the 'sigopt' keyword to the project.log_images_to parameter in the
+    config.yaml file. To save the images to the filesystem, add the 'local' keyword.
     """
 
     def __init__(self):
@@ -58,33 +57,28 @@ def _add_images(
         if "local" in module.hparams.project.log_images_to:
             visualizer.save(Path(module.hparams.project.path) / "images" / filename.parent.name / filename.name)
 
-    def on_test_epoch_end(self, _trainer: Trainer, pl_module: LightningModule) -> None:
-        """Log images at the end of training
-        Args:
-            _trainer (Trainer): Pytorch lightning trainer object (unused)
-            pl_module (LightningModule): Lightning modules derived from BaseAnomalyLightning object as
-            currently only they support logging images.
+    def on_test_batch_end(
+        self,
+        _trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        outputs: Dict,
+        _batch: Any,
+        _batch_idx: int,
+        _dataloader_idx: int,
+    ) -> None:
+        """
+        Visualize the results for the current batch of images.
         """
 
-        if isinstance(pl_module.results, SegmentationResults):
-            results = pl_module.results
-        else:
-            raise ValueError("Visualizer callback only supported for segmentation tasks.")
-
-        if results.images is None or results.true_masks is None or results.anomaly_maps is None:
-            raise ValueError("Result set cannot be empty!")
-
-        threshold, _ = compute_threshold_and_f1_score(results.true_masks, results.anomaly_maps)
-
-        for (filename, image, true_mask, anomaly_map) in tqdm(
-            zip(results.filenames, results.images, results.true_masks, results.anomaly_maps),
-            desc="Saving Results",
-            total=len(results.filenames),
+        for (filename, image, true_mask, anomaly_map) in zip(
+            outputs["image_path"], outputs["image"], outputs["mask"], outputs["anomaly_maps"]
         ):
-            image = Denormalize()(image)
+            image = Denormalize()(image.cpu())
+            true_mask = true_mask.cpu().numpy()
+            anomaly_map = anomaly_map.cpu().numpy()
 
             heat_map = superimpose_anomaly_map(anomaly_map, image)
-            pred_mask = compute_mask(anomaly_map, threshold)
+            pred_mask = compute_mask(anomaly_map, pl_module.threshold.item())
             vis_img = mark_boundaries(image, pred_mask, color=(1, 0, 0), mode="thick")
 
             visualizer = Visualizer(num_rows=1, num_cols=5, figure_size=(12, 3))
@@ -93,5 +87,5 @@ def on_test_epoch_end(self, _trainer: Trainer, pl_module: LightningModule) -> No
             visualizer.add_image(image=heat_map, title="Predicted Heat Map")
             visualizer.add_image(image=pred_mask, color_map="gray", title="Predicted Mask")
             visualizer.add_image(image=vis_img, title="Segmentation Result")
-            self._add_images(visualizer, pl_module, filename)
+            self._add_images(visualizer, pl_module, Path(filename))
             visualizer.close()

anomalib/core/metrics/__init__.py

@@ -0,0 +1,7 @@
+"""
+Custom anomaly evaluation metrics.
+"""
+from .auroc import AUROC
+from .opt_f1 import OptimalF1
+
+__all__ = ["AUROC", "OptimalF1"]

anomalib/core/metrics/auroc.py

@@ -0,0 +1,22 @@
+"""
+Implementation of AUROC metric based on TorchMetrics.
+"""
+from torch import Tensor
+from torchmetrics import ROC
+from torchmetrics.functional import auc
+
+
+class AUROC(ROC):
+    """
+    Area under the ROC curve.
+    """
+
+    def compute(self) -> Tensor:
+        """
+        First compute ROC curve, then compute area under the curve.
+
+        Returns:
+            Value of the AUROC metric
+        """
+        fpr, tpr, _thresholds = super().compute()
+        return auc(fpr, tpr)

anomalib/core/metrics/opt_f1.py

@@ -0,0 +1,29 @@
+"""
+Implementation of Optimal F1 score based on TorchMetrics.
+"""
+import torch
+from torch import Tensor
+from torchmetrics import PrecisionRecallCurve
+
+
+class OptimalF1(PrecisionRecallCurve):
+    """
+    Compute the optimal F1 score at the adaptive threshold, based on the F1 metric of the
+    true labels and the predicted anomaly scores.
+    """
+
+    threshold: torch.Tensor
+
+    def compute(self) -> Tensor:
+        """
+        Compute the F1 scores while varying the threshold. Store the optimal
+        threshold as attribute and return the maximum value of the F1 score.
+
+        Returns:
+            Value of the F1 score at the optimal threshold.
+        """
+        precision, recall, thresholds = super().compute()
+        f1_score = (2 * precision * recall) / (precision + recall + 1e-10)
+        self.threshold = thresholds[torch.argmax(f1_score)]
+        opt_f1_score = torch.max(f1_score)
+        return opt_f1_score

anomalib/core/model/anomaly_module.py

@@ -23,9 +23,9 @@
 from omegaconf import DictConfig, ListConfig
 from pytorch_lightning.callbacks.base import Callback
 from torch import nn
+from torchmetrics import F1, MetricCollection
 
-from anomalib.core.results import ClassificationResults, SegmentationResults
-from anomalib.utils.metrics import compute_threshold_and_f1_score
+from anomalib.core.metrics import AUROC, OptimalF1
 
 
 class AnomalyModule(pl.LightningModule):
@@ -44,18 +44,21 @@ def __init__(self, params: Union[DictConfig, ListConfig]):
         self.save_hyperparameters(params)
         self.loss: torch.Tensor
         self.callbacks: List[Callback]
-        self.register_buffer("threshold", torch.Tensor([params.model.threshold.default]))
+        self.register_buffer("threshold", torch.tensor(params.model.threshold.default))  # pylint: disable=not-callable
         self.threshold: torch.Tensor
 
         self.model: nn.Module
 
-        self.results: Union[ClassificationResults, SegmentationResults]
-        if params.dataset.task == "classification":
-            self.results = ClassificationResults()
-        elif params.dataset.task == "segmentation":
-            self.results = SegmentationResults()
+        # metrics
+        self.image_metrics = MetricCollection(
+            [AUROC(num_classes=1, pos_label=1, compute_on_step=False)], prefix="image_"
+        )
+        if params.model.threshold.adaptive:
+            self.image_metrics.add_metrics([OptimalF1(num_classes=1, compute_on_step=False)])
         else:
-            raise NotImplementedError("Only Classification and Segmentation tasks are supported in this version.")
+            self.image_metrics.add_metrics([F1(num_classes=1, compute_on_step=False, threshold=self.threshold.item())])
+        if self.hparams.dataset.task == "segmentation":
+            self.pixel_metrics = self.image_metrics.clone(prefix="pixel_")
 
     def forward(self, batch):  # pylint: disable=arguments-differ
         """
@@ -104,15 +107,19 @@ def validation_step_end(self, val_step_outputs):  # pylint: disable=arguments-di
         """
         Called at the end of each validation step.
         """
-        return self._post_process(val_step_outputs)
+        val_step_outputs = self._post_process(val_step_outputs)
+        self.image_metrics(val_step_outputs["pred_scores"], val_step_outputs["label"].int())
+        if self.hparams.dataset.task == "segmentation":
+            self.pixel_metrics(val_step_outputs["anomaly_maps"].flatten(), val_step_outputs["mask"].flatten().int())
+        return val_step_outputs
 
     def test_step_end(self, test_step_outputs):  # pylint: disable=arguments-differ
         """
         Called at the end of each validation step.
         """
-        return self._post_process(test_step_outputs)
+        return self.validation_step_end(test_step_outputs)
 
-    def validation_epoch_end(self, outputs):
+    def validation_epoch_end(self, _outputs):
         """
         Compute image-level performance metrics
 
@@ -121,23 +128,19 @@ def validation_epoch_end(self, outputs):
 
 
         """
-        self.results.store_outputs(outputs)
         if self.hparams.model.threshold.adaptive:
-            threshold, _ = compute_threshold_and_f1_score(self.results.true_labels, self.results.pred_scores)
-            self.threshold = torch.Tensor([threshold])
-        self.results.evaluate(self.threshold.item())
+            self.image_metrics.compute()
+            self.threshold = self.image_metrics.OptimalF1.threshold
         self._log_metrics()
 
-    def test_epoch_end(self, outputs):
+    def test_epoch_end(self, _outputs):
         """
         Compute and save anomaly scores of the test set.
 
         Args:
             outputs: Batch of outputs from the validation step
 
         """
-        self.results.store_outputs(outputs)
-        self.results.evaluate(self.threshold.item())
         self._log_metrics()
 
     def _post_process(self, outputs, predict_labels=False):
@@ -156,5 +159,6 @@ def _log_metrics(self):
         """
         Log computed performance metrics
         """
-        for name, value in self.results.performance.items():
-            self.log(name=name, value=value, on_epoch=True, prog_bar=True)
+        self.log_dict(self.image_metrics)
+        if self.hparams.dataset.task == "segmentation":
+            self.log_dict(self.pixel_metrics)

anomalib/core/model/kde.py

@@ -55,7 +55,7 @@ def forward(self, features: torch.Tensor) -> torch.Tensor:
         """
         features = torch.matmul(features, self.bw_transform)
 
-        estimate = torch.zeros(features.shape[0])
+        estimate = torch.zeros(features.shape[0]).to(features.device)
         for i in range(features.shape[0]):
             embedding = ((self.dataset - features[i]) ** 2).sum(dim=1)
             embedding = torch.exp(-embedding / 2) * self.norm