Source code for otx.api.utils.anomaly_utils
"""Detection Utils."""
# Copyright (C) 2021-2022 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
from typing import Dict, List
import cv2
import numpy as np
from otx.api.entities.annotation import Annotation
from otx.api.entities.label import LabelEntity
from otx.api.entities.scored_label import ScoredLabel
from otx.api.entities.shapes.rectangle import Rectangle
[docs]
def create_detection_annotation_from_anomaly_heatmap(
hard_prediction: np.ndarray,
soft_prediction: np.ndarray,
label_map: Dict[int, LabelEntity],
) -> List[Annotation]:
"""Create box annotation from the soft predictions.
Args:
hard_prediction: hard prediction containing the final label
index per pixel.
soft_prediction: soft prediction with shape
label_map: dictionary mapping labels to an index. It is assumed
that the first item in the dictionary corresponds to the
background label and will therefore be ignored.
Returns:
List of annotations.
"""
# pylint: disable=too-many-locals
if hard_prediction.ndim == 3 and hard_prediction.shape[0] == 1:
hard_prediction = hard_prediction.squeeze().astype(np.uint8)
if soft_prediction.ndim == 3 and soft_prediction.shape[0] == 1:
soft_prediction = soft_prediction.squeeze()
image_h, image_w = hard_prediction.shape[:2]
annotations: List[Annotation] = []
for label_index, label in label_map.items():
# Skip the normal label.
if label_index == 0:
continue
# cv2.connectedComponentsWithStats returns num_labels, labels, coordinates
# and centroids. This script only needs the coordinates.
_, connected_components, coordinates, _ = cv2.connectedComponentsWithStats(hard_prediction)
for i, coordinate in enumerate(coordinates):
# First row of the coordinates is always backround,
# so should be ignored.
if i == 0:
continue
# Last column of the coordinates is the area of the connected component.
# It could therefore be ignored.
comp_x, comp_y, comp_w, comp_h, _ = coordinate
# Compute the probability of each connected-component
component_hard_prediction = (connected_components == i).astype(np.uint8)
component_soft_prediction = cv2.bitwise_and(
soft_prediction, soft_prediction, mask=component_hard_prediction
)
# NOTE: Find the best approach to calculate the probability
probability = component_soft_prediction.reshape(-1).max()
# NOTE: NMS could be needed here.
# Create the annotation based on the box shape and the probability.
shape = Rectangle(
x1=comp_x / image_w,
y1=comp_y / image_h,
x2=(comp_x + comp_w) / image_w,
y2=(comp_y + comp_h) / image_h,
)
annotations.append(Annotation(shape=shape, labels=[ScoredLabel(label, float(probability))]))
return annotations