Visualizer show classification and segmentation

anomalib/post_processing/visualizer.py

@@ -64,6 +64,24 @@ def add_image(self, image: np.ndarray, title: str, color_map: Optional[str] = No
         self.axis[index].imshow(image, color_map)
         self.axis[index].title.set_text(title)
 
+    def add_text(self, image: np.ndarray, text: str, font: int = cv2.FONT_HERSHEY_PLAIN):
+        """Puts text on an image.
+
+        Args:
+            image (np.ndarray): Input image.
+            text (str): Text to add.
+            font (Optional[int]): cv2 font type. Defaults to 0.
+
+        Returns:
+            np.ndarray: Image with text.
+        """
+        image = image.copy()
+        font_size = image.shape[1] // 1024 + 1  # Text scale is calculated based on the reference size of 1024
+        (text_w, text_h), baseline = cv2.getTextSize(text, font, font_size, thickness=font_size // 2)
+        cv2.rectangle(image, (0, 0), (0 + text_w, 0 + text_h), (255, 255, 255), -1)
+        cv2.putText(image, text, (0, baseline // 2 + text_h), font, font_size, 0)
+        return image
+
     def show(self):
         """Show image on a matplotlib figure."""
         self.figure.show()

anomalib/utils/callbacks/__init__.py

@@ -81,7 +81,11 @@ def get_callbacks(config: Union[ListConfig, DictConfig]) -> List[Callback]:
             raise ValueError(f"Normalization method not recognized: {config.model.normalization_method}")
 
     if not config.project.log_images_to == []:
-        callbacks.append(VisualizerCallback(inputs_are_normalized=not config.model.normalization_method == "none"))
+        callbacks.append(
+            VisualizerCallback(
+                task=config.dataset.task, inputs_are_normalized=not config.model.normalization_method == "none"
+            )
+        )
 
     if "optimization" in config.keys():
         if config.optimization.nncf.apply:

anomalib/utils/callbacks/visualizer_callback.py

@@ -40,8 +40,9 @@ class VisualizerCallback(Callback):
     config.yaml file.
     """
 
-    def __init__(self, inputs_are_normalized: bool = True):
+    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(
@@ -111,23 +112,31 @@ def on_test_batch_end(
             normalize = True  # raw anomaly maps. Still need to normalize
         threshold = pl_module.pixel_metrics.F1.threshold
 
-        for (filename, image, true_mask, anomaly_map) in zip(
-            outputs["image_path"], outputs["image"], outputs["mask"], outputs["anomaly_maps"]
+        for i, (filename, image, anomaly_map) in enumerate(
+            zip(outputs["image_path"], outputs["image"], outputs["anomaly_maps"])
         ):
             image = Denormalize()(image.cpu())
-            true_mask = true_mask.cpu().numpy()
             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")
 
             visualizer = Visualizer(num_rows=1, num_cols=5, figure_size=(12, 3))
             visualizer.add_image(image=image, title="Image")
-            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")
+
+            if self.task == "classification":
+                image_classified = visualizer.add_text(
+                    image=image, text=f'Pred: {outputs["pred_scores"][i]:.3f} GT: {int(outputs["label"])}'
+                )
+                visualizer.add_image(image=image_classified, title="Classified Image")
+
+            if self.task == "segmentation":
+                true_mask = outputs["mask"][i].cpu().numpy()
+                visualizer.add_image(image=true_mask, color_map="gray", title="Ground Truth")
+
             self._add_images(visualizer, pl_module, Path(filename))
             visualizer.close()