File size: 5,546 Bytes
3de7bf6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
"""Post Process This module contains utils function to apply post-processing to the output predictions."""
# Copyright (C) 2022-2024 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
import math
import cv2
import numpy as np
from skimage import morphology
def add_label(
image: np.ndarray,
label_name: str,
color: tuple[int, int, int],
confidence: float | None = None,
font_scale: float = 5e-3,
thickness_scale: float = 1e-3,
) -> np.ndarray:
"""Add a label to an image.
Args:
image (np.ndarray): Input image.
label_name (str): Name of the label that will be displayed on the image.
color (tuple[int, int, int]): RGB values for background color of label.
confidence (float | None): confidence score of the label.
font_scale (float): scale of the font size relative to image size. Increase for bigger font.
thickness_scale (float): scale of the font thickness. Increase for thicker font.
Returns:
np.ndarray: Image with label.
"""
image = image.copy()
img_height, img_width, _ = image.shape
font = cv2.FONT_HERSHEY_PLAIN
text = label_name if confidence is None else f"{label_name} ({confidence*100:.0f}%)"
# get font sizing
font_scale = min(img_width, img_height) * font_scale
thickness = math.ceil(min(img_width, img_height) * thickness_scale)
(width, height), baseline = cv2.getTextSize(text, font, fontScale=font_scale, thickness=thickness)
# create label
label_patch = np.zeros((height + baseline, width + baseline, 3), dtype=np.uint8)
label_patch[:, :] = color
cv2.putText(
label_patch,
text,
(0, baseline // 2 + height),
font,
fontScale=font_scale,
thickness=thickness,
color=0,
lineType=cv2.LINE_AA,
)
# add label to image
image[: baseline + height, : baseline + width] = label_patch
return image
def add_normal_label(image: np.ndarray, confidence: float | None = None) -> np.ndarray:
"""Add the normal label to the image."""
return add_label(image, "normal", (225, 252, 134), confidence)
def add_anomalous_label(image: np.ndarray, confidence: float | None = None) -> np.ndarray:
"""Add the anomalous label to the image."""
return add_label(image, "anomalous", (255, 100, 100), confidence)
def anomaly_map_to_color_map(anomaly_map: np.ndarray, normalize: bool = True) -> np.ndarray:
"""Compute anomaly color heatmap.
Args:
anomaly_map (np.ndarray): Final anomaly map computed by the distance metric.
normalize (bool, optional): Bool to normalize the anomaly map prior to applying
the color map. Defaults to True.
Returns:
np.ndarray: [description]
"""
if normalize:
anomaly_map = (anomaly_map - anomaly_map.min()) / np.ptp(anomaly_map)
anomaly_map = anomaly_map * 255
anomaly_map = anomaly_map.astype(np.uint8)
anomaly_map = cv2.applyColorMap(anomaly_map, cv2.COLORMAP_JET)
return cv2.cvtColor(anomaly_map, cv2.COLOR_BGR2RGB)
def superimpose_anomaly_map(
anomaly_map: np.ndarray,
image: np.ndarray,
alpha: float = 0.4,
gamma: int = 0,
normalize: bool = False,
) -> np.ndarray:
"""Superimpose anomaly map on top of in the input image.
Args:
anomaly_map (np.ndarray): Anomaly map
image (np.ndarray): Input image
alpha (float, optional): Weight to overlay anomaly map
on the input image. Defaults to 0.4.
gamma (int, optional): Value to add to the blended image
to smooth the processing. Defaults to 0. Overall,
the formula to compute the blended image is
I' = (alpha*I1 + (1-alpha)*I2) + gamma
normalize: whether or not the anomaly maps should
be normalized to image min-max at image level
Returns:
np.ndarray: Image with anomaly map superimposed on top of it.
"""
anomaly_map = anomaly_map_to_color_map(anomaly_map.squeeze(), normalize=normalize)
return cv2.addWeighted(anomaly_map, alpha, image, (1 - alpha), gamma)
def compute_mask(anomaly_map: np.ndarray, threshold: float, kernel_size: int = 4) -> np.ndarray:
"""Compute anomaly mask via thresholding the predicted anomaly map.
Args:
anomaly_map (np.ndarray): Anomaly map predicted via the model
threshold (float): Value to threshold anomaly scores into 0-1 range.
kernel_size (int): Value to apply morphological operations to the predicted mask. Defaults to 4.
Returns:
Predicted anomaly mask
"""
anomaly_map = anomaly_map.squeeze()
mask: np.ndarray = np.zeros_like(anomaly_map).astype(np.uint8)
mask[anomaly_map > threshold] = 1
kernel = morphology.disk(kernel_size)
mask = morphology.opening(mask, kernel)
mask *= 255
return mask
def draw_boxes(image: np.ndarray, boxes: np.ndarray, color: tuple[int, int, int]) -> np.ndarray:
"""Draw bounding boxes on an image.
Args:
image (np.ndarray): Source image.
boxes (np.nparray): 2D array of shape (N, 4) where each row contains the xyxy coordinates of a bounding box.
color (tuple[int, int, int]): Color of the drawn boxes in RGB format.
Returns:
np.ndarray: Image showing the bounding boxes drawn on top of the source image.
"""
for box in boxes:
x_1, y_1, x_2, y_2 = box.astype(int)
image = cv2.rectangle(image, (x_1, y_1), (x_2, y_2), color=color, thickness=2)
return image
|