Add metric visualizations

anomalib/models/components/base/anomaly_module.py

@@ -159,7 +159,7 @@ def _collect_outputs(self, image_metric, pixel_metric, outputs):
             image_metric.update(output["pred_scores"], output["label"].int())
             if "mask" in output.keys() and "anomaly_maps" in output.keys():
                 pixel_metric.cpu()
-                pixel_metric.update(output["anomaly_maps"].flatten(), output["mask"].flatten().int())
+                pixel_metric.update(output["anomaly_maps"], output["mask"].int())
 
     def _post_process(self, outputs):
         """Compute labels based on model predictions."""

anomalib/post_processing/visualizer.py

@@ -235,4 +235,5 @@ def generate(self) -> np.ndarray:
         # convert canvas to numpy array to prepare for visualization with opencv
         img = np.frombuffer(self.figure.canvas.tostring_rgb(), dtype=np.uint8)
         img = img.reshape(self.figure.canvas.get_width_height()[::-1] + (3,))
+        plt.close(self.figure)
         return img

anomalib/utils/callbacks/__init__.py

@@ -31,7 +31,7 @@
 from .model_loader import LoadModelCallback
 from .tiler_configuration import TilerConfigurationCallback
 from .timer import TimerCallback
-from .visualizer_callback import VisualizerCallback
+from .visualizer import VisualizerCallbackImage, VisualizerCallbackMetric
 
 __all__ = [
     "CdfNormalizationCallback",
@@ -40,7 +40,8 @@
     "MinMaxNormalizationCallback",
     "TilerConfigurationCallback",
     "TimerCallback",
-    "VisualizerCallback",
+    "VisualizerCallbackImage",
+    "VisualizerCallbackMetric",
 ]
 
 
@@ -174,14 +175,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/__init__.py

@@ -0,0 +1,6 @@
+"""Callbacks to visualize anomaly detection performance."""
+from .visualizer_base import VisualizerCallbackBase
+from .visualizer_image import VisualizerCallbackImage
+from .visualizer_metric import VisualizerCallbackMetric
+
+__all__ = ["VisualizerCallbackBase", "VisualizerCallbackImage", "VisualizerCallbackMetric"]

anomalib/utils/callbacks/visualizer/visualizer_base.py

@@ -0,0 +1,109 @@
+"""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.post_processing import Visualizer
+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)
+
+        self.visualizer = Visualizer(mode, task)
+
+    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/visualizer_image.py

@@ -1,4 +1,4 @@
-"""Visualizer Callback."""
+"""Image Visualizer Callback."""
 
 # Copyright (C) 2020 Intel Corporation
 #
@@ -15,22 +15,19 @@
 # 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,
@@ -40,62 +37,6 @@ class VisualizerCallback(Callback):
     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)
-
-        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 +96,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/visualizer_metric.py

@@ -0,0 +1,65 @@
+"""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.
+
+from pathlib import Path
+
+import numpy as np
+import pytorch_lightning as pl
+from matplotlib import pyplot as plt
+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 metrics contained in pl_module.
+
+        In order to also plot custom metrics, they need to have implemented a `generate_figure` function that returns
+        Tuple[matplotlib.figure.Figure, str].
+
+        Args:
+            trainer (pl.Trainer): pytorch lightning trainer.
+            pl_module (AnomalyModule): pytorch lightning module.
+        """
+
+        if self.save_images or self.log_images:
+            for metrics in (pl_module.image_metrics, pl_module.pixel_metrics):
+                for metric in metrics.values():
+                    # `generate_figure` needs to be defined for every metric that should be plotted automatically
+                    if hasattr(metric, "generate_figure"):
+                        fig, log_name = metric.generate_figure()
+                        file_name = f"{metrics.prefix}{log_name}"
+                        if self.log_images:
+                            self._add_to_logger(fig, pl_module, trainer, file_name)
+
+                        if self.save_images:
+                            fig.canvas.draw()
+                            # convert figure to np.ndarray for saving via visualizer
+                            img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+                            img = img.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+                            self.visualizer.save(Path(self.image_save_path.joinpath(f"{file_name}.png")), img)
+                        plt.close(fig)
+        super().on_test_end(trainer, pl_module)

anomalib/utils/cli/cli.py

@@ -26,7 +26,7 @@
     ModelCheckpoint,
     TilerConfigurationCallback,
     TimerCallback,
-    VisualizerCallback,
+    VisualizerCallbackImage,
 )
 from anomalib.utils.loggers import configure_logger
 
@@ -115,7 +115,7 @@ def add_arguments_to_parser(self, parser: LightningArgumentParser) -> None:
             }
         )
 
-        parser.add_lightning_class_args(VisualizerCallback, "visualization")  # type: ignore
+        parser.add_lightning_class_args(VisualizerCallbackImage, "visualization")  # type: ignore
         parser.set_defaults(
             {
                 "visualization.mode": "full",

anomalib/utils/loggers/tensorboard.py

@@ -92,13 +92,8 @@ def add_image(self, image: Union[np.ndarray, Figure], name: Optional[str] = None
         if "global_step" not in kwargs:
             raise ValueError("`global_step` is required for tensorboard logger")
 
-        # Matplotlib Figure is not supported by tensorboard
+        # Need to call different functions of `SummaryWriter` for  Figure vs np.ndarray
         if isinstance(image, Figure):
-            axis = image.gca()
-            axis.axis("off")
-            axis.margins(0)
-            image.canvas.draw()  # cache the renderer
-            buffer = np.frombuffer(image.canvas.tostring_rgb(), dtype=np.uint8)
-            image = buffer.reshape(image.canvas.get_width_height()[::-1] + (3,))
-
-        self.experiment.add_image(img_tensor=image, tag=name, dataformats="HWC", **kwargs)
+            self.experiment.add_figure(figure=image, tag=name, close=False, **kwargs)
+        else:
+            self.experiment.add_image(img_tensor=image, tag=name, dataformats="HWC", **kwargs)

anomalib/utils/loggers/wandb.py

@@ -127,3 +127,4 @@ def save(self) -> None:
         super().save()
         if len(self.image_list) > 1:
             wandb.log({"Predictions": self.image_list})
+            self.image_list = []