| | |
| | import logging |
| | import numpy as np |
| | import cv2 |
| | import torch |
| |
|
| | Image = np.ndarray |
| | Boxes = torch.Tensor |
| |
|
| |
|
| | class MatrixVisualizer: |
| | """ |
| | Base visualizer for matrix data |
| | """ |
| |
|
| | def __init__( |
| | self, |
| | inplace=True, |
| | cmap=cv2.COLORMAP_PARULA, |
| | val_scale=1.0, |
| | alpha=0.7, |
| | interp_method_matrix=cv2.INTER_LINEAR, |
| | interp_method_mask=cv2.INTER_NEAREST, |
| | ): |
| | self.inplace = inplace |
| | self.cmap = cmap |
| | self.val_scale = val_scale |
| | self.alpha = alpha |
| | self.interp_method_matrix = interp_method_matrix |
| | self.interp_method_mask = interp_method_mask |
| |
|
| | def visualize(self, image_bgr, mask, matrix, bbox_xywh): |
| | self._check_image(image_bgr) |
| | self._check_mask_matrix(mask, matrix) |
| | if self.inplace: |
| | image_target_bgr = image_bgr |
| | else: |
| | image_target_bgr = image_bgr |
| | image_target_bgr *= 0 |
| | x, y, w, h = [int(v) for v in bbox_xywh] |
| | if w <= 0 or h <= 0: |
| | return image_bgr |
| | mask, matrix = self._resize(mask, matrix, w, h) |
| | mask_bg = np.tile((mask == 0)[:, :, np.newaxis], [1, 1, 3]) |
| | matrix_scaled = matrix.astype(np.float32) * self.val_scale |
| | _EPSILON = 1e-6 |
| | if np.any(matrix_scaled > 255 + _EPSILON): |
| | logger = logging.getLogger(__name__) |
| | logger.warning( |
| | f"Matrix has values > {255 + _EPSILON} after " f"scaling, clipping to [0..255]" |
| | ) |
| | matrix_scaled_8u = matrix_scaled.clip(0, 255).astype(np.uint8) |
| | matrix_vis = cv2.applyColorMap(matrix_scaled_8u, self.cmap) |
| | matrix_vis[mask_bg] = image_target_bgr[y : y + h, x : x + w, :][mask_bg] |
| | image_target_bgr[y : y + h, x : x + w, :] = ( |
| | image_target_bgr[y : y + h, x : x + w, :] * (1.0 - self.alpha) + matrix_vis * self.alpha |
| | ) |
| | return image_target_bgr.astype(np.uint8) |
| |
|
| | def _resize(self, mask, matrix, w, h): |
| | if (w != mask.shape[1]) or (h != mask.shape[0]): |
| | mask = cv2.resize(mask, (w, h), self.interp_method_mask) |
| | if (w != matrix.shape[1]) or (h != matrix.shape[0]): |
| | matrix = cv2.resize(matrix, (w, h), self.interp_method_matrix) |
| | return mask, matrix |
| |
|
| | def _check_image(self, image_rgb): |
| | assert len(image_rgb.shape) == 3 |
| | assert image_rgb.shape[2] == 3 |
| | assert image_rgb.dtype == np.uint8 |
| |
|
| | def _check_mask_matrix(self, mask, matrix): |
| | assert len(matrix.shape) == 2 |
| | assert len(mask.shape) == 2 |
| | assert mask.dtype == np.uint8 |
| |
|
| |
|
| | class RectangleVisualizer: |
| |
|
| | _COLOR_GREEN = (18, 127, 15) |
| |
|
| | def __init__(self, color=_COLOR_GREEN, thickness=1): |
| | self.color = color |
| | self.thickness = thickness |
| |
|
| | def visualize(self, image_bgr, bbox_xywh, color=None, thickness=None): |
| | x, y, w, h = bbox_xywh |
| | color = color or self.color |
| | thickness = thickness or self.thickness |
| | cv2.rectangle(image_bgr, (int(x), int(y)), (int(x + w), int(y + h)), color, thickness) |
| | return image_bgr |
| |
|
| |
|
| | class PointsVisualizer: |
| |
|
| | _COLOR_GREEN = (18, 127, 15) |
| |
|
| | def __init__(self, color_bgr=_COLOR_GREEN, r=5): |
| | self.color_bgr = color_bgr |
| | self.r = r |
| |
|
| | def visualize(self, image_bgr, pts_xy, colors_bgr=None, rs=None): |
| | for j, pt_xy in enumerate(pts_xy): |
| | x, y = pt_xy |
| | color_bgr = colors_bgr[j] if colors_bgr is not None else self.color_bgr |
| | r = rs[j] if rs is not None else self.r |
| | cv2.circle(image_bgr, (x, y), r, color_bgr, -1) |
| | return image_bgr |
| |
|
| |
|
| | class TextVisualizer: |
| |
|
| | _COLOR_GRAY = (218, 227, 218) |
| | _COLOR_WHITE = (255, 255, 255) |
| |
|
| | def __init__( |
| | self, |
| | font_face=cv2.FONT_HERSHEY_SIMPLEX, |
| | font_color_bgr=_COLOR_GRAY, |
| | font_scale=0.35, |
| | font_line_type=cv2.LINE_AA, |
| | font_line_thickness=1, |
| | fill_color_bgr=_COLOR_WHITE, |
| | fill_color_transparency=1.0, |
| | frame_color_bgr=_COLOR_WHITE, |
| | frame_color_transparency=1.0, |
| | frame_thickness=1, |
| | ): |
| | self.font_face = font_face |
| | self.font_color_bgr = font_color_bgr |
| | self.font_scale = font_scale |
| | self.font_line_type = font_line_type |
| | self.font_line_thickness = font_line_thickness |
| | self.fill_color_bgr = fill_color_bgr |
| | self.fill_color_transparency = fill_color_transparency |
| | self.frame_color_bgr = frame_color_bgr |
| | self.frame_color_transparency = frame_color_transparency |
| | self.frame_thickness = frame_thickness |
| |
|
| | def visualize(self, image_bgr, txt, topleft_xy): |
| | txt_w, txt_h = self.get_text_size_wh(txt) |
| | topleft_xy = tuple(map(int, topleft_xy)) |
| | x, y = topleft_xy |
| | if self.frame_color_transparency < 1.0: |
| | t = self.frame_thickness |
| | image_bgr[y - t : y + txt_h + t, x - t : x + txt_w + t, :] = ( |
| | image_bgr[y - t : y + txt_h + t, x - t : x + txt_w + t, :] |
| | * self.frame_color_transparency |
| | + np.array(self.frame_color_bgr) * (1.0 - self.frame_color_transparency) |
| | ).astype(float) |
| | if self.fill_color_transparency < 1.0: |
| | image_bgr[y : y + txt_h, x : x + txt_w, :] = ( |
| | image_bgr[y : y + txt_h, x : x + txt_w, :] * self.fill_color_transparency |
| | + np.array(self.fill_color_bgr) * (1.0 - self.fill_color_transparency) |
| | ).astype(float) |
| | cv2.putText( |
| | image_bgr, |
| | txt, |
| | topleft_xy, |
| | self.font_face, |
| | self.font_scale, |
| | self.font_color_bgr, |
| | self.font_line_thickness, |
| | self.font_line_type, |
| | ) |
| | return image_bgr |
| |
|
| | def get_text_size_wh(self, txt): |
| | ((txt_w, txt_h), _) = cv2.getTextSize( |
| | txt, self.font_face, self.font_scale, self.font_line_thickness |
| | ) |
| | return txt_w, txt_h |
| |
|
| |
|
| | class CompoundVisualizer: |
| | def __init__(self, visualizers): |
| | self.visualizers = visualizers |
| |
|
| | def visualize(self, image_bgr, data): |
| | assert len(data) == len( |
| | self.visualizers |
| | ), "The number of datas {} should match the number of visualizers" " {}".format( |
| | len(data), len(self.visualizers) |
| | ) |
| | image = image_bgr |
| | for i, visualizer in enumerate(self.visualizers): |
| | image = visualizer.visualize(image, data[i]) |
| | return image |
| |
|
| | def __str__(self): |
| | visualizer_str = ", ".join([str(v) for v in self.visualizers]) |
| | return "Compound Visualizer [{}]".format(visualizer_str) |
| |
|
