🐞 Fix visualizer for classification, mode=simple (#442)

* Fix visualizer for `classification`, `mode=simple`

Need to refer to `image_result.image` instead of `image_result.heat_map`
since AD-only methods don't produce heat-maps

* Add Integration test for AnomalyModule/Visualizer

Check that the combination of Visualizer in
[°segmentation", "classification"], ["full", "simple"] works with both
AD-only and segmentation-type models

* Simplify code

anomalib/post_processing/visualizer.py

@@ -158,9 +158,9 @@ def _visualize_simple(self, image_result):
             return (visualization * 255).astype(np.uint8)
         if self.task == "classification":
             if image_result.pred_label:
-                image_classified = add_anomalous_label(image_result.heat_map, image_result.pred_score)
+                image_classified = add_anomalous_label(image_result.image, image_result.pred_score)
             else:
-                image_classified = add_normal_label(image_result.heat_map, 1 - image_result.pred_score)
+                image_classified = add_normal_label(image_result.image, 1 - image_result.pred_score)
             return image_classified
         raise ValueError(f"Unknown task type: {self.task}")
 

tests/helpers/model.py

@@ -15,7 +15,7 @@
 # and limitations under the License.
 
 import os
-from typing import Dict, List, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import numpy as np
 from omegaconf import DictConfig, ListConfig
@@ -36,9 +36,10 @@ def setup_model_train(
     project_path: str,
     nncf: bool,
     category: str,
-    score_type: str = None,
-    weight_file: str = "weights/model.ckpt",
+    score_type: Optional[str] = None,
     fast_run: bool = False,
+    dataset_task: Optional[str] = None,
+    visualizer_mode: Optional[str] = None,
     device: Union[List[int], int] = [0],
 ) -> Tuple[Union[DictConfig, ListConfig], LightningDataModule, AnomalyModule, Trainer]:
     """Train the model based on the parameters passed.
@@ -50,9 +51,12 @@ def setup_model_train(
         nncf (bool): Add nncf callback.
         category (str): Category to train on.
         score_type (str, optional): Only used for DFM. Defaults to None.
-        weight_file (str, optional): Path to weight file.
         fast_run (bool, optional): If set to true, the model trains for only 1 epoch. We train for one epoch as
             this ensures that both anomalous and non-anomalous images are present in the validation step.
+        dataset_task (str, optional): Specify the type of task. Must be in ["classification", "segmentation"].
+            Used for integration testing of model / task / visualizer_mode.
+        visualizer_mode (str, optional): Specify the type of visualization. Must be in ["full", "simple"].
+            Used for integration testing of model / task / visualizer_mode.
         device (List[int], int, optional): Select which device you want to train the model on. Defaults to first GPU.
 
     Returns:
@@ -67,6 +71,13 @@ def setup_model_train(
     config.project.log_images_to = []
     config.trainer.devices = device
     config.trainer.accelerator = "gpu" if device != 0 else "cpu"
+    if dataset_task is not None:
+        config.dataset.task = dataset_task
+    if visualizer_mode is not None:
+        config.visualization.mode = visualizer_mode
+        config.visualization.save_images = True  # Enforce processing by Visualizer
+        if "pixel" in config.metrics and dataset_task == "classification":
+            del config.metrics.pixel
 
     # Remove legacy flags
     for legacy_device in ["num_processes", "gpus", "ipus", "tpu_cores"]:

tests/pre_merge/post_processing/test_visualizer.py

@@ -14,10 +14,15 @@
 # See the License for the specific language governing permissions
 # and limitations under the License.
 
+import tempfile
+
 import numpy as np
+import pytest
 from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
 
 from anomalib.post_processing.visualizer import ImageGrid
+from tests.helpers.dataset import TestDataset
+from tests.helpers.model import setup_model_train
 
 
 def test_visualize_fully_defected_masks():
@@ -36,3 +41,31 @@ def test_visualize_fully_defected_masks():
 
     # assert that the plotted image is completely white
     assert np.all(plotted_img[0][..., 0] == 255)
+
+
+class TestVisualizer:
+    @pytest.mark.parametrize(
+        ["model_name", "nncf"],
+        [
+            ("padim", False),
+            ("ganomaly", False),
+        ],
+    )
+    @pytest.mark.parametrize("task", ("classification", "segmentation"))
+    @pytest.mark.parametrize("mode", ("full", "simple"))
+    @TestDataset(num_train=20, num_test=10)
+    def test_model_visualizer_mode(self, model_name, nncf, task, mode, category="shapes", path=""):
+        """Test combination of model/visualizer/mode on only 1 epoch as a sanity check before merge."""
+        with tempfile.TemporaryDirectory() as project_path:
+            # Train test
+            datamodule, model, trainer = setup_model_train(
+                model_name,
+                dataset_path=path,
+                project_path=project_path,
+                nncf=nncf,
+                category=category,
+                fast_run=True,
+                dataset_task=task,
+                visualizer_mode=mode,
+            )[1:]
+            trainer.test(model=model, datamodule=datamodule)