Add segmentation mask to inference output

anomalib/deploy/inferencers/base.py

@@ -18,12 +18,14 @@
 from pathlib import Path
 from typing import Dict, Optional, Tuple, Union, cast
 
+import cv2
 import numpy as np
 from omegaconf import DictConfig, OmegaConf
+from skimage.segmentation import mark_boundaries
 from torch import Tensor
 
 from anomalib.data.utils import read_image
-from anomalib.post_processing import superimpose_anomaly_map
+from anomalib.post_processing import compute_mask, superimpose_anomaly_map
 from anomalib.post_processing.normalization.cdf import normalize as normalize_cdf
 from anomalib.post_processing.normalization.cdf import standardize
 from anomalib.post_processing.normalization.min_max import (
@@ -60,7 +62,11 @@ def post_process(
         raise NotImplementedError
 
     def predict(
-        self, image: Union[str, np.ndarray, Path], superimpose: bool = True, meta_data: Optional[dict] = None
+        self,
+        image: Union[str, np.ndarray, Path],
+        superimpose: bool = True,
+        meta_data: Optional[dict] = None,
+        overlay_mask: bool = False,
     ) -> Tuple[np.ndarray, float]:
         """Perform a prediction for a given input image.
 
@@ -74,6 +80,8 @@ def predict(
                 will be superimposed onto the original image. If false, `predict`
                 method will return the raw heatmap.
 
+            overlay_mask (bool): If this is set to True, output segmentation mask on top of image.
+
         Returns:
             np.ndarray: Output predictions to be visualized.
         """
@@ -83,18 +91,42 @@ def predict(
             else:
                 meta_data = {}
         if isinstance(image, (str, Path)):
-            image = read_image(image)
-        meta_data["image_shape"] = image.shape[:2]
+            image_arr: np.ndarray = read_image(image)
+        else:  # image is already a numpy array. Kept for mypy compatibility.
+            image_arr = image
+        meta_data["image_shape"] = image_arr.shape[:2]
 
-        processed_image = self.pre_process(image)
+        processed_image = self.pre_process(image_arr)
         predictions = self.forward(processed_image)
         anomaly_map, pred_scores = self.post_process(predictions, meta_data=meta_data)
 
+        # Overlay segmentation mask using raw predictions
+        if overlay_mask and meta_data is not None:
+            image_arr = self._superimpose_segmentation_mask(meta_data, anomaly_map, image_arr)
+
         if superimpose is True:
-            anomaly_map = superimpose_anomaly_map(anomaly_map, image)
+            anomaly_map = superimpose_anomaly_map(anomaly_map, image_arr)
 
         return anomaly_map, pred_scores
 
+    def _superimpose_segmentation_mask(self, meta_data: dict, anomaly_map: np.ndarray, image: np.ndarray):
+        """Superimpose segmentation mask on top of image.
+
+        Args:
+            meta_data (dict): Metadata of the image which contains the image size.
+            anomaly_map (np.ndarray): Anomaly map which is used to extract segmentation mask.
+            image (np.ndarray): Image on which segmentation mask is to be superimposed.
+
+        Returns:
+            np.ndarray: Image with segmentation mask superimposed.
+        """
+        pred_mask = compute_mask(anomaly_map, 0.5)  # assumes predictions are normalized.
+        image_height = meta_data["image_shape"][0]
+        image_width = meta_data["image_shape"][1]
+        pred_mask = cv2.resize(pred_mask, (image_width, image_height))
+        image = mark_boundaries(image, pred_mask, color=(1, 0, 0), mode="thick")
+        return (image * 255).astype(np.uint8)
+
     def __call__(self, image: np.ndarray) -> Tuple[np.ndarray, float]:
         """Call predict on the Image.
 

tools/inference.py

@@ -124,7 +124,7 @@ def infer(image_path: Path, inferencer: Inferencer, save_path: Optional[Path] =
     # path is provided, `predict` method will read the image from
     # file for convenience. We set the superimpose flag to True
     # to overlay the predicted anomaly map on top of the input image.
-    output = inferencer.predict(image=image_path, superimpose=True)
+    output = inferencer.predict(image=image_path, superimpose=True, overlay_mask=True)
 
     # Incase both anomaly map and scores are returned add scores to the image.
     if isinstance(output, tuple):