diff --git a/src/anomalib/post_processing/visualizer.py b/src/anomalib/post_processing/visualizer.py index 05050b1c02..b39b04e6b3 100644 --- a/src/anomalib/post_processing/visualizer.py +++ b/src/anomalib/post_processing/visualizer.py @@ -168,7 +168,14 @@ def _visualize_full(self, image_result: ImageResult) -> np.ndarray: visualization.add_image(image=image_result.gt_mask, color_map="gray", title="Ground Truth") visualization.add_image(image_result.heat_map, "Predicted Heat Map") visualization.add_image(image=image_result.pred_mask, color_map="gray", title="Predicted Mask") - visualization.add_image(image=image_result.segmentations, title="Segmentation Result") + + # add classification result to segmented image + if image_result.pred_label: + image_segmented = add_anomalous_label(image_result.segmentations, image_result.pred_score) + else: + image_segmented = add_normal_label(image_result.segmentations, 1 - image_result.pred_score) + + visualization.add_image(image=image_segmented, title="Segmentation Result") elif self.task == TaskType.CLASSIFICATION: visualization.add_image(image_result.image, title="Image") if hasattr(image_result, "heat_map"):