Add metric visualizations

anomalib/utils/callbacks/__init__.py

@@ -31,7 +31,8 @@
 from .model_loader import LoadModelCallback
 from .tiler_configuration import TilerConfigurationCallback
 from .timer import TimerCallback
-from .visualizer_callback import VisualizerCallback
+from .visualizer_image import VisualizerCallbackImage
+from .visualizer_metric import VisualizerCallbackMetric
 
 __all__ = [
     "CdfNormalizationCallback",
@@ -40,7 +41,8 @@
     "MinMaxNormalizationCallback",
     "TilerConfigurationCallback",
     "TimerCallback",
-    "VisualizerCallback",
+    "VisualizerCallbackImage",
+    "VisualizerCallbackMetric",
 ]
 
 
@@ -174,14 +176,15 @@ def add_visualizer_callback(callbacks: List[Callback], config: Union[DictConfig,
             if config.visualization.image_save_path
             else config.project.path + "/images"
         )
-        callbacks.append(
-            VisualizerCallback(
-                task=config.dataset.task,
-                mode=config.visualization.mode,
-                image_save_path=image_save_path,
-                inputs_are_normalized=not config.model.normalization_method == "none",
-                show_images=config.visualization.show_images,
-                log_images=config.visualization.log_images,
-                save_images=config.visualization.save_images,
+        for callback in (VisualizerCallbackImage, VisualizerCallbackMetric):
+            callbacks.append(
+                callback(
+                    task=config.dataset.task,
+                    mode=config.visualization.mode,
+                    image_save_path=image_save_path,
+                    inputs_are_normalized=not config.model.normalization_method == "none",
+                    show_images=config.visualization.show_images,
+                    log_images=config.visualization.log_images,
+                    save_images=config.visualization.save_images,
+                )
             )
-        )

anomalib/utils/callbacks/visualizer_base.py

@@ -0,0 +1,106 @@
+"""Base Visualizer Callback."""
+
+# Copyright (C) 2020 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions
+# and limitations under the License.
+
+from pathlib import Path
+from typing import Union, cast
+
+import numpy as np
+import pytorch_lightning as pl
+from pytorch_lightning import Callback
+
+from anomalib.models.components import AnomalyModule
+from anomalib.utils.loggers import AnomalibWandbLogger
+from anomalib.utils.loggers.base import ImageLoggerBase
+
+
+class VisualizerCallbackBase(Callback):
+    """Callback that visualizes the results of a model.
+
+    To save the images to the filesystem, add the 'local' keyword to the `project.log_images_to` parameter in the
+    config.yaml file.
+    """
+
+    def __init__(
+        self,
+        task: str,
+        mode: str,
+        image_save_path: str,
+        inputs_are_normalized: bool = True,
+        show_images: bool = False,
+        log_images: bool = True,
+        save_images: bool = True,
+    ):
+        """Visualizer callback."""
+        if mode not in ["full", "simple"]:
+            raise ValueError(f"Unknown visualization mode: {mode}. Please choose one of ['full', 'simple']")
+        self.mode = mode
+        if task not in ["classification", "segmentation"]:
+            raise ValueError(f"Unknown task type: {mode}. Please choose one of ['classification', 'segmentation']")
+        self.task = task
+        self.inputs_are_normalized = inputs_are_normalized
+        self.show_images = show_images
+        self.log_images = log_images
+        self.save_images = save_images
+        self.image_save_path = Path(image_save_path)
+
+    def _add_to_logger(
+        self,
+        image: np.ndarray,
+        module: AnomalyModule,
+        trainer: pl.Trainer,
+        filename: Union[Path, str],
+    ):
+        """Log image from a visualizer to each of the available loggers in the project.
+
+        Args:
+            image (np.ndarray): Image that should be added to the loggers.
+            module (AnomalyModule): Anomaly module.
+            trainer (Trainer): Pytorch Lightning trainer which holds reference to `logger`
+            filename (Path): Path of the input image. This name is used as name for the generated image.
+        """
+        # Store names of logger and the logger in a dict
+        available_loggers = {
+            type(logger).__name__.lower().rstrip("logger").lstrip("anomalib"): logger for logger in trainer.loggers
+        }
+        # save image to respective logger
+        if self.log_images:
+            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
+                    if isinstance(filename, Path):
+                        _name = filename.parent.name + "_" + filename.name
+                    elif isinstance(filename, str):
+                        _name = filename
+                    logger.add_image(
+                        image=image,
+                        name=_name,
+                        global_step=module.global_step,
+                    )
+
+    def on_test_end(self, trainer: pl.Trainer, pl_module: AnomalyModule) -> None:
+        """Sync logs.
+
+        Currently only ``AnomalibWandbLogger.save`` is called from this method.
+        This is because logging as a single batch ensures that all images appear as part of the same step.
+
+        Args:
+            trainer (pl.Trainer): Pytorch Lightning trainer
+            pl_module (AnomalyModule): Anomaly module (unused)
+        """
+        for logger in trainer.loggers:
+            if isinstance(logger, AnomalibWandbLogger):
+                logger.save()

anomalib/utils/callbacks/visualizer_callback.py → anomalib/utils/callbacks/visualizer_image.py

@@ -1,4 +1,4 @@
-"""Visualizer Callback."""
+"""Image Visualizer Callback."""
 
 # Copyright (C) 2020 Intel Corporation
 #
@@ -15,22 +15,20 @@
 # and limitations under the License.
 
 from pathlib import Path
-from typing import Any, Optional, cast
+from typing import Any, Optional
 
-import numpy as np
 import pytorch_lightning as pl
-from pytorch_lightning import Callback
 from pytorch_lightning.utilities.cli import CALLBACK_REGISTRY
 from pytorch_lightning.utilities.types import STEP_OUTPUT
 
 from anomalib.models.components import AnomalyModule
 from anomalib.post_processing import Visualizer
-from anomalib.utils.loggers import AnomalibWandbLogger
-from anomalib.utils.loggers.base import ImageLoggerBase
+
+from .visualizer_base import VisualizerCallbackBase
 
 
 @CALLBACK_REGISTRY
-class VisualizerCallback(Callback):
+class VisualizerCallbackImage(VisualizerCallbackBase):
     """Callback that visualizes the inference results of a model.
 
     The callback generates a figure showing the original image, the ground truth segmentation mask,
@@ -51,51 +49,18 @@ def __init__(
         save_images: bool = True,
     ):
         """Visualizer callback."""
-        if mode not in ["full", "simple"]:
-            raise ValueError(f"Unknown visualization mode: {mode}. Please choose one of ['full', 'simple']")
-        self.mode = mode
-        if task not in ["classification", "segmentation"]:
-            raise ValueError(f"Unknown task type: {mode}. Please choose one of ['classification', 'segmentation']")
-        self.task = task
-        self.inputs_are_normalized = inputs_are_normalized
-        self.show_images = show_images
-        self.log_images = log_images
-        self.save_images = save_images
-        self.image_save_path = Path(image_save_path)
+        super().__init__(
+            task=task,
+            mode=mode,
+            image_save_path=image_save_path,
+            inputs_are_normalized=inputs_are_normalized,
+            show_images=show_images,
+            log_images=log_images,
+            save_images=save_images,
+        )
 
         self.visualizer = Visualizer(mode, task)
 
-    def _add_to_logger(
-        self,
-        image: np.ndarray,
-        module: AnomalyModule,
-        trainer: pl.Trainer,
-        filename: Path,
-    ):
-        """Log image from a visualizer to each of the available loggers in the project.
-
-        Args:
-            image (np.ndarray): Image that should be added to the loggers.
-            module (AnomalyModule): Anomaly module.
-            trainer (Trainer): Pytorch Lightning trainer which holds reference to `logger`
-            filename (Path): Path of the input image. This name is used as name for the generated image.
-        """
-        # Store names of logger and the logger in a dict
-        available_loggers = {
-            type(logger).__name__.lower().rstrip("logger").lstrip("anomalib"): logger for logger in trainer.loggers
-        }
-        # save image to respective logger
-        if self.log_images:
-            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
-                    logger.add_image(
-                        image=image,
-                        name=filename.parent.name + "_" + filename.name,
-                        global_step=module.global_step,
-                    )
-
     def on_predict_batch_end(
         self,
         _trainer: pl.Trainer,
@@ -155,16 +120,3 @@ def on_test_batch_end(
                 self._add_to_logger(image, pl_module, trainer, filename)
             if self.show_images:
                 self.visualizer.show(str(filename), image)
-
-    def on_test_end(self, _trainer: pl.Trainer, pl_module: AnomalyModule) -> None:
-        """Sync logs.
-
-        Currently only ``AnomalibWandbLogger`` is called from this method. This is because logging as a single batch
-        ensures that all images appear as part of the same step.
-
-        Args:
-            _trainer (pl.Trainer): Pytorch Lightning trainer (unused)
-            pl_module (AnomalyModule): Anomaly module
-        """
-        if pl_module.logger is not None and isinstance(pl_module.logger, AnomalibWandbLogger):
-            pl_module.logger.save()

anomalib/utils/callbacks/visualizer_metric.py

@@ -0,0 +1,40 @@
+"""Metric Visualizer Callback."""
+
+# Copyright (C) 2020 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions
+# and limitations under the License.
+
+import pytorch_lightning as pl
+from pytorch_lightning.utilities.cli import CALLBACK_REGISTRY
+
+from anomalib.models.components import AnomalyModule
+
+from .visualizer_base import VisualizerCallbackBase
+
+
+@CALLBACK_REGISTRY
+class VisualizerCallbackMetric(VisualizerCallbackBase):
+    """Callback that visualizes the metric results of a model by plotting the corresponding curves.
+
+    To save the images to the filesystem, add the 'local' keyword to the `project.log_images_to` parameter in the
+    config.yaml file.
+    """
+
+    def on_test_end(self, trainer: pl.Trainer, pl_module: AnomalyModule) -> None:
+        """Log images of the metric scoires for all appropriate.
+
+        Args:
+            trainer (pl.Trainer): pytorch lightning trainer.
+            pl_module (AnomalyModule): pytorch lightning module.
+        """
+        super().on_batch_end(trainer, pl_module)