return pred masks in predict step (#103)

openvinotoolkit · Feb 15, 2022 · dbca627 · dbca627
1 parent 4bf4970
commit dbca627
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/anomalib/models/components/base/anomaly_module.py b/anomalib/models/components/base/anomaly_module.py
@@ -95,6 +95,8 @@ def predict_step(self, batch: Any, batch_idx: int, _dataloader_idx: Optional[int
         outputs = self.validation_step(batch, batch_idx)
         self._post_process(outputs)
         outputs["pred_labels"] = outputs["pred_scores"] >= self.image_threshold.value
+        if "anomaly_maps" in outputs.keys():
+            outputs["pred_masks"] = outputs["anomaly_maps"] >= self.pixel_threshold.value
         return outputs
 
     def test_step(self, batch, _):  # pylint: disable=arguments-differ