anomalib/utils/callbacks/visualizer_callback.py

"""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 Any, Optional, cast
from warnings import warn

import pytorch_lightning as pl
from pytorch_lightning import Callback
from pytorch_lightning.utilities.types import STEP_OUTPUT
from skimage.segmentation import mark_boundaries

from anomalib.models.components import AnomalyModule
from anomalib.post_processing import Visualizer, compute_mask, superimpose_anomaly_map
from anomalib.pre_processing.transforms import Denormalize
from anomalib.utils import loggers
from anomalib.utils.loggers import AnomalibWandbLogger
from anomalib.utils.loggers.base import ImageLoggerBase


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.
    """

    def __init__(self, task: str, inputs_are_normalized: bool = True):
        """Visualizer callback."""
        self.task = task
        self.inputs_are_normalized = inputs_are_normalized

    def _add_images(
        self,
        visualizer: Visualizer,
        module: AnomalyModule,
        trainer: pl.Trainer,
        filename: Path,
    ):
        """Save image to logger/local storage.

        Saves the image in `visualizer.figure` to the respective loggers and local storage if specified in
        `log_images_to` in `config.yaml` of the models.

        Args:
            visualizer (Visualizer): Visualizer object from which the `figure` is saved/logged.
            module (AnomalyModule): Anomaly module which holds reference to `hparams`.
            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
        for log_to in module.hparams.project.log_images_to:
            if log_to in loggers.AVAILABLE_LOGGERS:
                # check if logger object is same as the requested object
                if log_to in available_loggers and isinstance(available_loggers[log_to], ImageLoggerBase):
                    logger: ImageLoggerBase = cast(ImageLoggerBase, available_loggers[log_to])  # placate mypy
                    logger.add_image(
                        image=visualizer.figure,
                        name=filename.parent.name + "_" + filename.name,
                        global_step=module.global_step,
                    )
                else:
                    warn(
                        f"Requested {log_to} logging but logger object is of type: {type(module.logger)}."
                        f" Skipping logging to {log_to}"
                    )
            else:
                warn(f"{log_to} not in the list of supported image loggers.")

        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_batch_end(
        self,
        trainer: pl.Trainer,
        pl_module: AnomalyModule,
        outputs: Optional[STEP_OUTPUT],
        _batch: Any,
        _batch_idx: int,
        _dataloader_idx: int,
    ) -> None:
        """Log images at the end of every batch.

        Args:
            trainer (Trainer): Pytorch lightning trainer object (unused).
            pl_module (LightningModule): Lightning modules derived from BaseAnomalyLightning object as
            currently only they support logging images.
            outputs (Dict[str, Any]): Outputs of the current test step.
            _batch (Any): Input batch of the current test step (unused).
            _batch_idx (int): Index of the current test batch (unused).
            _dataloader_idx (int): Index of the dataloader that yielded the current batch (unused).
        """
        assert outputs is not None

        if self.inputs_are_normalized:
            normalize = False  # anomaly maps are already normalized
        else:
            normalize = True  # raw anomaly maps. Still need to normalize
        threshold = pl_module.pixel_metrics.threshold

        for i, (filename, image, anomaly_map, pred_score, gt_label) in enumerate(
            zip(
                outputs["image_path"],
                outputs["image"],
                outputs["anomaly_maps"],
                outputs["pred_scores"],
                outputs["label"],
            )
        ):
            image = Denormalize()(image.cpu())
            anomaly_map = anomaly_map.cpu().numpy()
            heat_map = superimpose_anomaly_map(anomaly_map, image, normalize=normalize)
            pred_mask = compute_mask(anomaly_map, threshold)
            vis_img = mark_boundaries(image, pred_mask, color=(1, 0, 0), mode="thick")

            num_cols = 6 if self.task == "segmentation" else 5
            visualizer = Visualizer(num_rows=1, num_cols=num_cols, figure_size=(12, 3))
            visualizer.add_image(image=image, title="Image")

            if "mask" in outputs:
                true_mask = outputs["mask"][i].cpu().numpy() * 255
                visualizer.add_image(image=true_mask, color_map="gray", title="Ground Truth")

            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")

            image_classified = visualizer.add_text(
                image=image,
                text=f"""Pred: { "anomalous" if pred_score > threshold else "normal"}({pred_score:.3f}) \n
                GT: {"anomalous" if bool(gt_label) else "normal"}""",
            )
            visualizer.add_image(image=image_classified, title="Classified Image")

            self._add_images(visualizer, pl_module, trainer, Path(filename))
            visualizer.close()

    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()