Source code for

# Copyright (C) 2019-2021 Intel Corporation
# SPDX-License-Identifier: MIT

import logging as log
import os
import os.path as osp
import warnings
from collections import Counter
from enum import Enum, auto
from typing import TYPE_CHECKING, Union

import cv2
import numpy as np

from datumaro.components.annotation import AnnotationType, LabelCategories
from datumaro.components.dataset import IDataset
from import Image
from datumaro.util import parse_str_enum_value
from datumaro.util.image import save_image
from datumaro.util.import_util import lazy_import

    with warnings.catch_warnings():
        import tensorboardX as tb
    tb = lazy_import("tensorboardX")

[docs] class DistanceCompareVisualizer:
[docs] class OutputFormat(Enum): simple = auto() tensorboard = auto()
DEFAULT_FORMAT = OutputFormat.simple _UNMATCHED_LABEL = -1 def __init__( self, comparator, save_dir: str, output_format: Union[None, str, OutputFormat] = None ): self._cmp = comparator self._output_format = parse_str_enum_value( output_format, self.OutputFormat, default=self.DEFAULT_FORMAT ) self._save_dir = save_dir def __enter__(self): os.makedirs(self._save_dir, exist_ok=True) if self._output_format is self.OutputFormat.tensorboard: logdir = osp.join(self._save_dir, "logs", "compare") self._file_writer = tb.SummaryWriter(logdir) elif self._output_format is self.OutputFormat.simple: self._label_diff_writer = None self._a_classes = {} self._b_classes = {} self.label_confusion_matrix = Counter() self.bbox_confusion_matrix = Counter() self.polygon_confusion_matrix = Counter() self.mask_confusion_matrix = Counter() return self def __exit__(self, *args, **kwargs): if self._output_format is self.OutputFormat.tensorboard: self._file_writer.flush() self._file_writer.close() elif self._output_format is self.OutputFormat.simple: if self._label_diff_writer: self._label_diff_writer.flush() self._label_diff_writer.close()
[docs] def save(self, a: IDataset, b: IDataset): if len(a) != len(b): print("Datasets have different lengths: %s vs %s" % (len(a), len(b))) a_classes = set(a.get_label_cat_names()) b_classes = set(b.get_label_cat_names()) if a_classes ^ b_classes: print("Datasets have mismatching labels:") for idx, diff in enumerate(a_classes - b_classes): print(" #%s: > %s" % (idx, diff)) for idx, diff in enumerate(b_classes - a_classes, start=len((a_classes - b_classes))): print(" #%s: < %s" % (idx, diff)) self._a_classes = a.categories().get(AnnotationType.label) self._b_classes = b.categories().get(AnnotationType.label) ids_a = set((, item.subset) for item in a) ids_b = set((, item.subset) for item in b) ids = ids_a & ids_b if len(ids) != len(ids_a): print("Unmatched items in the first dataset: ") print(ids_a - ids) if len(ids) != len(ids_b): print("Unmatched items in the second dataset: ") print(ids_b - ids) for item_id, item_subset in ids: item_a = a.get(item_id, item_subset) item_b = b.get(item_id, item_subset) label_diff = self._cmp.match_labels(item_a, item_b) self.update_label_confusion(label_diff) bbox_diff = self._cmp.match_boxes(item_a, item_b) self.update_bbox_confusion(bbox_diff) polygon_diff = self._cmp.match_polygons(item_a, item_b) self.update_polygon_confusion(polygon_diff) mask_diff = self._cmp.match_masks(item_a, item_b) self.update_mask_confusion(mask_diff) self.save_item_label_diff(item_a, item_b, label_diff) if ( a.media_type() and issubclass(a.media_type(), Image) and b.media_type() and issubclass(b.media_type(), Image) ): self.save_item_bbox_diff(item_a, item_b, bbox_diff) if len(self.label_confusion_matrix) != 0: self.save_conf_matrix(self.label_confusion_matrix, "label_confusion.png") if len(self.bbox_confusion_matrix) != 0: self.save_conf_matrix(self.bbox_confusion_matrix, "bbox_confusion.png") if len(self.polygon_confusion_matrix) != 0: self.save_conf_matrix(self.polygon_confusion_matrix, "polygon_confusion.png") if len(self.mask_confusion_matrix) != 0: self.save_conf_matrix(self.mask_confusion_matrix, "mask_confusion.png")
[docs] def update_label_confusion(self, label_diff): matches, a_unmatched, b_unmatched = label_diff for label in matches: self.label_confusion_matrix[(label, label)] += 1 for a_label in a_unmatched: self.label_confusion_matrix[(a_label, self._UNMATCHED_LABEL)] += 1 for b_label in b_unmatched: self.label_confusion_matrix[(self._UNMATCHED_LABEL, b_label)] += 1
@classmethod def _update_segment_confusion(cls, matrix, diff): matches, mispred, a_unmatched, b_unmatched = diff for a_segm, b_segm in matches: matrix[(a_segm.label, b_segm.label)] += 1 for a_segm, b_segm in mispred: matrix[(a_segm.label, b_segm.label)] += 1 for a_segm in a_unmatched: matrix[(a_segm.label, cls._UNMATCHED_LABEL)] += 1 for b_segm in b_unmatched: matrix[(cls._UNMATCHED_LABEL, b_segm.label)] += 1
[docs] def update_bbox_confusion(self, diff): self._update_segment_confusion(self.bbox_confusion_matrix, diff)
[docs] def update_polygon_confusion(self, diff): self._update_segment_confusion(self.polygon_confusion_matrix, diff)
[docs] def update_mask_confusion(self, diff): self._update_segment_confusion(self.mask_confusion_matrix, diff)
[docs] @classmethod def draw_text_with_background( cls, frame, text, origin, font=None, scale=1.0, color=(0, 0, 0), thickness=1, bgcolor=(1, 1, 1), ): if not font: font = cv2.FONT_HERSHEY_SIMPLEX text_size, baseline = cv2.getTextSize(text, font, scale, thickness) cv2.rectangle( frame, tuple((origin + (0, baseline)).astype(int)), tuple((origin + (text_size[0], -text_size[1])).astype(int)), bgcolor, cv2.FILLED, ) cv2.putText(frame, text, tuple(origin.astype(int)), font, scale, color, thickness) return text_size, baseline
[docs] def draw_detection_roi(self, frame, x, y, w, h, label, conf, color): cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2) text = "%s %.2f%%" % (label, 100.0 * conf) text_scale = 0.5 font = cv2.FONT_HERSHEY_SIMPLEX text_size = cv2.getTextSize(text, font, text_scale, 1) line_height = np.array([0, text_size[0][1]]) self.draw_text_with_background( frame, text, np.array([x, y]) - line_height * 0.5, font, scale=text_scale, color=[255 - c for c in color], )
[docs] def get_a_label(self, label_id): return self._get_label(self._a_classes, label_id)
[docs] def get_b_label(self, label_id): return self._get_label(self._b_classes, label_id)
@staticmethod def _get_label(cat: LabelCategories, label_id): if cat is None: return str(label_id) return cat[label_id].name
[docs] def draw_bbox(self, img, shape, label, color): x, y, w, h = shape.get_bbox() self.draw_detection_roi( img, int(x), int(y), int(w), int(h), label, shape.attributes.get("score", 1), color )
[docs] def get_label_diff_file(self): if self._label_diff_writer is None: self._label_diff_writer = open( osp.join(self._save_dir, "label_diff.txt"), "w", encoding="utf-8" ) return self._label_diff_writer
[docs] def save_item_label_diff(self, item_a, item_b, diff): _, a_unmatched, b_unmatched = diff if 0 < len(a_unmatched) + len(b_unmatched): if self._output_format is self.OutputFormat.simple: f = self.get_label_diff_file() f.write( + "\n") for a_label in a_unmatched: f.write(" >%s\n" % self.get_a_label(a_label)) for b_label in b_unmatched: f.write(" <%s\n" % self.get_b_label(b_label)) elif self._output_format is self.OutputFormat.tensorboard: tag = for a_label in a_unmatched: self._file_writer.add_text(tag, ">%s\n" % self.get_a_label(a_label)) for b_label in b_unmatched: self._file_writer.add_text(tag, "<%s\n" % self.get_b_label(b_label))
[docs] def save_item_bbox_diff(self, item_a, item_b, diff): _, mispred, a_unmatched, b_unmatched = diff if 0 < len(a_unmatched) + len(b_unmatched) + len(mispred): if not isinstance(, Image) or not log.warning("Item %s: item has no image data, " "it will be skipped" % ( return img_a = img_b = img_a.copy() for a_bbox, b_bbox in mispred: self.draw_bbox(img_a, a_bbox, self.get_a_label(a_bbox.label), (0, 255, 0)) self.draw_bbox(img_b, b_bbox, self.get_b_label(b_bbox.label), (0, 0, 255)) for a_bbox in a_unmatched: self.draw_bbox(img_a, a_bbox, self.get_a_label(a_bbox.label), (255, 255, 0)) for b_bbox in b_unmatched: self.draw_bbox(img_b, b_bbox, self.get_b_label(b_bbox.label), (255, 255, 0)) img = np.hstack([img_a, img_b]) path = osp.join(self._save_dir, item_a.subset, if self._output_format is self.OutputFormat.simple: save_image(path + ".png", img, create_dir=True) elif self._output_format is self.OutputFormat.tensorboard: self.save_as_tensorboard(img, path)
[docs] def save_as_tensorboard(self, img, name): img = img[:, :, ::-1] # to RGB img = np.transpose(img, (2, 0, 1)) # to (C, H, W) img = img.astype(dtype=np.uint8) self._file_writer.add_image(name, img)
[docs] def save_conf_matrix(self, conf_matrix, filename): import matplotlib.pyplot as plt def _get_class_map(label_categories): classes = None if label_categories is not None: classes = {id: for id, c in enumerate(label_categories.items)} if classes is None: classes = {c: "label_%s" % c for c, _ in conf_matrix} classes[self._UNMATCHED_LABEL] = "unmatched" classes[None] = "no_class" return classes a_classes = _get_class_map(self._a_classes) b_classes = _get_class_map(self._b_classes) a_class_idx = {id: i for i, id in enumerate(a_classes)} b_class_idx = {id: i for i, id in enumerate(b_classes)} matrix = np.zeros((len(a_classes), len(b_classes)), dtype=int) for idx_pair in conf_matrix: index = (a_class_idx[idx_pair[0]], b_class_idx[idx_pair[1]]) matrix[index] = conf_matrix[idx_pair] a_labels = [label for id, label in a_classes.items()] b_labels = [label for id, label in b_classes.items()] fig = plt.figure() fig.add_subplot(111) table = plt.table(cellText=matrix, rowLabels=a_labels, colLabels=b_labels, loc="center") table.auto_set_font_size(False) table.set_fontsize(8) table.scale(3, 3) # Removing ticks and spines enables you to get the figure only with table plt.tick_params(axis="x", which="both", bottom=False, top=False, labelbottom=False) plt.tick_params(axis="y", which="both", right=False, left=False, labelleft=False) for pos in ["right", "top", "bottom", "left"]: plt.gca().spines[pos].set_visible(False) for idx_pair in conf_matrix: i = a_class_idx[idx_pair[0]] j = b_class_idx[idx_pair[1]] if conf_matrix[idx_pair] != 0: if a_classes[idx_pair[0]] == b_classes[idx_pair[1]]: table._cells[(i + 1, j)].set_facecolor("#00FF00") else: table._cells[(i + 1, j)].set_facecolor("#FF0000") plt.savefig(osp.join(self._save_dir, filename), bbox_inches="tight", pad_inches=0.05)