Detectron2:Visualization

# -*- coding: utf-8 -*-

from typing import Final, List, Tuple

from cv2 import FILLED as cv2_FILLED  # noqa
from cv2 import FONT_HERSHEY_SIMPLEX as cv2_FONT_HERSHEY_SIMPLEX  # noqa
from cv2 import LINE_4 as cv2_LINE_4  # noqa
from cv2 import LINE_8 as cv2_LINE_8  # noqa
from cv2 import LINE_AA as cv2_LINE_AA  # noqa
from cv2 import addWeighted as cv2_addWeighted  # noqa
from cv2 import fillPoly as cv2_fillPoly  # noqa
from cv2 import getTextSize as cv2_getTextSize  # noqa
from cv2 import polylines as cv2_polylines  # noqa
from cv2 import putText as cv2_putText  # noqa
from cv2 import rectangle as cv2_rectangle  # noqa
from numpy import ndarray

LINE_4: Final[int] = cv2_LINE_4
LINE_8: Final[int] = cv2_LINE_8
LINE_AA: Final[int] = cv2_LINE_AA
FILLED: Final[int] = cv2_FILLED
FONT_HERSHEY_SIMPLEX: Final[int] = cv2_FONT_HERSHEY_SIMPLEX

DEFAULT_COLOR: Final[Tuple[int, int, int]] = (0, 0, 255)
DEFAULT_THICKNESS: Final[int] = 2
DEFAULT_LINE_TYPE: Final[int] = LINE_AA
DEFAULT_FONT: Final[int] = FONT_HERSHEY_SIMPLEX
DEFAULT_FONT_SCALE: Final[float] = 1.0


def fill_poly(
    image: ndarray,
    points: List[ndarray],
    color=DEFAULT_COLOR,
    line_type=DEFAULT_LINE_TYPE,
) -> None:
    cv2_fillPoly(image, points, color, line_type)


def polylines(
    image: ndarray,
    points: List[ndarray],
    closed=True,
    color=DEFAULT_COLOR,
    thickness=DEFAULT_THICKNESS,
    line_type=DEFAULT_LINE_TYPE,
) -> None:
    cv2_polylines(image, points, closed, color, thickness, line_type)


def add_weighted(
    src1: ndarray,
    alpha: float,
    src2: ndarray,
    beta: float,
    gamma=0.0,
) -> ndarray:
    return cv2_addWeighted(src1, alpha, src2, beta, gamma)


def rectangle(
    image: ndarray,
    point1: Tuple[int, int],
    point2: Tuple[int, int],
    color=DEFAULT_COLOR,
    thickness=DEFAULT_THICKNESS,
    line_type=DEFAULT_LINE_TYPE,
) -> None:
    cv2_rectangle(image, point1, point2, color, thickness, line_type)


def put_text(
    image: ndarray,
    text: str,
    bottom_left_corner: Tuple[int, int],
    font_face=DEFAULT_FONT,
    font_scale=DEFAULT_FONT_SCALE,
    color=DEFAULT_COLOR,
    thickness=DEFAULT_THICKNESS,
    line_type=DEFAULT_LINE_TYPE,
) -> None:
    cv2_putText(
        image,
        text,
        bottom_left_corner,
        font_face,
        font_scale,
        color,
        thickness,
        line_type,
        bottomLeftOrigin=False,
    )


def get_text_size(
    text: str,
    font_face=DEFAULT_FONT,
    font_scale=DEFAULT_FONT_SCALE,
    thickness=DEFAULT_THICKNESS,
) -> Tuple[Tuple[int, int], int]:
    return cv2_getTextSize(text, font_face, font_scale, thickness)

# -*- coding: utf-8 -*-

from colorsys import hls_to_rgb, rgb_to_hls
from typing import Tuple


def adjust_color(r: int, g: int, b: int, factor: float) -> Tuple[int, int, int]:
    h, l, s = rgb_to_hls(r / 255.0, g / 255.0, b / 255.0)
    adjust_l = max(min(l * factor, 1.0), 0.0)
    fr, fg, fb = hls_to_rgb(h, adjust_l, s)
    return int(fr * 255), int(fg * 255), int(fb * 255)


def lighten_color(r: int, g: int, b: int, factor=0.1):
    return adjust_color(r, g, b, 1 + factor)


def darken_color(r: int, g: int, b: int, factor=0.1):
    return adjust_color(r, g, b, 1 - factor)

# -*- coding: utf-8 -*-

from typing import Tuple


def invert_color(r: int, g: int, b: int) -> Tuple[int, int, int]:
    return abs(r - 255), abs(g - 255), abs(b - 255)

# -*- coding: utf-8 -*-
# https://flatuicolors.com/palette/us
# American Palette by Kevin Yang from

from typing import Tuple


def _rgb_to_bgr(r: int, g: int, b: int) -> Tuple[int, int, int]:
    return b, g, r


LIGHT_GREENISH_BLUE = _rgb_to_bgr(85, 239, 196)
FADED_POSTER = _rgb_to_bgr(129, 236, 236)
GREEN_DARNER_TAIL = _rgb_to_bgr(116, 185, 255)
SHY_MOMENT = _rgb_to_bgr(162, 155, 254)

CITY_LIGHTS = _rgb_to_bgr(223, 230, 233)
MINT_LEAF = _rgb_to_bgr(0, 184, 148)
ROBINS_EGG_BLUE = _rgb_to_bgr(0, 206, 201)
ELECTRON_BLUE = _rgb_to_bgr(9, 132, 227)

EXODUS_FRUIT = _rgb_to_bgr(108, 92, 231)
SOOTHING_BREEZE = _rgb_to_bgr(178, 190, 195)
SOUR_LEMON = _rgb_to_bgr(255, 234, 167)
FIRST_DATE = _rgb_to_bgr(250, 177, 160)

PINK_GLAMOUR = _rgb_to_bgr(255, 118, 117)
PICO_8_PINK = _rgb_to_bgr(253, 121, 168)
AMERICAN_RIVER = _rgb_to_bgr(99, 110, 114)
BRIGHT_YARROW = _rgb_to_bgr(253, 203, 110)

ORANGEVILLE = _rgb_to_bgr(225, 112, 85)
CHI_GONG = _rgb_to_bgr(214, 48, 49)
PRUNUS_AVIUM = _rgb_to_bgr(232, 67, 147)
DRACULA_ORCHID = _rgb_to_bgr(45, 52, 54)

AMERICAN_PALETTE = [
    LIGHT_GREENISH_BLUE,
    FADED_POSTER,
    GREEN_DARNER_TAIL,
    SHY_MOMENT,
    CITY_LIGHTS,
    MINT_LEAF,
    ROBINS_EGG_BLUE,
    ELECTRON_BLUE,
    EXODUS_FRUIT,
    SOOTHING_BREEZE,
    SOUR_LEMON,
    FIRST_DATE,
    PINK_GLAMOUR,
    PICO_8_PINK,
    AMERICAN_RIVER,
    BRIGHT_YARROW,
    ORANGEVILLE,
    CHI_GONG,
    PRUNUS_AVIUM,
    DRACULA_ORCHID,
]

# -*- coding: utf-8 -*-

from dataclasses import dataclass
from io import StringIO
from typing import Final, List, Optional, Tuple

from adjust import darken_color
from american_palette import AMERICAN_PALETTE
from invert import invert_color
from draw import (
    FILLED,
    add_weighted,
    fill_poly,
    get_text_size,
    polylines,
    put_text,
    rectangle,
)
from numpy import array, float64, int32, ndarray

Polyline = List[float]  # x0, y0, x1, y1, x2, y2, ... xN, yN
Polylines = List[Polyline]
PolylinesList = List[Polylines]

Box = List[float]  # x1, y1, x2, y2
Boxes = List[Box]

BLACK_COLOR: Final[Tuple[int, int, int]] = (0, 0, 0)
WHITE_COLOR: Final[Tuple[int, int, int]] = (255, 255, 255)

DEFAULT_LINE_WIDTH: Final[int] = 2
DEFAULT_OVERLAY_ALPHA: Final[float] = 0.3
DEFAULT_FONT_SCALE: Final[float] = 0.7
DEFAULT_LABEL_BOX_PADDING: Final[int] = 2
DEFAULT_DARKEN_FACTOR: Final[float] = 0.6


@dataclass
class Point:
    x: float
    y: float


@dataclass
class Rect:
    x1: float
    y1: float
    x2: float
    y2: float

    def to_string(self):
        return f"{self.__class__.__name__}[{self.x1},{self.y1},{self.x2},{self.y2}]"

    def __repr__(self):
        return self.to_string()

    def __str__(self):
        return self.to_string()

    @property
    def left(self):
        return min(self.x1, self.x2)

    @property
    def right(self):
        return max(self.x1, self.x2)

    @property
    def top(self):
        return min(self.y1, self.y2)

    @property
    def bottom(self):
        return max(self.y1, self.y2)

    @property
    def width(self):
        return abs(self.x2 - self.x1)

    @property
    def height(self):
        return abs(self.y2 - self.y1)

    def intersection(self, obj: "Rect") -> bool:
        x1 = max(self.left, obj.left)
        y1 = max(self.top, obj.top)
        x2 = min(self.right, obj.right)
        y2 = min(self.bottom, obj.bottom)
        return x1 < x2 and y1 < y2


@dataclass
class ColorInfo:
    line_bgr: Tuple[int, int, int]
    fill_bgr: Tuple[int, int, int]
    label_bgr: Tuple[int, int, int]

    @classmethod
    def from_base_bgr(cls, base_color: Tuple[int, int, int], darken_factor: float):
        fill_bgr = base_color
        b, g, r = fill_bgr
        line_bgr = darken_color(r, g, b, darken_factor)[::-1]
        label_bgr = invert_color(line_bgr[0], line_bgr[1], line_bgr[2])
        return cls(line_bgr, fill_bgr, label_bgr)


@dataclass
class LabelBoxInfo:
    label_box_left: int = 0
    label_box_top: int = 0
    label_box_right: int = 0
    label_box_bottom: int = 0

    label_left: int = 0
    label_bottom: int = 0

    def as_label_box_left_top(self) -> Tuple[int, int]:
        return self.label_box_left, self.label_box_top

    def as_label_box_right_bottom(self) -> Tuple[int, int]:
        return self.label_box_right, self.label_box_bottom

    def as_label_left_bottom(self) -> Tuple[int, int]:
        return self.label_left, self.label_bottom

    def adjust_screen_position(
        self,
        width: Optional[int] = None,
        height: Optional[int] = None,
    ) -> None:
        """If the 'label box' is off the screen, adjust its position."""
        if self.label_box_left < 0:
            # Move to right
            move_right = abs(self.label_box_left)
            self.label_box_left += move_right
            self.label_box_right += move_right
            self.label_left += move_right
        elif width is not None and self.label_box_right > width:
            # Move to left
            move_left = self.label_box_right - width
            assert move_left > 0
            self.label_box_left -= move_left
            self.label_box_right -= move_left
            self.label_left -= move_left

        if self.label_box_top < 0:
            # Move to down
            move_down = abs(self.label_box_top)
            self.label_box_top += move_down
            self.label_box_bottom += move_down
            self.label_bottom += move_down
        elif height is not None and self.label_box_bottom > height:
            # Move to up
            move_up = self.label_box_bottom - height
            assert move_up > 0
            self.label_box_top -= move_up
            self.label_box_bottom -= move_up
            self.label_bottom -= move_up


@dataclass
class PredictInstance:
    index: int  # aka class
    score: float
    label: Optional[str] = None
    box: Optional[Rect] = None
    polylines: Optional[Polylines] = None

    def to_string(self):
        io = StringIO()
        io.write(self.__class__.__name__)
        io.write(f"[index={self.index},score={self.score:.2f}")
        if self.label is not None:
            io.write(f",label={self.label}")
        if self.box is not None:
            io.write(f",box={self.box}")
        io.write("]")
        return io.getvalue()

    def __repr__(self):
        return self.to_string()

    def __str__(self):
        return self.to_string()

    def to_drawable_label(self) -> str:
        if self.label:
            return "{} {:.0f}%".format(self.label, self.score * 100)
        else:
            return "#{} {:.0f}%".format(self.index, self.score * 100)

    def as_numpy_polylines(self) -> List[ndarray]:
        if not self.polylines:
            return list()

        if not isinstance(self.polylines[0], list):
            element_typename = type(self.polylines[0]).__name__
            raise TypeError(f"Unexpected polylines element: {element_typename}")

        result = []
        for polyline in self.polylines:
            assert isinstance(polyline, list)
            arr = array(polyline, float64).astype(int32).reshape((-1, 2))
            result.append(arr)
        return result

    def get_text_size(self, font_scale: float) -> Tuple[int, int, int]:
        label_size = get_text_size(self.to_drawable_label(), font_scale=font_scale)
        label_width, label_height = label_size[0]
        label_base_line = label_size[1]
        return label_width, label_height, label_base_line

    def get_label_box_info(
        self,
        font_scale: float,
        label_box_padding: int,
        image_width: Optional[int] = None,
        image_height: Optional[int] = None,
    ) -> LabelBoxInfo:
        label_width, label_height, label_base_line = self.get_text_size(font_scale)

        box_left = int(self.box.x1) if self.box else 0
        box_top = int(self.box.y1) if self.box else 0

        label_box_left = box_left
        label_box_top = box_top - label_height - (label_base_line * 2)
        label_box_right = box_left + label_width + (label_box_padding * 2)
        label_box_bottom = box_top

        label_left = box_left + label_box_padding
        label_bottom = box_top - label_base_line

        result = LabelBoxInfo(
            label_box_left,
            label_box_top,
            label_box_right,
            label_box_bottom,
            label_left,
            label_bottom,
        )
        result.adjust_screen_position(image_width, image_height)
        return result

    def draw_polylines(
        self,
        image: ndarray,
        line_width: int,
        fill_color: Tuple[int, int, int],
        line_color: Tuple[int, int, int],
    ) -> None:
        numpy_polylines = self.as_numpy_polylines()
        fill_poly(image, numpy_polylines, fill_color)
        polylines(image, numpy_polylines, True, line_color, line_width)

    def draw_box(
        self,
        image: ndarray,
        line_width: int,
        line_color: Tuple[int, int, int],
    ) -> None:
        box_left = int(self.box.x1) if self.box else 0
        box_right = int(self.box.x2) if self.box else 0
        box_top = int(self.box.y1) if self.box else 0
        box_bottom = int(self.box.y2) if self.box else 0
        rectangle(
            image,
            (box_left, box_top),
            (box_right, box_bottom),
            line_color,
            line_width,
        )

    def draw_label_box(
        self,
        image: ndarray,
        font_scale: float,
        label_box_padding: int,
        fill_color: Tuple[int, int, int],
    ) -> None:
        label_box_info = self.get_label_box_info(
            font_scale=font_scale,
            label_box_padding=label_box_padding,
            image_width=image.shape[1],
            image_height=image.shape[0],
        )
        rectangle(
            image,
            label_box_info.as_label_box_left_top(),
            label_box_info.as_label_box_right_bottom(),
            fill_color,
            FILLED,
        )

    def draw_label(
        self,
        image: ndarray,
        font_scale: float,
        label_box_padding: int,
        label_color: Tuple[int, int, int],
    ) -> None:
        label_box_info = self.get_label_box_info(
            font_scale=font_scale,
            label_box_padding=label_box_padding,
            image_width=image.shape[1],
            image_height=image.shape[0],
        )
        put_text(
            image,
            self.to_drawable_label(),
            label_box_info.as_label_left_bottom(),
            font_scale=font_scale,
            color=label_color,
        )

    def draw_overlay(
        self,
        image: ndarray,
        fill_color: Tuple[int, int, int],
        line_color: Tuple[int, int, int],
        line_width: int,
        overlay_alpha: float,
        font_scale: float,
        label_box_padding: int,
        with_polylines=True,
        with_box=True,
        with_label=True,
    ) -> None:
        overlay = image.copy()
        if with_polylines and self.polylines:
            self.draw_polylines(overlay, line_width, fill_color, line_color)
        if with_box and self.box is not None:
            self.draw_box(overlay, line_width, line_color)
        if with_label and self.label is not None and self.box is not None:
            self.draw_label_box(overlay, font_scale, label_box_padding, line_color)

        alpha = 1 - overlay_alpha
        image[:] = add_weighted(overlay, alpha, image, overlay_alpha)

    def draw(
        self,
        image: ndarray,
        line_width=DEFAULT_LINE_WIDTH,
        overlay_alpha=DEFAULT_OVERLAY_ALPHA,
        font_scale=DEFAULT_FONT_SCALE,
        label_box_padding=DEFAULT_LABEL_BOX_PADDING,
        darken_factor=DEFAULT_DARKEN_FACTOR,
        with_polylines=True,
        with_box=True,
        with_label=True,
        palette: List[Tuple[int, int, int]] = AMERICAN_PALETTE,
    ) -> None:
        color_info = ColorInfo.from_base_bgr(palette[self.index], darken_factor)
        self.draw_overlay(
            image=image,
            fill_color=color_info.fill_bgr,
            line_color=color_info.line_bgr,
            line_width=line_width,
            overlay_alpha=overlay_alpha,
            font_scale=font_scale,
            label_box_padding=label_box_padding,
            with_polylines=with_polylines,
            with_box=with_box,
            with_label=with_label,
        )

        # If the label is printed as an overlay, readability is reduced,
        # so it should be drawn opaque.
        self.draw_label(
            image=image,
            font_scale=font_scale,
            label_box_padding=label_box_padding,
            label_color=color_info.label_bgr,
        )


@dataclass
class PredictQ:
    threshold: Optional[float] = None
    roi: Optional[Rect] = None


@dataclass
class PredictA:
    predictions: List[PredictInstance]
    iteration: Optional[int] = None
    elapsed: Optional[float] = None

    def to_string(self):
        io = StringIO()
        io.write(self.__class__.__name__)
        io.write(f"[predictions={len(self.predictions)}")
        if self.iteration is not None:
            io.write(f",iteration={self.iteration}")
        if self.elapsed is not None:
            io.write(f",elapsed={self.elapsed:.2f}s")
        io.write("]")
        return io.getvalue()

    def __repr__(self):
        return self.to_string()

    def __str__(self):
        return self.to_string()

    @classmethod
    def from_raws(
        cls,
        classes: List[int],
        labels: List[str],
        scores: List[float],
        boxes: Boxes,
        polylines_list: PolylinesList,
        iteration: Optional[int] = None,
        elapsed: Optional[float] = None,
    ):
        len_labels = len(labels)
        len_scores = len(scores)
        len_boxes = len(boxes)
        len_polylines_list = len(polylines_list)
        if not (len_labels == len_scores == len_boxes == len_polylines_list):
            raise IndexError("All raw data must be the same length")

        predictions = list()
        for i in range(len_labels):
            x1, y1, x2, y2 = boxes[i]
            instance = PredictInstance(
                index=classes[i],
                score=scores[i],
                label=labels[i],
                box=Rect(x1, y1, x2, y2),
                polylines=polylines_list[i],
            )
            predictions.append(instance)

        return cls(
            predictions=predictions,
            iteration=iteration,
            elapsed=elapsed,
        )

    def draw(
        self,
        image: ndarray,
        line_width=DEFAULT_LINE_WIDTH,
        overlay_alpha=DEFAULT_OVERLAY_ALPHA,
        font_scale=DEFAULT_FONT_SCALE,
        label_box_padding=DEFAULT_LABEL_BOX_PADDING,
        darken_factor=DEFAULT_DARKEN_FACTOR,
        with_polylines=True,
        with_box=True,
        with_label=True,
        palette: List[Tuple[int, int, int]] = AMERICAN_PALETTE,
    ) -> None:
        # If the printed 'label' overlaps,
        # an object with a higher score is drawn on top.
        predictions = self.predictions.copy()
        predictions.sort(key=lambda x: x.score)

        for predict in predictions:
            predict.draw(
                image=image,
                line_width=line_width,
                overlay_alpha=overlay_alpha,
                font_scale=font_scale,
                label_box_padding=label_box_padding,
                darken_factor=darken_factor,
                with_polylines=with_polylines,
                with_box=with_box,
                with_label=with_label,
                palette=palette,
            )

from io import StringIO
from typing import List, Optional


def create_text_labels(
    classes: List[int],
    scores: Optional[List[float]] = None,
    is_crowd: Optional[List[bool]] = None,
    *,
    labels: Optional[List[str]] = None,
) -> List[str]:
    result = list()
    for i, class_number in enumerate(classes):
        io = StringIO()
        io.write(labels[class_number] if labels else f"[{class_number}]")
        if scores:
            io.write(" {:.0f}%".format(scores[i] * 100))
        if is_crowd and is_crowd[i]:
            io.write("|crowd")
        result.append(io.getvalue())
    return result


def read_labels_file(path: str) -> List[str]:
    try:
        with open(path, mode="r") as f:
            return f.read().strip().split("\n")
    except BaseException:  # noqa
        return list()


cpu_device = device("cpu")
label_table = read_labels_file(labels_file if labels_file else str())


def predictions_to_answer(
    predictions: Dict[str, Any],
    width: int,
    height: int,
) -> PredictA:
    assert "instances" in predictions
    objs = predictions["instances"].to(cpu_device)
    assert isinstance(objs, Instances)

    assert objs.has("pred_boxes")
    assert isinstance(objs.pred_boxes, Boxes)
    boxes = [b.tolist() for b in objs.pred_boxes]
    assert isinstance(boxes, list)

    assert objs.has("scores")
    assert isinstance(objs.scores, Tensor)
    scores = objs.scores.tolist()
    assert isinstance(scores, list)

    assert objs.has("pred_classes")
    classes = objs.pred_classes.tolist()
    assert isinstance(classes, list)
    labels = create_text_labels(
        classes=classes,
        scores=None,
        is_crowd=None,
        labels=label_table,
    )

    # objs.has("pred_keypoints")
    # keypoints = objs.pred_keypoints if objs.has("pred_keypoints") else None
    # assert keypoints is None

    assert objs.has("pred_masks")
    pred_masks = asarray(objs.pred_masks)
    masks = [GenericMask(x, height, width) for x in pred_masks]

    polylines_list = list()
    for mask in masks:
        polylines = list()
        for polygon in mask.polygons:
            raw_polyline = polygon.tolist()
            assert isinstance(raw_polyline, list)
            polylines.append(raw_polyline)
        polylines_list.append(polylines)

    return PredictA.from_raws(
        classes,
        labels,
        scores,
        boxes,
        polylines_list,
    )


cfg = ...
image = ...

height, width, channel = image.shape
assert height >= 1
assert width >= 1
assert channel == 3

from detectron2.engine.defaults import DefaultPredictor
predictor = DefaultPredictor(cfg)
predictions = predictor(image)
answer = predictions_to_answer(predictions, width, height)
answer.draw(image)